scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 133

refactor(rust): remove per-request enablement arguments (#39928)

Browse source 
* refactor(rust): remove per-request enablement arguments

* fix(rust): remove ignored transcription enablement

* refactor(rust): remove OCR-specific bridge controls
This commit is contained in:
yujonglee 2026-09-07 10:43:45 -07:00 • committed by GitHub
parent a2b7868a5b
commit 217cb12623
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 109 additions and 357 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -21,13 +21,10 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
}
)
# The per-deployment Rust opt-in.
RUST_KWARG_KEY: Final = "rust"
# Keys `completion()` forwards from its own kwargs into `get_litellm_params`,
# which are otherwise invisible to it because that call site passes explicit
# named arguments rather than `**kwargs`.
FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY})
FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS
# Pre-define optional kwargs keys as frozenset for O(1) lookups
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
@ -58,10 +55,6 @@ OPTIONAL_KWARGS_KEYS: Final = (
"itpm",
"otpm",
"use_xai_oauth",
# The per-deployment Rust opt-in. `all_litellm_params` keeps it out
# of the provider body; this keeps it *in* litellm_params, which is
# where the chat completions handlers read it from.
RUST_KWARG_KEY,
}
)
| AWS_CREDENTIAL_KWARGS_KEYS

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

@ -163,13 +163,10 @@ from litellm.utils import (
def _rust_responses_websocket_enabled(
custom_llm_provider: str | None,
litellm_params: GenericLiteLLMParams,
) -> bool:
from litellm.rust_bridge.configuration import rust_enabled
raw_request_override: Final = litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
return custom_llm_provider == "openai" and rust_enabled(request_override=request_override)
return custom_llm_provider == "openai" and rust_enabled()
from .http_handler import get_shared_realtime_ssl_context
@ -2403,9 +2400,7 @@ class BaseLLMHTTPHandler:
return None
from litellm.rust_bridge.configuration import rust_enabled
raw_request_override: Final = litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
if not rust_enabled(request_override=request_override):
if not rust_enabled():
return None
if has_agentic_hook:
return None
@ -6514,7 +6509,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: Final = await rust_responses_websocket.connect(

11
litellm/ocr/main.py
View file

@ -29,6 +29,7 @@ from litellm.llms.base_llm.ocr.transformation import (
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.configuration import rust_enabled
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -196,12 +197,6 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool:
raw_request_override: Final = prepared_request.litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override)
def _rust_bridge_optional_params(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
@ -430,7 +425,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_enabled(prepared):
if _rust_ocr_supported(prepared) and rust_enabled():
from litellm.secret_managers.main import get_secret_str
rust_response: Final = await _run_rust_aocr(
@ -702,7 +697,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_enabled(prepared):
if _rust_ocr_supported(prepared) and rust_enabled():
from litellm.secret_managers.main import get_secret_str
rust_response: Final = _run_rust_ocr(

3
litellm/rust_bridge/chat_completions.py
View file

@ -247,8 +247,7 @@ def rust_chat_completions_accepts(
return False
if stream:
return False
request_override: Final = litellm_params.get("rust") if litellm_params is not None else None
if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None):
if not rust_enabled():
return False
if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params):
verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path")

33
litellm/rust_bridge/configuration.py
View file

@ -1,13 +1,11 @@
from __future__ import annotations
import os
import warnings
from typing import Final
DEFAULT_RUST_ENABLED: Final = False
_TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"})
_GLOBAL_ENV_NAME: Final = "LITELLM_RUST"
_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR"
class _RustConfiguration:
@ -26,49 +24,24 @@ def _parse_env_bool(value: str | None) -> bool | None:
def resolve_rust_enabled(
*,
request_override: bool | None,
process_override: bool | None,
environment_override: bool | None,
legacy_environment_override: bool | None = None,
release_default: bool = DEFAULT_RUST_ENABLED,
) -> bool:
if request_override is not None:
return request_override
if process_override is not None:
return process_override
if environment_override is not None:
return environment_override
if legacy_environment_override is not None:
return legacy_environment_override
return release_default
def rust_enabled(*, request_override: bool | None = None) -> bool:
if request_override is not None:
return request_override
process_override: Final = _CONFIGURATION.override
if process_override is not None:
return process_override
global_override: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME))
legacy_override: Final = None if global_override is not None else _parse_env_bool(os.getenv(_LEGACY_OCR_ENV_NAME))
if legacy_override is not None:
warnings.warn(
f"{_LEGACY_OCR_ENV_NAME} is deprecated; use {_GLOBAL_ENV_NAME} instead",
DeprecationWarning,
stacklevel=2,
)
def rust_enabled() -> bool:
return resolve_rust_enabled(
request_override=None,
process_override=None,
environment_override=global_override,
legacy_environment_override=legacy_override,
process_override=_CONFIGURATION.override,
environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)),
)
def rust_ocr_enabled(*, request_override: bool | None = None) -> bool:
return rust_enabled(request_override=request_override)
def reset_rust_configuration() -> None:
_CONFIGURATION.override = None

46
litellm/rust_bridge/ocr.py
View file

@ -7,12 +7,9 @@ from typing import Final, Protocol, cast # noqa: TID251 # native extension exp
import httpx
from litellm.rust_bridge import configuration as _configuration
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
rust_ocr_enabled = _configuration.rust_ocr_enabled
rust = _configuration.rust
class RustOcr(Protocol):
def __call__(
@ -44,49 +41,24 @@ class RustAocr(Protocol):
raise NotImplementedError
class _Unset:
pass
def _as_ocr(value: object) -> RustOcr | None:
return cast(RustOcr, value) if callable(value) else None
_UNSET: Final[_Unset] = _Unset()
def _as_aocr(value: object) -> RustAocr | None:
return cast(RustAocr, value) if callable(value) else None
_rust_ocr_impl: RustOcr | None = None
_rust_aocr_impl: RustAocr | None = None
def set_rust_ocr(
*,
ocr: RustOcr | None | _Unset = _UNSET,
aocr: RustAocr | None | _Unset = _UNSET,
) -> None:
global _rust_ocr_impl, _rust_aocr_impl
if not isinstance(ocr, _Unset):
_rust_ocr_impl = ocr
if not isinstance(aocr, _Unset):
_rust_aocr_impl = aocr
_OCR: Final = NativeBinding("ocr", validate=_as_ocr)
_AOCR: Final = NativeBinding("aocr", validate=_as_aocr)
def load_rust_ocr() -> RustOcr | None:
if _rust_ocr_impl is not None:
return _rust_ocr_impl
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustOcr, native_bridge.ocr)
return _OCR.load()
def load_rust_aocr() -> RustAocr | None:
if _rust_aocr_impl is not None:
return _rust_aocr_impl
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustAocr, getattr(native_bridge, "aocr", None))
return _AOCR.load()
def ocr(

1
litellm/rust_bridge/transcription.py
View file

@ -56,7 +56,6 @@ _STATE: Final = _RustTranscriptionState()
def configure_rust_transcription(
enabled: bool = True,
*,
transcription: RustTranscription | None | _Unset = _UNSET,
atranscription: RustAtranscription | None | _Unset = _UNSET,

1
litellm/types/router.py
View file

@ -307,7 +307,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
"""
custom_llm_provider: str | None = None
rust: bool | None = None
tpm: int | None = None
rpm: int | None = None
itpm: int | None = None

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

@ -126,16 +126,6 @@ def test_load_rust_messages_returns_injected_impl():
assert rust_messages.load_rust_messages() is bridge
def test_bare_rust_still_toggles_ocr():
from litellm.rust_bridge.ocr import rust_ocr_enabled
litellm.rust(True)
assert rust_ocr_enabled() is True
litellm.rust(False)
assert rust_ocr_enabled() is False
def test_load_rust_amessages_returns_injected_impl():
bridge = RecordingAsyncMessages()
litellm.rust(True)
@ -214,7 +204,7 @@ async def test_amessages_wrapper_forwards_args():
def _gate(**overrides):
kwargs = {
"custom_llm_provider": "azure_ai",
"litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True),
"litellm_params": GenericLiteLLMParams(api_key="sk-azure"),
"has_agentic_hook": False,
"model": "claude-sonnet-4-5",
"api_key": "sk-azure",
@ -282,18 +272,6 @@ async def test_gate_uses_process_enable_without_request_override():
assert bridge.calls[0]["custom_llm_provider"] == "azure_ai"
@pytest.mark.asyncio
async def test_gate_skips_rust_when_flag_false():
bridge = ExplodingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(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()
@ -302,7 +280,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"),
api_key="sk-ant",
api_base="https://api.anthropic.com",
headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"},

29
tests/test_litellm/litellm_core_utils/test_get_litellm_params.py
View file

@ -215,32 +215,3 @@ class TestMetadataFallsBackToLitellmMetadata:
assert result["metadata"] is not litellm_metadata
result["metadata"].pop("trace_id")
assert litellm_metadata == {"trace_id": "trace-1"}
class TestRustOptIn:
"""`rust: true` is a litellm param, so it has to reach `litellm_params`.
`all_litellm_params` keeps it out of the provider body; without it also
being carried into `litellm_params` the chat completions handlers cannot
see the opt-in and the Rust path is silently never taken.
"""
def test_rust_is_an_optional_kwargs_key(self):
assert "rust" in _OPTIONAL_KWARGS_KEYS
def test_rust_is_forwarded_from_completion_kwargs(self):
from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
assert "rust" in FORWARDED_KWARGS_KEYS
def test_rust_survives_into_litellm_params(self):
params = get_litellm_params(rust=True)
assert params["rust"] is True
def test_rust_is_absent_when_the_deployment_did_not_set_it(self):
assert "rust" not in get_litellm_params()
def test_rust_stays_out_of_the_provider_body(self):
from litellm.types.utils import all_litellm_params
assert "rust" in all_litellm_params

8
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
View file

@ -2256,7 +2256,7 @@ class TestRustChatCompletionsHook:
def _reset_bridge(self, monkeypatch):
from litellm.rust_bridge import chat_completions as bridge
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
@ -2282,7 +2282,7 @@ class TestRustChatCompletionsHook:
"logging_obj": MagicMock(),
"optional_params": {"max_tokens": 16},
"timeout": 30.0,
"litellm_params": {"rust": True},
"litellm_params": {},
"acompletion": False,
"headers": {},
"client": None,
@ -2366,7 +2366,8 @@ class TestRustChatCompletionsHook:
)
assert seen["call"][0]["optional_params"]["max_tokens"] == 7
def test_without_the_opt_in_the_core_is_never_consulted(self):
def test_without_the_opt_in_the_core_is_never_consulted(self, monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "0")
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@ -2587,6 +2588,7 @@ class TestRustChatCompletionsHook:
def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch):
"""The suppression must not swallow the log on the ordinary path."""
monkeypatch.setenv("LITELLM_RUST", "0")
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig

12
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -48,7 +48,7 @@ RESOLVED_CREDENTIALS = Credentials(
@pytest.fixture(autouse=True)
def reset_bridge(monkeypatch):
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
@ -87,7 +87,7 @@ def _completion_kwargs(**overrides):
"optional_params": {"maxTokens": 16},
"acompletion": False,
"timeout": 30.0,
"litellm_params": {"rust": True},
"litellm_params": {},
"extra_headers": None,
"client": None,
"api_key": None,
@ -157,7 +157,8 @@ def test_the_core_receives_the_untranslated_openai_messages():
]
def test_without_the_opt_in_the_core_is_never_consulted():
def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "0")
seen = _inject()
try:
_run(litellm_params={})
@ -401,9 +402,10 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines():
assert logging_obj.pre_call.call_count == 1
def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in():
def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch):
"""The suppression must not swallow the log on a request the gate declined,
so a deployment with no `rust` flag keeps exactly the log it always had."""
monkeypatch.setenv("LITELLM_RUST", "0")
logging_obj = MagicMock()
response = _run(
logging_obj=logging_obj,
@ -491,6 +493,7 @@ def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monke
"""With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no
credentials at all. Preparing the Rust handoff must not dereference that
None: the bearer token signs the request on its own."""
monkeypatch.setenv("LITELLM_RUST", "0")
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
client = _sync_client_returning_converse_response()
@ -520,6 +523,7 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
raises; a bearer-token deployment must still serve the request, since the
bearer token alone signs it."""
monkeypatch.setenv("LITELLM_RUST", "0")
if configured_through == "env_var":
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
else:

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

@ -2689,20 +2689,18 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h
@pytest.mark.parametrize(
"custom_llm_provider, litellm_params, expected",
[
("openai", GenericLiteLLMParams(rust=True), True),
("openai", GenericLiteLLMParams(), False),
("openai", GenericLiteLLMParams(rust=False), False),
("azure", GenericLiteLLMParams(rust=True), False),
("hosted_vllm", GenericLiteLLMParams(rust=True), False),
(None, GenericLiteLLMParams(rust=True), False),
],
"custom_llm_provider, enabled, expected",
[("openai", True, True), ("openai", False, False), ("azure", True, False),
("hosted_vllm", True, False), (None, True, False)],
)
def test_the_rust_responses_websocket_needs_both_openai_and_the_rust_flag(
custom_llm_provider, litellm_params, expected
def test_the_rust_responses_websocket_needs_openai_and_process_enablement(
custom_llm_provider, enabled, expected, monkeypatch
):
assert _rust_responses_websocket_enabled(custom_llm_provider, litellm_params) is expected
from litellm.rust_bridge import configuration
configuration.reset_rust_configuration()
monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
assert _rust_responses_websocket_enabled(custom_llm_provider) is expected
def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch):

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

@ -17,6 +17,7 @@ from litellm.rust_bridge import configuration
# explicitly via importlib rather than attribute traversal.
ocr_main = importlib.import_module("litellm.ocr.main")
rust_bridge = importlib.import_module("litellm.rust_bridge.ocr")
rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings")
rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
MODEL = "mistral/mistral-ocr-latest"
@ -215,11 +216,13 @@ def build_prepared_request(
@pytest.fixture(autouse=True)
def _reset_rust_flag():
"""Keep the global toggle isolated between tests."""
rust_bridge.set_rust_ocr(ocr=None, aocr=None)
rust_bridge._OCR.reset()
rust_bridge._AOCR.reset()
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
yield
rust_bridge.set_rust_ocr(ocr=None, aocr=None)
rust_bridge._OCR.reset()
rust_bridge._AOCR.reset()
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
@ -229,7 +232,7 @@ def fake_bridge():
"""Enable the Rust path with an injected recording bridge (no native wheel)."""
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
return bridge
@ -238,34 +241,14 @@ def fake_async_bridge():
"""Enable the async Rust path with an injected recording bridge."""
bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(aocr=bridge)
rust_bridge._AOCR.override(bridge)
return bridge
def test_rust_toggles_flag():
assert rust_bridge.rust_ocr_enabled() is False
litellm.rust(True)
assert rust_bridge.rust_ocr_enabled() is True
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")
with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"):
assert rust_bridge.rust_ocr_enabled() is True
def test_explicit_false_overrides_process_enable():
litellm.rust(True)
assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False
def test_load_rust_ocr_returns_injected_impl():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
assert rust_bridge.load_rust_ocr() is bridge
@ -329,7 +312,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch):
def test_load_rust_aocr_returns_injected_impl():
bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(aocr=bridge)
rust_bridge._AOCR.override(bridge)
assert rust_bridge.load_rust_aocr() is bridge
@ -338,7 +321,8 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
bridge = RecordingBridge()
async_bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge)
rust_bridge._OCR.override(bridge)
rust_bridge._AOCR.override(async_bridge)
litellm.rust(False)
assert rust_bridge.load_rust_ocr() is bridge
@ -350,16 +334,18 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
rust_bridge_bindings,
"get_native_bridge",
lambda: None,
)
bridge = RecordingBridge()
async_bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge)
rust_bridge._OCR.override(bridge)
rust_bridge._AOCR.override(async_bridge)
rust_bridge.set_rust_ocr(ocr=None, aocr=None)
rust_bridge._OCR.override(None)
rust_bridge._AOCR.override(None)
assert rust_bridge.load_rust_ocr() is None
assert rust_bridge.load_rust_aocr() is None
@ -368,7 +354,7 @@ def test_load_rust_ocr_none_when_extension_absent(monkeypatch):
"""With no injected impl and no compiled wheel, the loader returns None so the
caller degrades to the Python path instead of raising ImportError."""
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
rust_bridge_bindings,
"get_native_bridge",
lambda: None,
)
@ -385,7 +371,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
rust_bridge_bindings,
"get_native_bridge",
lambda: fake_module,
)
@ -406,7 +392,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response():
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
response = rust_bridge.ocr(
model="mistral-ocr-latest",
document=DOCUMENT,
@ -441,7 +427,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response():
litellm.rust(True)
rust_bridge.set_rust_ocr(aocr=bridge)
rust_bridge._AOCR.override(bridge)
response = await rust_bridge.aocr(
model="mistral-ocr-maas",
document=DOCUMENT,
@ -470,7 +456,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
bridge = RecordingBridge()
logging_obj = RecordingLogging()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
response = ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -503,7 +489,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(api_key=None, timeout=None),
@ -516,7 +502,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
def test_run_rust_ocr_prefers_explicit_key_over_resolver():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
def _resolver(name: str) -> str | None:
raise AssertionError(f"resolver should not be called for {name}")
@ -536,7 +522,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
bridge = RecordingBridge()
resolver_calls = []
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
def _resolver(name):
resolver_calls.append(name)
@ -559,7 +545,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -586,7 +572,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
def _resolver(name: str) -> str | None:
return {
@ -610,7 +596,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana
def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -628,7 +614,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -649,7 +635,7 @@ def test_run_rust_ocr_runs_pre_call_logging():
logging_obj = RecordingLogging()
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge._OCR.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -737,7 +723,7 @@ def test_ocr_exception_type_uses_resolved_provider_context(
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=RaisingBridge())
rust_bridge._OCR.override(RaisingBridge())
with pytest.raises(CapturedException):
litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
@ -783,7 +769,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context(
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
litellm.rust(True)
rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge())
rust_bridge._AOCR.override(RaisingAsyncBridge())
with pytest.raises(CapturedException):
await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
@ -812,9 +798,7 @@ 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.rust(False)
rust_bridge.set_rust_ocr(ocr=bridge)
assert rust_bridge.rust_ocr_enabled() is False
rust_bridge._OCR.override(bridge)
# 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 == []

21
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 configuration, responses_websocket
from litellm.types.router import GenericLiteLLMParams
class _FakeNativeConnection:
@ -48,22 +47,12 @@ def reset_responses_websocket():
configuration.reset_rust_configuration()
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_explicit_false_overrides_process_enable() -> None:
def test_rust_websocket_bridge_uses_process_enablement() -> None:
configuration.rust(False)
assert not _rust_responses_websocket_enabled("openai")
configuration.rust(True)
assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False))
def test_process_enable_applies_without_request_override() -> None:
configuration.rust(True)
assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams())
assert _rust_responses_websocket_enabled("openai")
assert not _rust_responses_websocket_enabled("anthropic")
@pytest.mark.asyncio

38
tests/test_litellm/rust_bridge/test_chat_completions.py
View file

@ -10,8 +10,8 @@ from __future__ import annotations
import pytest
import litellm
from litellm.rust_bridge import chat_completions as bridge
from litellm.rust_bridge import configuration
from litellm.rust_bridge import chat_completions as bridge
from litellm.types.utils import ModelResponse
RUST_RESPONSE = {
@ -67,10 +67,11 @@ def _hide_native_bridge(monkeypatch):
@pytest.fixture(autouse=True)
def reset_bridge():
def reset_bridge(monkeypatch):
"""Every test starts with no injected callables, and leaves none behind."""
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
configuration.reset_rust_configuration()
monkeypatch.setenv("LITELLM_RUST", "1")
yield
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
configuration.reset_rust_configuration()
@ -112,7 +113,7 @@ def _accepts(**overrides) -> bool:
"messages": MESSAGES,
"optional_params": {"max_tokens": 16},
"custom_llm_provider": "anthropic",
"litellm_params": {"rust": True},
"litellm_params": {},
"stream": None,
}
kwargs.update(overrides)
@ -126,23 +127,16 @@ class TestGate:
bridge.set_rust_chat_completions(decline=gate)
assert _accepts(litellm_params={}) is False
assert _accepts(litellm_params=None) is False
assert _accepts(litellm_params={"rust": False}) is False
assert gate.calls == [], "the gate must not be consulted before opt-in"
def test_accepts_when_the_deployment_opted_in_and_the_core_agrees(self, monkeypatch):
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
assert _accepts() is True
assert gate.calls[0]["model"] == "claude-sonnet-4-5"
assert gate.calls[0]["custom_llm_provider"] == "anthropic"
def test_explicit_false_overrides_process_enable(self):
bridge.set_rust_chat_completions(decline=_RecordingDecline())
configuration.rust(True)
assert _accepts(litellm_params={"rust": False}) is False
def test_process_enable_applies_without_request_override(self):
bridge.set_rust_chat_completions(decline=_RecordingDecline())
configuration.rust(True)
@ -155,7 +149,7 @@ class TestGate:
assert _accepts(litellm_params={}) is True
def test_declines_streaming_and_providers_off_the_path(self, monkeypatch):
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
assert _accepts(stream=True) is False
@ -170,10 +164,10 @@ class TestGate:
handed `optional_params` only, so accepting here would send the request
to Anthropic with the abuse-detection attribution silently missing.
"""
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}) is False
assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False
assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of"
# Bedrock's Converse transform reads no `user_id`, and an Anthropic request
@ -182,13 +176,13 @@ class TestGate:
_accepts(
custom_llm_provider="bedrock",
model="bedrock/us-east-1/anthropic.claude-v2",
litellm_params={"rust": True, "metadata": {"user_id": "u-123"}},
litellm_params={"metadata": {"user_id": "u-123"}},
)
is True
)
assert _accepts(litellm_params={"rust": True, "metadata": {"trace_id": "t-1"}}) is True
assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": None}}) is True
assert _accepts(litellm_params={"rust": True, "metadata": None}) is True
assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True
assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True
assert _accepts(litellm_params={"metadata": None}) is True
def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch):
"""`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the
@ -196,7 +190,7 @@ class TestGate:
evicting a caller-supplied one. The core can do neither, so an operator
who armed `bedrock_request_metadata_fields` keeps the Python path.
"""
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
bedrock = {
@ -213,17 +207,17 @@ class TestGate:
assert _accepts(**bedrock) is True, "the decline follows the operator's opt-in alone"
def test_declines_when_the_core_declines(self, monkeypatch):
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
bridge.set_rust_chat_completions(decline=_RecordingDecline("streaming"))
assert _accepts() is False
def test_declines_when_the_bridge_is_unavailable(self, monkeypatch):
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
_hide_native_bridge(monkeypatch)
assert _accepts() is False
def test_declines_when_the_gate_itself_raises(self, monkeypatch):
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1")
def exploding(**_kwargs):
raise RuntimeError("boom")

61
tests/test_litellm/rust_bridge/test_configuration.py
View file

@ -10,7 +10,6 @@ from typing import Final
import pytest
from litellm.rust_bridge import configuration
from litellm.rust_bridge import ocr as rust_ocr
@pytest.fixture(autouse=True)
@ -19,42 +18,31 @@ def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest
) -> Generator[None]:
configuration.reset_rust_configuration()
monkeypatch.delenv("LITELLM_RUST", raising=False)
monkeypatch.delenv("LITELLM_USE_RUST_OCR", raising=False)
rust_ocr.set_rust_ocr(ocr=None, aocr=None)
yield
configuration.reset_rust_configuration()
rust_ocr.set_rust_ocr(ocr=None, aocr=None)
@pytest.mark.parametrize(
("request_override", "process", "environment", "legacy_environment", "release_default", "expected"),
("process", "environment", "release_default", "expected"),
(
(False, True, True, True, True, False),
(True, False, False, False, False, True),
(None, False, True, True, True, False),
(None, True, False, False, False, True),
(None, None, False, True, True, False),
(None, None, True, False, False, True),
(None, None, None, False, True, False),
(None, None, None, True, False, True),
(None, None, None, None, False, False),
(None, None, None, None, True, True),
(False, True, True, False),
(True, False, False, True),
(None, False, True, False),
(None, True, False, True),
(None, None, False, False),
(None, None, True, True),
),
)
def test_resolution_precedence(
request_override: bool | None,
process: bool | None,
environment: bool | None,
legacy_environment: bool | None,
release_default: bool,
expected: bool,
) -> None:
assert (
configuration.resolve_rust_enabled(
request_override=request_override,
process_override=process,
environment_override=environment,
legacy_environment_override=legacy_environment,
release_default=release_default,
)
is expected
@ -71,7 +59,6 @@ def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch)
configuration.rust(True)
assert configuration.rust_enabled() is True
assert configuration.rust_enabled(request_override=False) is False
def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPatch) -> None:
@ -83,18 +70,8 @@ def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPat
@pytest.mark.parametrize("value", ("", " ", "sometimes", "2"))
def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
monkeypatch.setenv("LITELLM_RUST", value)
monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
assert configuration.rust_enabled() is False
assert configuration.rust_ocr_enabled() is False
@pytest.mark.parametrize("value", ("", " ", "sometimes", "2"))
def test_invalid_legacy_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
monkeypatch.setenv("LITELLM_USE_RUST_OCR", value)
with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"):
assert configuration.rust_enabled() is False
def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None:
@ -104,40 +81,20 @@ def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytes
assert executor.submit(configuration.rust_enabled).result() is True
configuration.rust(False)
assert executor.submit(configuration.rust_enabled).result() is False
assert executor.submit(configuration.rust_ocr_enabled).result() is False
configuration.reset_rust_configuration()
assert executor.submit(configuration.rust_enabled).result() is True
assert executor.submit(configuration.rust_ocr_enabled).result() is True
def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_RUST", "sometimes")
assert configuration.rust_enabled(request_override=False) is False
configuration.rust(True)
assert configuration.rust_enabled() is True
def test_legacy_ocr_environment_is_deprecated_and_global(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"):
assert configuration.rust_enabled() is True
with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"):
assert configuration.rust_ocr_enabled() is True
def test_global_environment_precedes_legacy_ocr_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
assert configuration.rust_enabled() is False
@pytest.mark.parametrize("environment_name", ("LITELLM_RUST", "LITELLM_USE_RUST_OCR"))
@pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False")))
def test_environment_controls_startup(environment_name: str, value: str, expected: str) -> None:
environment: Final = {**os.environ, environment_name: value}
def test_environment_controls_startup(value: str, expected: str) -> None:
environment: Final = {**os.environ, "LITELLM_RUST": value}
result: Final = subprocess.run(
(
sys.executable,

4
tests/test_litellm/test_audio_transcription_rust_bridge.py
View file

@ -44,7 +44,7 @@ class AsyncBridge:
def test_enabled_sync_bridge_receives_audio() -> None:
bridge = SyncBridge()
rust_bridge.configure_rust_transcription(True, transcription=bridge)
rust_bridge.configure_rust_transcription(transcription=bridge)
result = rust_bridge.transcription(
model="mistral.voxtral-mini-3b-2507",
audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"},
@ -61,7 +61,7 @@ def test_enabled_sync_bridge_receives_audio() -> None:
@pytest.mark.asyncio
async def test_enabled_async_bridge() -> None:
rust_bridge.configure_rust_transcription(True, atranscription=AsyncBridge())
rust_bridge.configure_rust_transcription(atranscription=AsyncBridge())
result = await rust_bridge.atranscription(
model="mistral.voxtral-mini-3b-2507",
audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"},

46
tests/test_litellm/test_utils.py
View file

@ -5239,52 +5239,6 @@ def test_client_side_timeout_marker_never_reaches_the_provider():
)
def test_rust_flag_not_forwarded_as_provider_param():
forwarded = get_non_default_completion_params({"rust": True, "temperature": 0.5})
assert "rust" not in forwarded
def test_completion_does_not_leak_rust_flag_into_provider_request_body():
mock_response = MagicMock()
mock_response.model_dump.return_value = {
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 1234567890,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 1,
"completion_tokens": 1,
"total_tokens": 2,
},
}
mock_raw_response = MagicMock()
mock_raw_response.headers = {}
mock_raw_response.parse.return_value = mock_response
mock_client = MagicMock()
mock_client.chat.completions.with_raw_response.create.return_value = mock_raw_response
litellm.completion(
model="openai/gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
rust=True,
api_key="sk-test",
client=mock_client,
)
create_kwargs = mock_client.chat.completions.with_raw_response.create.call_args.kwargs
assert "rust" not in create_kwargs
assert "rust" not in (create_kwargs.get("extra_body") or {})
class _RecordingDeploymentFailureLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()

4
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -29562,8 +29562,6 @@ export interface components {
regional_processing_uplift_multiplier_us?: number | null;
/** Rpm */
rpm?: number | null;
/** Rust */
rust?: boolean | null;
/** S3 Bucket Name */
s3_bucket_name?: string | null;
/** S3 Encryption Key Id */
@ -39729,8 +39727,6 @@ export interface components {
regional_processing_uplift_multiplier_us?: number | null;
/** Rpm */
rpm?: number | null;
/** Rust */
rust?: boolean | null;
/** S3 Bucket Name */
s3_bucket_name?: string | null;
/** S3 Encryption Key Id */