scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 134

feat(rust): route OCR, messages, and responses-websocket through Rust by default; remove flags (#34035)

Browse source 
* feat(rust): route OCR, messages, and responses-websocket through Rust by default; remove flags and error-fallback

* fix(rust): drop now-unused os import; route Python-handler tests off the default Rust path

* test(rust): assert native-anthropic gate routes to Rust even when litellm_params.rust is false

The gate no longer reads litellm_params.rust after the flag removal, so
passing rust=False here proves Rust runs unconditionally and guards against
anyone re-introducing a rust gate on the native-anthropic path.
This commit is contained in:
Yassin Kortam 2026-07-21 17:14:46 -07:00 • committed by GitHub
parent 2b2ae4ca49
commit e5b19c5583
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 35 additions and 162 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -1,6 +1,5 @@
import asyncio
import json
import os
import ssl
from contextlib import asynccontextmanager
from functools import lru_cache
@ -152,9 +151,8 @@ from litellm.utils import (
def _rust_responses_websocket_enabled(
custom_llm_provider: str | None,
litellm_params: GenericLiteLLMParams,
) -> bool:
return custom_llm_provider == "openai" and litellm_params.get("rust") is True
return custom_llm_provider == "openai"
from .http_handler import get_shared_realtime_ssl_context
@ -2257,10 +2255,6 @@ class BaseLLMHTTPHandler:
"anthropic_messages",
)
@staticmethod
def _rust_env_enabled() -> bool:
return os.getenv("LITELLM_RUST", "").strip().lower() in {"1", "true", "yes", "on"}
@staticmethod
async def _maybe_rust_anthropic_messages(
*,
@ -2277,30 +2271,21 @@ class BaseLLMHTTPHandler:
) -> AnthropicMessagesResponse | None:
if custom_llm_provider not in ("azure_ai", "anthropic"):
return None
if litellm_params.get("rust") is not True and not BaseLLMHTTPHandler._rust_env_enabled():
return None
if stream and not rust_stream_eligible:
return None
from litellm.rust_bridge import messages as rust_messages_bridge
upstream_body = {key: value for key, value in request_body.items() if key != "stream"}
try:
rust_response = await rust_messages_bridge.amessages(
model=model,
body=upstream_body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
)
except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path
verbose_logger.debug(
"Rust Anthropic messages bridge raised %s; falling back to Python path",
type(rust_error).__name__,
)
return None
rust_response = await rust_messages_bridge.amessages(
model=model,
body=upstream_body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
)
if rust_response is None:
return None
@ -6232,7 +6217,7 @@ class BaseLLMHTTPHandler:
@asynccontextmanager
async def _backend_connection():
if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params):
if _rust_responses_websocket_enabled(custom_llm_provider):
from litellm.rust_bridge import responses_websocket as rust_responses_websocket
rust_backend = await rust_responses_websocket.connect(

4
litellm/ocr/main.py
View file

@ -396,7 +396,7 @@ async def aocr(
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled():
if _rust_ocr_supported(prepared):
from litellm.secret_managers.main import get_secret_str
rust_response = await _run_rust_aocr(
@ -664,7 +664,7 @@ def ocr(
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled():
if _rust_ocr_supported(prepared):
from litellm.secret_managers.main import get_secret_str
rust_response = _run_rust_ocr(

20
litellm/rust_bridge/ocr.py
View file

@ -2,7 +2,6 @@
from __future__ import annotations
import os
from typing import TYPE_CHECKING, Any, Awaitable, Final, Protocol, Union, cast
import httpx
@ -50,16 +49,6 @@ class _Unset:
_UNSET: Final[_Unset] = _Unset()
def _env_enables_rust_ocr() -> bool:
return os.getenv("LITELLM_USE_RUST_OCR", "").strip().lower() in {
"1",
"true",
"yes",
"on",
}
_rust_ocr_enabled = _env_enables_rust_ocr()
_rust_ocr_impl: RustOcr | None = None
_rust_aocr_impl: RustAocr | None = None
@ -75,13 +64,10 @@ def use_litellm_rust(
transcription: Any | None | _Unset = _UNSET,
atranscription: Any | None | _Unset = _UNSET,
) -> None:
global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl
configuring_ocr = not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset)
global _rust_ocr_impl, _rust_aocr_impl
configuring_messages = not isinstance(messages, _Unset) or not isinstance(amessages, _Unset)
configuring_responses_websocket = not isinstance(responses_websocket, _Unset)
configuring_transcription = not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset)
if configuring_ocr or (not configuring_messages and not configuring_responses_websocket):
_rust_ocr_enabled = enabled
if not isinstance(ocr, _Unset):
_rust_ocr_impl = ocr
if not isinstance(aocr, _Unset):
@ -111,10 +97,6 @@ def use_litellm_rust(
set_rust_responses_websocket(connection=responses_websocket)
def rust_ocr_enabled() -> bool:
return _rust_ocr_enabled
def load_rust_ocr() -> RustOcr | None:
if _rust_ocr_impl is not None:
return _rust_ocr_impl

5
tests/documentation_tests/test_env_keys.py
View file

@ -26,10 +26,7 @@ EXCLUDED_GUARD_ONLY_VARS = {
# Temporary/internal rollout flags are intentionally not added to the public
# environment settings docs until the feature is ready for broad use.
EXCLUDED_ROLLOUT_FLAGS = {
"LITELLM_USE_RUST_OCR",
"LITELLM_RUST",
}
EXCLUDED_ROLLOUT_FLAGS: set[str] = set()
EXCLUDED_TERMINAL_VARS = {
"TERM",

82
tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
View file

@ -122,27 +122,6 @@ def test_load_rust_messages_returns_injected_impl():
assert rust_messages.load_rust_messages() is bridge
def test_configuring_messages_does_not_enable_ocr():
from litellm.rust_bridge.ocr import rust_ocr_enabled
litellm.use_litellm_rust(False)
assert rust_ocr_enabled() is False
litellm.use_litellm_rust(True, messages=RecordingMessages())
assert rust_ocr_enabled() is False
def test_bare_use_litellm_rust_still_toggles_ocr():
from litellm.rust_bridge.ocr import rust_ocr_enabled
litellm.use_litellm_rust(True)
assert rust_ocr_enabled() is True
litellm.use_litellm_rust(False)
assert rust_ocr_enabled() is False
def test_load_rust_amessages_returns_injected_impl():
bridge = RecordingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
@ -252,38 +231,15 @@ async def test_gate_invokes_rust_and_marks_response_header():
@pytest.mark.asyncio
async def test_gate_falls_back_to_python_when_bridge_raises():
async def test_gate_surfaces_rust_error_without_falling_back():
bridge = RaisingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
response = await _gate()
assert response is None
with pytest.raises(RuntimeError):
await _gate()
assert bridge.calls == 1
@pytest.mark.asyncio
async def test_gate_skips_rust_when_flag_absent():
bridge = ExplodingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure"))
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_gate_skips_rust_when_flag_false():
bridge = ExplodingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False))
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_gate_invokes_rust_for_native_anthropic_provider():
bridge = RecordingAsyncMessages()
@ -291,7 +247,7 @@ async def test_gate_invokes_rust_for_native_anthropic_provider():
response = await _gate(
custom_llm_provider="anthropic",
litellm_params=GenericLiteLLMParams(api_key="sk-ant", rust=True),
litellm_params=GenericLiteLLMParams(api_key="sk-ant", rust=False),
api_key="sk-ant",
api_base="https://api.anthropic.com",
headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"},
@ -303,36 +259,6 @@ async def test_gate_invokes_rust_for_native_anthropic_provider():
assert bridge.calls[0]["api_key"] == "sk-ant"
@pytest.mark.asyncio
async def test_gate_invokes_rust_when_env_var_set(monkeypatch):
bridge = RecordingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
monkeypatch.setenv("LITELLM_RUST", "1")
response = await _gate(
custom_llm_provider="anthropic",
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
)
assert response is not None
assert bridge.calls[0]["custom_llm_provider"] == "anthropic"
@pytest.mark.asyncio
async def test_gate_env_var_falsey_does_not_enable(monkeypatch):
bridge = ExplodingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
monkeypatch.setenv("LITELLM_RUST", "0")
response = await _gate(
custom_llm_provider="anthropic",
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
)
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_gate_skips_rust_for_unsupported_provider():
bridge = ExplodingAsyncMessages()

16
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -361,6 +361,14 @@ def test_fingerprint_agentic_tools_is_deterministic():
assert handler._fingerprint_agentic_tools(tools_a) == handler._fingerprint_agentic_tools(tools_b)
@pytest.fixture
def force_python_messages_path(monkeypatch):
import litellm.rust_bridge.messages as rust_messages_bridge
monkeypatch.setattr(rust_messages_bridge, "load_rust_amessages", lambda: None)
monkeypatch.setattr(rust_messages_bridge, "load_rust_messages", lambda: None)
@pytest.mark.asyncio
async def test_async_anthropic_messages_handler_extra_headers():
"""
@ -443,7 +451,7 @@ async def test_async_anthropic_messages_handler_extra_headers():
@pytest.mark.asyncio
async def test_async_anthropic_messages_handler_streaming_forwards_provider_response_headers():
async def test_async_anthropic_messages_handler_streaming_forwards_provider_response_headers(force_python_messages_path):
"""
Regression test for LIT-3724 (issue 2): streaming /v1/messages responses
dropped the upstream provider's HTTP response headers, so Bedrock's
@ -1129,7 +1137,7 @@ def test_resolve_anthropic_messages_timeout(
@pytest.mark.asyncio
async def test_async_anthropic_messages_handler_forwards_request_timeout(monkeypatch):
async def test_async_anthropic_messages_handler_forwards_request_timeout(monkeypatch, force_python_messages_path):
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
monkeypatch.setattr(litellm, "callbacks", [])
@ -1177,7 +1185,7 @@ async def test_async_anthropic_messages_handler_forwards_request_timeout(monkeyp
@pytest.mark.asyncio
async def test_async_anthropic_messages_handler_forwards_stream_timeout(monkeypatch):
async def test_async_anthropic_messages_handler_forwards_stream_timeout(monkeypatch, force_python_messages_path):
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
monkeypatch.setattr(litellm, "callbacks", [])
@ -1582,7 +1590,7 @@ def test_async_compact_handler_sends_json_when_not_signed():
@pytest.mark.asyncio
async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks():
async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(force_python_messages_path):
"""
Regression: async_anthropic_messages_handler must inject api_key into the
kwargs dict forwarded to _call_agentic_completion_hooks.

24
tests/test_litellm/ocr/test_rust_bridge.py
View file

@ -237,19 +237,6 @@ def fake_async_bridge():
return bridge
def test_use_litellm_rust_toggles_flag():
assert rust_bridge.rust_ocr_enabled() is False
litellm.use_litellm_rust()
assert rust_bridge.rust_ocr_enabled() is True
litellm.use_litellm_rust(False)
assert rust_bridge.rust_ocr_enabled() is False
def test_env_var_enables_rust_ocr(monkeypatch):
monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
assert rust_bridge._env_enables_rust_ocr() is True
def test_load_rust_ocr_returns_injected_impl():
bridge = RecordingBridge()
litellm.use_litellm_rust(True, ocr=bridge)
@ -769,17 +756,6 @@ def test_ocr_passes_default_request_timeout_to_rust(fake_bridge):
assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout)
def test_ocr_does_not_route_to_rust_when_disabled():
"""With the flag off, the bridge must not be consulted even if an impl exists."""
bridge = RecordingBridge()
litellm.use_litellm_rust(False, ocr=bridge)
assert rust_bridge.rust_ocr_enabled() is False
# The impl stays available for injection, but the disabled flag gates usage,
# so ocr() never reaches the Rust path (asserted via the enabled-path test).
assert bridge.calls == []
def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch):
"""Rust enabled but no bridge available (no injected impl, no compiled wheel):
ocr() must degrade to the Python HTTP handler instead of raising."""

9
tests/test_litellm/responses/test_rust_bridge_websocket.py
View file

@ -4,7 +4,6 @@ import pytest
from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled
from litellm.rust_bridge import responses_websocket
from litellm.types.router import GenericLiteLLMParams
class _FakeNativeConnection:
@ -39,10 +38,10 @@ class _FakeNativeBridge:
return _FakeNativeConnection()
def test_rust_websocket_bridge_is_disabled_without_flag() -> None:
assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams())
assert not _rust_responses_websocket_enabled("anthropic", GenericLiteLLMParams(rust=True))
assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True))
def test_rust_websocket_bridge_enabled_for_openai_only() -> None:
assert _rust_responses_websocket_enabled("openai")
assert not _rust_responses_websocket_enabled("anthropic")
assert not _rust_responses_websocket_enabled(None)
@pytest.mark.asyncio