diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c48d75439a7..a98769cc2d9 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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( diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 38f3f804e10..823d65f18bd 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -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( diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 91aac6c1232..fc6063600ed 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -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 diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 3bf2c88a848..14f36b39885 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -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", diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index b745ca8eadf..2b24c960e6f 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -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() diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index fddd8d09dfc..667ce24e63e 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -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. diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index acad249a2bb..7f0d4782a99 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -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.""" diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index c9a5b988be6..cc94220b3af 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -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