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

refactor(python): centralize Rust dispatch policy

Browse source
This commit is contained in:
Yujong Lee 2026-09-04 13:10:49 -07:00
parent df68edca76
commit e312c50542
21 changed files with 1093 additions and 733 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/anthropic/chat/handler.py
View file

@ -465,6 +465,7 @@ class AnthropicChatCompletion(BaseLLM):
timeout=timeout,
on_response=log_rust_post_call,
python_fallback=python_fallback,
request_override=rust_chat_completions_bridge.rust_request_override(litellm_params),
)
rust_response: Final = rust_chat_completions_bridge.chat_completions(
model=model,
@ -477,6 +478,7 @@ class AnthropicChatCompletion(BaseLLM):
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
request_override=rust_chat_completions_bridge.rust_request_override(litellm_params),
)
if rust_response is not None:
return rust_response

4
litellm/llms/bedrock/audio_transcription/__init__.py
View file

@ -53,8 +53,6 @@ class BedrockAudioTranscriptionRustDispatch:
optional_params=optional_params,
timeout=timeout,
)
if rust_response is None:
raise RuntimeError("Rust audio transcription bridge is unavailable")
return TranscriptionResponse(**rust_response)
async def async_audio_transcriptions(
@ -79,6 +77,4 @@ class BedrockAudioTranscriptionRustDispatch:
optional_params=optional_params,
timeout=timeout,
)
if rust_response is None:
raise RuntimeError("Rust audio transcription bridge is unavailable")
return TranscriptionResponse(**rust_response)

2
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -452,6 +452,7 @@ class BedrockConverseLLM(BaseAWSLLM):
api_key=api_key,
skip_pre_call_logging=True,
),
request_override=rust_chat_completions_bridge.rust_request_override(litellm_params),
)
rust_response: Final = rust_chat_completions_bridge.chat_completions(
model=model,
@ -464,6 +465,7 @@ class BedrockConverseLLM(BaseAWSLLM):
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
request_override=rust_chat_completions_bridge.rust_request_override(litellm_params),
)
if rust_response is not None:
return rust_response

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

@ -2401,34 +2401,24 @@ class BaseLLMHTTPHandler:
) -> AnthropicMessagesResponse | None:
if custom_llm_provider not in ("azure_ai", "anthropic"):
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):
return None
if has_agentic_hook:
return None
from litellm.rust_bridge import messages as rust_messages_bridge
upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"}
try:
rust_response: Final = 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: Final = 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,
request_override=request_override,
)
if rust_response is None:
return None
@ -6517,10 +6507,12 @@ class BaseLLMHTTPHandler:
if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params):
from litellm.rust_bridge import responses_websocket as rust_responses_websocket
raw_rust_override: Final = litellm_params.get("rust")
rust_backend: Final = await rust_responses_websocket.connect(
url=ws_url,
headers={str(key): str(value) for key, value in headers.items()},
timeout=timeout,
request_override=raw_rust_override if isinstance(raw_rust_override, bool) else None,
)
if rust_backend is not None:
yield rust_backend

97
litellm/ocr/main.py
View file

@ -52,14 +52,6 @@ class _PreparedOCRRequest:
litellm_logging_obj: LiteLLMLoggingObj
@dataclass
class _PreparedRustOCRCall:
api_key: str | None
api_base: str | None
headers: dict[str, object]
optional_params: dict[str, object]
_RUST_OCR_PROVIDERS: Final = {
"mistral",
"azure_ai",
@ -194,10 +186,9 @@ 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:
def _rust_request_override(prepared_request: _PreparedOCRRequest) -> bool | None:
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)
return raw_request_override if isinstance(raw_request_override, bool) else None
def _rust_bridge_optional_params(
@ -242,7 +233,7 @@ def _rust_bridge_api_base(
def _prepare_rust_ocr_call(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> _PreparedRustOCRCall:
) -> rust_ocr_bridge.RustOCRRequest:
provider_config: Final = prepared_request.provider_config
api_key_env_var: Final = provider_config.get_api_key_env_var()
resolved_api_key: Final = prepared_request.api_key or (
@ -276,11 +267,15 @@ def _prepare_rust_ocr_call(
"headers": resolved_headers,
},
)
return _PreparedRustOCRCall(
return rust_ocr_bridge.RustOCRRequest(
model=prepared_request.model,
document=prepared_request.document,
api_key=resolved_api_key,
api_base=rust_api_base,
headers=cast(dict[str, object], resolved_headers),
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=cast(dict[str, object], resolved_headers),
optional_params=rust_optional_params,
timeout=prepared_request.effective_timeout,
)
@ -288,21 +283,15 @@ def _run_rust_ocr(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> OCRResponse | None:
if rust_ocr_bridge.load_rust_ocr() is None:
return None
prepared: Final = _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
)
rust_response: Final = rust_ocr_bridge.ocr(
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
),
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
provider=prepared_request.custom_llm_provider,
request_override=_rust_request_override(prepared_request),
eligible=_rust_ocr_supported(prepared_request),
)
if rust_response is None:
return None
@ -313,21 +302,15 @@ async def _run_rust_aocr(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> OCRResponse | None:
if rust_ocr_bridge.load_rust_aocr() is None:
return None
prepared: Final = _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
)
rust_response: Final = await rust_ocr_bridge.aocr(
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
),
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
provider=prepared_request.custom_llm_provider,
request_override=_rust_request_override(prepared_request),
eligible=_rust_ocr_supported(prepared_request),
)
if rust_response is None:
return None
@ -428,17 +411,14 @@ 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):
from litellm.secret_managers.main import get_secret_str
from litellm.secret_managers.main import get_secret_str
rust_response: Final = await _run_rust_aocr(
prepared_request=prepared,
resolve_api_key=get_secret_str,
)
if rust_response is None:
verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path")
else:
return rust_response
rust_response: Final = await _run_rust_aocr(
prepared_request=prepared,
resolve_api_key=get_secret_str,
)
if rust_response is not None:
return rust_response
response = base_llm_http_handler.ocr(
model=prepared.model,
@ -700,17 +680,14 @@ 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):
from litellm.secret_managers.main import get_secret_str
from litellm.secret_managers.main import get_secret_str
rust_response: Final = _run_rust_ocr(
prepared_request=prepared,
resolve_api_key=get_secret_str,
)
if rust_response is None:
verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path")
else:
return rust_response
rust_response: Final = _run_rust_ocr(
prepared_request=prepared,
resolve_api_key=get_secret_str,
)
if rust_response is not None:
return rust_response
response: Final = base_llm_http_handler.ocr(
model=prepared.model,

16
litellm/rust_bridge/bindings.py
View file

@ -1,7 +1,7 @@
from __future__ import annotations
from collections.abc import Callable
from typing import Final, Generic, TypeVar
from typing import Final, Generic, TypeVar, cast
from litellm.rust_bridge.loader import get_native_bridge
@ -15,6 +15,13 @@ class _Unset:
_UNSET: Final = _Unset()
class Unchanged:
pass
UNCHANGED: Final = Unchanged()
class NativeBinding(Generic[BindingT]):
"""Resolve one native attribute with an explicit, resettable test override."""
@ -23,6 +30,13 @@ class NativeBinding(Generic[BindingT]):
self._validate: Final = validate
self._override: BindingT | None | _Unset = _UNSET
@classmethod
def callable(cls, attribute: str) -> NativeBinding[BindingT]:
return cls(
attribute,
validate=lambda value: cast(BindingT, value) if callable(value) else None,
)
def load(self) -> BindingT | None:
if not isinstance(self._override, _Unset):
return self._override

254
litellm/rust_bridge/chat_completions.py
View file

@ -14,20 +14,24 @@ from __future__ import annotations
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.exceptions import APIError
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_model_response_object,
)
from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
RustBridge,
RustEndpoint,
async_none,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.utils import ModelResponse
@ -126,67 +130,42 @@ def response_logger(
return log
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustChatCompletionsState:
chat_completions: RustChatCompletions | None = None
achat_completions: RustAchatCompletions | None = None
decline: RustChatCompletionsDecline | None = None
_STATE: Final[_RustChatCompletionsState] = _RustChatCompletionsState()
_CHAT: Final[RustEndpoint[RustChatCompletions, RustAchatCompletions]] = RustEndpoint.native(
route="chat completions",
sync="chat_completions",
asynchronous="achat_completions",
enabled=rust_enabled,
)
_CHAT_PREFLIGHT: Final[RustBridge[RustChatCompletionsDecline]] = RustBridge.native(
route="chat completions preflight",
attribute="chat_completions_decline",
enabled=rust_enabled,
)
def set_rust_chat_completions(
*,
chat_completions: RustChatCompletions | None | _Unset = _UNSET,
achat_completions: RustAchatCompletions | None | _Unset = _UNSET,
decline: RustChatCompletionsDecline | None | _Unset = _UNSET,
chat_completions: RustChatCompletions | None | Unchanged = UNCHANGED,
achat_completions: RustAchatCompletions | None | Unchanged = UNCHANGED,
decline: RustChatCompletionsDecline | None | Unchanged = UNCHANGED,
) -> None:
"""Inject the native callables, so tests can supply a double instead of
patching module attributes."""
if not isinstance(chat_completions, _Unset):
_STATE.chat_completions = chat_completions
if not isinstance(achat_completions, _Unset):
_STATE.achat_completions = achat_completions
if not isinstance(decline, _Unset):
_STATE.decline = decline
def load_rust_chat_completions() -> RustChatCompletions | None:
if _STATE.chat_completions is not None:
return _STATE.chat_completions
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustChatCompletions | None = getattr(native_bridge, "chat_completions", None)
return loaded
def load_rust_achat_completions() -> RustAchatCompletions | None:
if _STATE.achat_completions is not None:
return _STATE.achat_completions
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustAchatCompletions | None = getattr(native_bridge, "achat_completions", None)
return loaded
def _load_rust_decline() -> RustChatCompletionsDecline | None:
if _STATE.decline is not None:
return _STATE.decline
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustChatCompletionsDecline | None = getattr(native_bridge, "chat_completions_decline", None)
return loaded
if not isinstance(chat_completions, Unchanged):
if chat_completions is None:
_CHAT.sync.reset()
else:
_CHAT.sync.override(chat_completions)
if not isinstance(achat_completions, Unchanged):
if achat_completions is None:
_CHAT.asynchronous.reset()
else:
_CHAT.asynchronous.override(achat_completions)
if not isinstance(decline, Unchanged):
if decline is None:
_CHAT_PREFLIGHT.reset()
else:
_CHAT_PREFLIGHT.override(decline)
def _anthropic_user_id_reaches_the_body(litellm_params: Mapping[str, object] | None) -> bool:
@ -227,6 +206,11 @@ def _litellm_metadata_reaches_the_provider(
return False
def rust_request_override(litellm_params: Mapping[str, object] | None) -> bool | None:
raw_request_override: Final = litellm_params.get("rust") if litellm_params is not None else None
return raw_request_override if isinstance(raw_request_override, bool) else None
def rust_chat_completions_accepts(
*,
model: str,
@ -247,82 +231,17 @@ 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):
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")
return False
decline: Final = _load_rust_decline()
if decline is None:
return False
try:
reason: Final = decline(
return _CHAT_PREFLIGHT.accepts(
request_override=rust_request_override(litellm_params),
check=lambda decline: decline(
model=model,
messages=messages,
optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
)
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 chat completions gate raised %s; staying on the Python path",
type(rust_error).__name__,
)
return False
if reason is not None:
verbose_logger.debug("Rust chat completions declined (%s); using the Python path", reason)
return False
return True
def _rust_bridge_exceptions() -> tuple[type[BaseException], type[BaseException]] | None:
"""`(declined, upstream_failed)` from the native module, or None when absent."""
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
declined: Final = getattr(native_bridge, "RustBridgeDeclined", None)
upstream: Final = getattr(native_bridge, "RustUpstreamError", None)
if declined is None or upstream is None:
return None
return declined, upstream
def _reraise_or_decline(
rust_error: BaseException,
*,
model: str,
custom_llm_provider: str | None,
) -> None:
"""Re-raise a failure the provider already saw, or return so the caller declines.
A request that never reached the provider is safe to serve on the Python
path. One that did is not: the provider has already done the work, so a
second attempt bills for it twice. Those surface as an `APIError` carrying
the upstream status, which LiteLLM's exception mapping already understands.
"""
exceptions: Final = _rust_bridge_exceptions()
if exceptions is None:
verbose_logger.debug(
"Rust chat completions bridge raised %s; falling back to Python path",
type(rust_error).__name__,
)
return
declined, upstream_failed = exceptions
if isinstance(rust_error, upstream_failed):
args: Final = rust_error.args
status: Final = args[0] if args else 0
message: Final = args[1] if len(args) > 1 else ""
raise APIError(
status_code=int(status) or 500,
message=f"litellm rust chat completions: {message}",
llm_provider=custom_llm_provider or "",
model=model,
)
if not isinstance(rust_error, declined):
raise rust_error
verbose_logger.debug(
"Rust chat completions declined before calling the provider (%s); using the Python path",
rust_error,
),
)
@ -352,12 +271,14 @@ def chat_completions(
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
request_override: bool | None,
) -> ModelResponse | None:
rust_chat_completions: Final = load_rust_chat_completions()
if rust_chat_completions is None:
return None
try:
rust_response: Final = rust_chat_completions(
def adapt(rust_response: Mapping[str, object]) -> ModelResponse:
on_response(rust_response)
return _build_model_response(rust_response, model_response)
return _CHAT.invoke(
call=lambda rust_chat_completions: rust_chat_completions(
model=model,
messages=messages,
optional_params=optional_params,
@ -366,12 +287,12 @@ def chat_completions(
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)
except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw
_reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider)
return None
on_response(rust_response)
return _build_model_response(rust_response, model_response)
),
fallback=lambda: None,
adapt=adapt,
context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
request_override=request_override,
)
async def achat_completions(
@ -386,12 +307,14 @@ async def achat_completions(
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
request_override: bool | None,
) -> ModelResponse | None:
rust_achat_completions: Final = load_rust_achat_completions()
if rust_achat_completions is None:
return None
try:
rust_response: Final = await rust_achat_completions(
def adapt(rust_response: Mapping[str, object]) -> ModelResponse:
on_response(rust_response)
return _build_model_response(rust_response, model_response)
return await _CHAT.ainvoke(
call=lambda rust_achat_completions: rust_achat_completions(
model=model,
messages=messages,
optional_params=optional_params,
@ -400,12 +323,12 @@ async def achat_completions(
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)
except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw
_reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider)
return None
on_response(rust_response)
return _build_model_response(rust_response, model_response)
),
fallback=async_none,
adapt=adapt,
context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
request_override=request_override,
)
async def achat_completions_or_fallback(
@ -421,6 +344,7 @@ async def achat_completions_or_fallback(
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
python_fallback: Callable[[], Awaitable[object]],
request_override: bool | None,
) -> object:
"""Await the Rust path, falling back to the caller's own Python path when
the bridge is unavailable or the call fails.
@ -430,18 +354,24 @@ async def achat_completions_or_fallback(
already returned a coroutine by the time a Rust failure surfaces, and so
cannot fall back on its own.
"""
response: Final = await achat_completions(
model=model,
messages=messages,
optional_params=optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout=timeout,
on_response=on_response,
def adapt(rust_response: Mapping[str, object]) -> object:
on_response(rust_response)
return _build_model_response(rust_response, model_response)
return await _CHAT.ainvoke(
call=lambda rust_achat_completions: rust_achat_completions(
model=model,
messages=messages,
optional_params=optional_params,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
),
fallback=python_fallback,
adapt=adapt,
context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
request_override=request_override,
)
if response is not None:
return response
return await python_fallback()

123
litellm/rust_bridge/messages.py
View file

@ -3,11 +3,18 @@
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, Protocol, cast
from typing import Final, Protocol
import httpx
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
RustEndpoint,
async_none,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
@ -39,53 +46,29 @@ class RustAmessages(Protocol):
raise NotImplementedError
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustMessagesState:
messages: RustMessages | None = None
amessages: RustAmessages | None = None
_STATE: Final[_RustMessagesState] = _RustMessagesState()
_MESSAGES: Final[RustEndpoint[RustMessages, RustAmessages]] = RustEndpoint.native(
route="messages",
sync="messages",
asynchronous="amessages",
enabled=rust_enabled,
)
def set_rust_messages(
*,
messages: RustMessages | None | _Unset = _UNSET,
amessages: RustAmessages | None | _Unset = _UNSET,
messages: RustMessages | None | Unchanged = UNCHANGED,
amessages: RustAmessages | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(messages, _Unset):
_STATE.messages = messages
if not isinstance(amessages, _Unset):
_STATE.amessages = amessages
def load_rust_messages() -> RustMessages | None:
if _STATE.messages is not None:
return _STATE.messages
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustMessages, getattr(native_bridge, "messages", None))
def load_rust_amessages() -> RustAmessages | None:
if _STATE.amessages is not None:
return _STATE.amessages
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustAmessages, getattr(native_bridge, "amessages", None))
if not isinstance(messages, Unchanged):
if messages is None:
_MESSAGES.sync.reset()
else:
_MESSAGES.sync.override(messages)
if not isinstance(amessages, Unchanged):
if amessages is None:
_MESSAGES.asynchronous.reset()
else:
_MESSAGES.asynchronous.override(amessages)
def messages(
@ -97,18 +80,22 @@ def messages(
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
request_override: bool | None = None,
) -> dict[str, object] | None:
rust_messages: Final = load_rust_messages()
if rust_messages is None:
return None
return rust_messages(
model=model,
body=body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
return _MESSAGES.invoke(
call=lambda rust_messages: rust_messages(
model=model,
body=body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
),
fallback=lambda: None,
adapt=identity,
context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
request_override=request_override,
)
@ -121,16 +108,20 @@ async def amessages(
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
request_override: bool | None = None,
) -> dict[str, object] | None:
rust_amessages: Final = load_rust_amessages()
if rust_amessages is None:
return None
return await rust_amessages(
model=model,
body=body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
return await _MESSAGES.ainvoke(
call=lambda rust_amessages: rust_amessages(
model=model,
body=body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
),
fallback=async_none,
adapt=identity,
context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
request_override=request_override,
)

170
litellm/rust_bridge/ocr.py
View file

@ -2,18 +2,38 @@
from __future__ import annotations
from collections.abc import Awaitable
from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from typing import Final, Protocol
import httpx
from litellm.rust_bridge import configuration as _configuration
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
RustEndpoint,
async_none,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
rust_ocr_enabled = _configuration.rust_ocr_enabled
rust = _configuration.rust
@dataclass(frozen=True, slots=True)
class RustOCRRequest:
model: str
document: dict[str, object]
api_key: str | None
api_base: str | None
custom_llm_provider: str | None
extra_headers: dict[str, object] | None
optional_params: dict[str, object]
timeout: float | httpx.Timeout | None
class RustOcr(Protocol):
def __call__(
self,
@ -44,98 +64,88 @@ class RustAocr(Protocol):
raise NotImplementedError
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
_rust_ocr_impl: RustOcr | None = None
_rust_aocr_impl: RustAocr | None = None
_OCR: Final[RustEndpoint[RustOcr, RustAocr]] = RustEndpoint.native(
route="ocr",
sync="ocr",
asynchronous="aocr",
enabled=_configuration.rust_ocr_enabled,
)
def set_rust_ocr(
*,
ocr: RustOcr | None | _Unset = _UNSET,
aocr: RustAocr | None | _Unset = _UNSET,
ocr: RustOcr | None | Unchanged = UNCHANGED,
aocr: RustAocr | None | Unchanged = UNCHANGED,
) -> 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
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)
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))
if not isinstance(ocr, Unchanged):
if ocr is None:
_OCR.sync.reset()
else:
_OCR.sync.override(ocr)
if not isinstance(aocr, Unchanged):
if aocr is None:
_OCR.asynchronous.reset()
else:
_OCR.asynchronous.override(aocr)
def ocr(
*,
prepare: Callable[[], RustOCRRequest],
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_ocr: Final = load_rust_ocr()
if rust_ocr is None:
return None
return rust_ocr(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=_timeout_to_seconds(timeout),
provider: str,
request_override: bool | None,
eligible: bool,
) -> Mapping[str, object] | None:
return _OCR.invoke(
call=lambda rust_ocr: _call_ocr(rust_ocr, prepare()),
fallback=lambda: None,
adapt=identity,
context=BridgeErrorContext(provider=provider, model=model),
request_override=request_override,
eligible=eligible,
)
async def aocr(
*,
prepare: Callable[[], RustOCRRequest],
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_aocr: Final = load_rust_aocr()
if rust_aocr is None:
return None
return await rust_aocr(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=_timeout_to_seconds(timeout),
provider: str,
request_override: bool | None,
eligible: bool,
) -> Mapping[str, object] | None:
return await _OCR.ainvoke(
call=lambda rust_aocr: _call_aocr(rust_aocr, prepare()),
fallback=async_none,
adapt=identity,
context=BridgeErrorContext(provider=provider, model=model),
request_override=request_override,
eligible=eligible,
)
def _call_ocr(rust_ocr: RustOcr, request: RustOCRRequest) -> Mapping[str, object]:
return rust_ocr(
model=request.model,
document=request.document,
api_key=request.api_key,
api_base=request.api_base,
custom_llm_provider=request.custom_llm_provider,
extra_headers=request.extra_headers,
optional_params=request.optional_params,
timeout_seconds=_timeout_to_seconds(request.timeout),
)
def _call_aocr(rust_aocr: RustAocr, request: RustOCRRequest) -> Awaitable[Mapping[str, object]]:
return rust_aocr(
model=request.model,
document=request.document,
api_key=request.api_key,
api_base=request.api_base,
custom_llm_provider=request.custom_llm_provider,
extra_headers=request.extra_headers,
optional_params=request.optional_params,
timeout_seconds=_timeout_to_seconds(request.timeout),
)

55
litellm/rust_bridge/responses.py Normal file
View file

@ -0,0 +1,55 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, Protocol
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.runtime import BridgeErrorContext, RustEndpoint, async_none, identity
class RustResponses(Protocol):
def __call__(self, request: Mapping[str, object]) -> Mapping[str, object]: ...
class RustAresponses(Protocol):
def __call__(self, request: Mapping[str, object]) -> Awaitable[Mapping[str, object]]: ...
_RESPONSES: Final[RustEndpoint[RustResponses, RustAresponses]] = RustEndpoint.native(
route="responses",
sync="responses",
asynchronous="aresponses",
enabled=rust_enabled,
)
def responses(
*,
prepare: Callable[[], Mapping[str, object]],
model: str,
provider: str,
request_override: bool | None = None,
) -> Mapping[str, object] | None:
return _RESPONSES.invoke(
call=lambda rust_responses: rust_responses(prepare()),
fallback=lambda: None,
adapt=identity,
context=BridgeErrorContext(provider=provider, model=model),
request_override=request_override,
)
async def aresponses(
*,
prepare: Callable[[], Mapping[str, object]],
model: str,
provider: str,
request_override: bool | None = None,
) -> Mapping[str, object] | None:
return await _RESPONSES.ainvoke(
call=lambda rust_aresponses: rust_aresponses(prepare()),
fallback=async_none,
adapt=identity,
context=BridgeErrorContext(provider=provider, model=model),
request_override=request_override,
)

66
litellm/rust_bridge/responses_websocket.py
View file

@ -2,13 +2,18 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, Protocol
import httpx
from websockets.exceptions import ConnectionClosedOK
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
RustBridge,
async_none,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
@ -30,39 +35,22 @@ class RustResponsesWebSocketConnection(Protocol):
) -> RustResponsesWebSocket: ...
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustResponsesWebSocketState:
connection: RustResponsesWebSocketConnection | None = None
_STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState()
_RESPONSES_WEBSOCKET: Final[RustBridge[RustResponsesWebSocketConnection]] = RustBridge.native(
route="responses websocket",
attribute="ResponsesWebSocketConnection",
enabled=rust_enabled,
)
def set_rust_responses_websocket(
*,
connection: RustResponsesWebSocketConnection | None | _Unset = _UNSET,
connection: RustResponsesWebSocketConnection | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(connection, _Unset):
_STATE.connection = connection
def load_rust_responses_websocket() -> RustResponsesWebSocketConnection | None:
if _STATE.connection is not None:
return _STATE.connection
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
connection_type: Final[RustResponsesWebSocketConnection | None] = getattr(
native_bridge, "ResponsesWebSocketConnection", None
)
return connection_type
if not isinstance(connection, Unchanged):
if connection is None:
_RESPONSES_WEBSOCKET.reset()
else:
_RESPONSES_WEBSOCKET.override(connection)
class _ConnectionAdapter:
@ -87,16 +75,16 @@ async def connect(
url: str,
headers: dict[str, str],
timeout: float | httpx.Timeout | None,
request_override: bool | None = None,
) -> _ConnectionAdapter | None:
connection_type: Final = load_rust_responses_websocket()
if connection_type is None:
return None
try:
connection: Final = await connection_type.connect(
return await _RESPONSES_WEBSOCKET.ainvoke(
call=lambda connection_type: connection_type.connect(
url=url,
headers=headers,
timeout_seconds=timeout_to_seconds(timeout),
)
except Exception: # noqa: BLE001 # bridge failures must fall back to Python
return None
return _ConnectionAdapter(connection)
),
fallback=async_none,
adapt=_ConnectionAdapter,
context=BridgeErrorContext(provider="openai", model="responses websocket"),
request_override=request_override,
)

526
litellm/rust_bridge/runtime.py
View file

@ -1,20 +1,29 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from dataclasses import dataclass, field
from enum import Enum
from typing import Final, Generic, NoReturn, TypeAlias, TypeVar
from typing import Final, Generic, NoReturn, Protocol, TypeAlias, TypeVar
from litellm.exceptions import APIError
from litellm.rust_bridge.bindings import native_exception_types
from litellm.rust_bridge.bindings import (
UNCHANGED,
NativeBinding,
Unchanged,
native_exception_types,
)
BindingT = TypeVar("BindingT")
NativeT = TypeVar("NativeT")
ResultT = TypeVar("ResultT")
SyncBindingT = TypeVar("SyncBindingT")
AsyncBindingT = TypeVar("AsyncBindingT")
class FallbackMode(Enum):
PYTHON = "python"
RUST_REQUIRED = "rust_required"
class PythonFallbackReason(Enum):
RUST_DISABLED = "rust_disabled"
RUST_UNAVAILABLE = "rust_unavailable"
RUST_DECLINED = "rust_declined"
@dataclass(frozen=True, slots=True)
@ -23,153 +32,410 @@ class RustHandled(Generic[ResultT]):
@dataclass(frozen=True, slots=True)
class RustDeclined:
reason: str
class PythonFallback:
reason: PythonFallbackReason
detail: str | None = None
@dataclass(frozen=True, slots=True)
class RustUnavailable:
pass
RustAttempt: TypeAlias = RustHandled[ResultT] | RustDeclined | RustUnavailable
RustDispatchResult: TypeAlias = RustHandled[ResultT] | PythonFallback
@dataclass(frozen=True, slots=True)
class BridgeErrorContext:
route: str
provider: str
model: str
def invoke(
*,
native_call: Callable[[], NativeT] | None,
fallback: Callable[[], ResultT],
adapt: Callable[[NativeT], ResultT],
mode: FallbackMode,
context: BridgeErrorContext,
) -> ResultT:
result: Final = attempt(native_call=native_call, adapt=adapt, context=context)
if isinstance(result, RustHandled):
return result.value
if mode is FallbackMode.PYTHON:
return fallback()
_raise_required(result, context)
class RustEnablement(Protocol):
def __call__(self, *, request_override: bool | None = None) -> bool: ...
async def ainvoke(
*,
native_call: Callable[[], Awaitable[NativeT]] | None,
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
mode: FallbackMode,
context: BridgeErrorContext,
) -> ResultT:
result: Final = await aattempt(native_call=native_call, adapt=adapt, context=context)
if isinstance(result, RustHandled):
return result.value
if mode is FallbackMode.PYTHON:
return await fallback()
_raise_required(result, context)
@dataclass(frozen=True, slots=True)
class RustBridge(Generic[BindingT]):
route: str
load: Callable[[], BindingT | None]
enabled: RustEnablement
_native_binding: NativeBinding[BindingT] | None = field(default=None, repr=False)
@classmethod
def native(
cls,
*,
route: str,
attribute: str,
enabled: RustEnablement,
) -> RustBridge[BindingT]:
binding: NativeBinding[BindingT] = NativeBinding.callable(attribute)
return cls(
route=route,
load=binding.load,
enabled=enabled,
_native_binding=binding,
)
def override(self, value: BindingT | None) -> None:
if self._native_binding is None:
raise RuntimeError("only native Rust bridges support binding overrides")
self._native_binding.override(value)
def reset(self) -> None:
if self._native_binding is None:
raise RuntimeError("only native Rust bridges support binding resets")
self._native_binding.reset()
def _attempt(
self,
*,
call: Callable[[BindingT], NativeT],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> RustDispatchResult[ResultT]:
binding_or_fallback: Final = self._binding_or_python_fallback(
request_override=request_override,
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return binding_or_fallback
return self._attempt_call(
call=lambda: call(binding_or_fallback),
adapt=adapt,
context=context,
)
async def _aattempt(
self,
*,
call: Callable[[BindingT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> RustDispatchResult[ResultT]:
binding_or_fallback: Final = self._binding_or_python_fallback(
request_override=request_override,
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return binding_or_fallback
return await self._attempt_acall(
call=lambda: call(binding_or_fallback),
adapt=adapt,
context=context,
)
def invoke(
self,
*,
call: Callable[[BindingT], NativeT],
fallback: Callable[[], ResultT],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> ResultT:
result: Final = self._attempt(
call=call,
adapt=adapt,
context=context,
request_override=request_override,
eligible=eligible,
)
match result:
case RustHandled(value=value):
return value
case PythonFallback():
return fallback()
async def ainvoke(
self,
*,
call: Callable[[BindingT], Awaitable[NativeT]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> ResultT:
result: Final = await self._aattempt(
call=call,
adapt=adapt,
context=context,
request_override=request_override,
eligible=eligible,
)
match result:
case RustHandled(value=value):
return value
case PythonFallback():
return await fallback()
def require(
self,
*,
call: Callable[[BindingT], NativeT],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> ResultT:
result: Final = self._attempt(
call=call,
adapt=adapt,
context=context,
request_override=request_override,
eligible=eligible,
)
match result:
case RustHandled(value=value):
return value
case PythonFallback():
self._raise_required(result)
async def arequire(
self,
*,
call: Callable[[BindingT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> ResultT:
result: Final = await self._aattempt(
call=call,
adapt=adapt,
context=context,
request_override=request_override,
eligible=eligible,
)
match result:
case RustHandled(value=value):
return value
case PythonFallback():
self._raise_required(result)
def accepts(
self,
*,
check: Callable[[BindingT], str | None],
request_override: bool | None = None,
eligible: bool = True,
) -> bool:
binding_or_fallback: Final = self._binding_or_python_fallback(
request_override=request_override,
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return False
try:
reason: Final = check(binding_or_fallback)
except Exception: # noqa: BLE001 # preflight performs no provider I/O, so Python handoff is safe
return False
return reason is None
def _binding_or_python_fallback(
self,
*,
request_override: bool | None,
eligible: bool,
) -> BindingT | PythonFallback:
if not eligible or not self.enabled(request_override=request_override):
return PythonFallback(PythonFallbackReason.RUST_DISABLED)
binding: Final = self.load()
if binding is None:
return PythonFallback(PythonFallbackReason.RUST_UNAVAILABLE)
return binding
def _attempt_call(
self,
*,
call: Callable[[], NativeT],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
) -> RustDispatchResult[ResultT]:
exceptions: Final = native_exception_types()
if exceptions is None:
return RustHandled(adapt(call()))
declined, upstream = exceptions
try:
value: Final = call()
except declined as error:
return PythonFallback(PythonFallbackReason.RUST_DECLINED, _error_message(error))
except upstream as error:
self._raise_upstream(error, context)
return RustHandled(adapt(value))
async def _attempt_acall(
self,
*,
call: Callable[[], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
) -> RustDispatchResult[ResultT]:
exceptions: Final = native_exception_types()
if exceptions is None:
return RustHandled(adapt(await call()))
declined, upstream = exceptions
try:
value: Final = await call()
except declined as error:
return PythonFallback(PythonFallbackReason.RUST_DECLINED, _error_message(error))
except upstream as error:
self._raise_upstream(error, context)
return RustHandled(adapt(value))
def _raise_required(self, fallback: PythonFallback) -> NoReturn:
detail: Final = f": {fallback.detail}" if fallback.detail else ""
reason: Final = _required_reason(fallback.reason)
raise RuntimeError(f"Rust {self.route} bridge {reason}{detail}")
def _raise_upstream(self, error: BaseException, context: BridgeErrorContext) -> NoReturn:
args: Final[tuple[object, ...]] = error.args
attribute_status: Final = getattr(error, "status_code", None)
attribute_message: Final = getattr(error, "message", None)
status_value: Final = attribute_status if isinstance(attribute_status, int) else (args[0] if args else 0)
message_value: Final = (
attribute_message if isinstance(attribute_message, str) else (args[1] if len(args) > 1 else str(error))
)
status: Final = status_value if isinstance(status_value, int) else 0
message: Final = message_value if isinstance(message_value, str) else str(message_value)
raise APIError(
status_code=status or 500,
message=f"litellm rust {self.route}: {message}",
llm_provider=context.provider,
model=context.model,
) from error
def attempt(
*,
native_call: Callable[[], NativeT] | None,
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
) -> RustAttempt[ResultT]:
if native_call is None:
return RustUnavailable()
exceptions: Final = native_exception_types()
if exceptions is None:
return RustHandled(adapt(native_call()))
declined, upstream = exceptions
try:
value: Final = native_call()
except declined as error:
return RustDeclined(reason=_decline_reason(error))
except upstream as error:
_raise_upstream(error, context)
return RustHandled(adapt(value))
@dataclass(frozen=True, slots=True)
class RustEndpoint(Generic[SyncBindingT, AsyncBindingT]):
sync: RustBridge[SyncBindingT]
asynchronous: RustBridge[AsyncBindingT]
@classmethod
def native(
cls,
*,
route: str,
sync: str,
asynchronous: str,
enabled: RustEnablement,
) -> RustEndpoint[SyncBindingT, AsyncBindingT]:
return cls(
sync=RustBridge.native(route=route, attribute=sync, enabled=enabled),
asynchronous=RustBridge.native(
route=route,
attribute=asynchronous,
enabled=enabled,
),
)
def override(
self,
*,
sync: SyncBindingT | None | Unchanged = UNCHANGED,
asynchronous: AsyncBindingT | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(sync, Unchanged):
self.sync.override(sync)
if not isinstance(asynchronous, Unchanged):
self.asynchronous.override(asynchronous)
def reset(self) -> None:
self.sync.reset()
self.asynchronous.reset()
def invoke(
self,
*,
call: Callable[[SyncBindingT], NativeT],
fallback: Callable[[], ResultT],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> ResultT:
return self.sync.invoke(
call=call,
fallback=fallback,
adapt=adapt,
context=context,
request_override=request_override,
eligible=eligible,
)
async def ainvoke(
self,
*,
call: Callable[[AsyncBindingT], Awaitable[NativeT]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> ResultT:
return await self.asynchronous.ainvoke(
call=call,
fallback=fallback,
adapt=adapt,
context=context,
request_override=request_override,
eligible=eligible,
)
def require(
self,
*,
call: Callable[[SyncBindingT], NativeT],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> ResultT:
return self.sync.require(
call=call,
adapt=adapt,
context=context,
request_override=request_override,
eligible=eligible,
)
async def arequire(
self,
*,
call: Callable[[AsyncBindingT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
request_override: bool | None = None,
eligible: bool = True,
) -> ResultT:
return await self.asynchronous.arequire(
call=call,
adapt=adapt,
context=context,
request_override=request_override,
eligible=eligible,
)
async def aattempt(
*,
native_call: Callable[[], Awaitable[NativeT]] | None,
adapt: Callable[[NativeT], ResultT],
context: BridgeErrorContext,
) -> RustAttempt[ResultT]:
if native_call is None:
return RustUnavailable()
exceptions: Final = native_exception_types()
if exceptions is None:
return RustHandled(adapt(await native_call()))
declined, upstream = exceptions
try:
value: Final = await native_call()
except declined as error:
return RustDeclined(reason=_decline_reason(error))
except upstream as error:
_raise_upstream(error, context)
return RustHandled(adapt(value))
def call(operation: Callable[[], ResultT], context: BridgeErrorContext) -> ResultT:
exceptions: Final = native_exception_types()
if exceptions is None:
return operation()
upstream: Final = exceptions[1]
try:
return operation()
except upstream as error:
_raise_upstream(error, context)
async def acall(operation: Callable[[], Awaitable[ResultT]], context: BridgeErrorContext) -> ResultT:
exceptions: Final = native_exception_types()
if exceptions is None:
return await operation()
upstream: Final = exceptions[1]
try:
return await operation()
except upstream as error:
_raise_upstream(error, context)
def _decline_reason(error: BaseException) -> str:
def _error_message(error: BaseException) -> str:
reason: Final[object] = error.args[0] if error.args else str(error)
return reason if isinstance(reason, str) else str(reason)
def _raise_required(
result: RustDeclined | RustUnavailable,
context: BridgeErrorContext,
) -> NoReturn:
raise RuntimeError(f"Rust {context.route} bridge {_required_reason(result)}")
def _required_reason(result: RustDeclined | RustUnavailable) -> str:
match result:
case RustUnavailable():
def _required_reason(reason: PythonFallbackReason) -> str:
match reason:
case PythonFallbackReason.RUST_DISABLED:
return "is disabled"
case PythonFallbackReason.RUST_UNAVAILABLE:
return "is unavailable"
case RustDeclined(reason=reason):
return f"declined the request: {reason}"
case PythonFallbackReason.RUST_DECLINED:
return "declined the request"
def _raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoReturn:
args: Final[tuple[object, ...]] = error.args
status_value: Final = args[0] if args else 0
message_value: Final = args[1] if len(args) > 1 else str(error)
status: Final = status_value if isinstance(status_value, int) else 0
message: Final = message_value if isinstance(message_value, str) else str(message_value)
raise APIError(
status_code=status or 500,
message=f"litellm rust {context.route}: {message}",
llm_provider=context.provider,
model=context.model,
) from error
def always_enabled(*, request_override: bool | None = None) -> bool:
return True
def identity(value: ResultT) -> ResultT:

132
litellm/rust_bridge/transcription.py
View file

@ -1,11 +1,17 @@
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, Protocol, cast
from typing import Final, Protocol
import httpx
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
RustEndpoint,
always_enabled,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
@ -39,62 +45,30 @@ class RustAtranscription(Protocol):
raise NotImplementedError
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass
class _RustTranscriptionState:
transcription: RustTranscription | None = None
atranscription: RustAtranscription | None = None
_STATE: Final = _RustTranscriptionState()
_TRANSCRIPTION: Final[RustEndpoint[RustTranscription, RustAtranscription]] = RustEndpoint.native(
route="audio transcription",
sync="transcription",
asynchronous="atranscription",
enabled=always_enabled,
)
def configure_rust_transcription(
enabled: bool = True,
*,
transcription: RustTranscription | None | _Unset = _UNSET,
atranscription: RustAtranscription | None | _Unset = _UNSET,
transcription: RustTranscription | None | Unchanged = UNCHANGED,
atranscription: RustAtranscription | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(transcription, _Unset):
_STATE.transcription = transcription
if not isinstance(atranscription, _Unset):
_STATE.atranscription = atranscription
def load_rust_transcription() -> RustTranscription | None:
if _STATE.transcription is not None:
return _STATE.transcription
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
return (
None
if native_bridge is None
else cast( # cast-ok: native extension protocol is runtime-defined
RustTranscription, getattr(native_bridge, "transcription", None)
)
)
def load_rust_atranscription() -> RustAtranscription | None:
if _STATE.atranscription is not None:
return _STATE.atranscription
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
return (
None
if native_bridge is None
else cast( # cast-ok: native extension protocol is runtime-defined
RustAtranscription, getattr(native_bridge, "atranscription", None)
)
)
if not isinstance(transcription, Unchanged):
if transcription is None:
_TRANSCRIPTION.sync.reset()
else:
_TRANSCRIPTION.sync.override(transcription)
if not isinstance(atranscription, Unchanged):
if atranscription is None:
_TRANSCRIPTION.asynchronous.reset()
else:
_TRANSCRIPTION.asynchronous.override(atranscription)
def transcription(
@ -107,19 +81,20 @@ def transcription(
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_transcription: Final = load_rust_transcription()
if rust_transcription is None:
return None
return rust_transcription(
model=model,
audio=audio,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
) -> dict[str, object]:
return _TRANSCRIPTION.require(
call=lambda rust_transcription: rust_transcription(
model=model,
audio=audio,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
),
adapt=identity,
context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -133,17 +108,18 @@ async def atranscription(
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_atranscription: Final = load_rust_atranscription()
if rust_atranscription is None:
return None
return await rust_atranscription(
model=model,
audio=audio,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
) -> dict[str, object]:
return await _TRANSCRIPTION.arequire(
call=lambda rust_atranscription: rust_atranscription(
model=model,
audio=audio,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
),
adapt=identity,
context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)

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

@ -8,7 +8,7 @@ import pytest
import litellm
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import configuration
from litellm.rust_bridge import bindings, configuration
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -123,7 +123,7 @@ def test_load_rust_messages_returns_injected_impl():
bridge = RecordingMessages()
litellm.rust(True)
rust_messages.set_rust_messages(messages=bridge)
assert rust_messages.load_rust_messages() is bridge
assert rust_messages._MESSAGES.sync.load() is bridge
def test_bare_rust_still_toggles_ocr():
@ -140,17 +140,16 @@ def test_load_rust_amessages_returns_injected_impl():
bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
assert rust_messages.load_rust_amessages() is bridge
assert rust_messages._MESSAGES.asynchronous.load() is bridge
def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
bindings,
"get_native_bridge",
lambda: None,
)
litellm.rust(True)
assert rust_messages.load_rust_messages() is None
result = rust_messages.messages(
model="claude",
body=REQUEST_BODY,
@ -248,14 +247,14 @@ 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_does_not_handoff_on_unknown_bridge_failure():
bridge = RaisingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate()
with pytest.raises(RuntimeError, match="upstream request failed"):
await _gate()
assert response is None
assert bridge.calls == 1
@ -405,7 +404,7 @@ async def test_fake_stream_wraps_rust_response_as_anthropic_sse():
@pytest.mark.asyncio
async def test_gate_falls_back_when_bridge_unavailable(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
bindings,
"get_native_bridge",
lambda: None,
)

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

@ -3,14 +3,15 @@
import builtins
import importlib
import types
from typing import Any
from typing import Any, Final
from unittest.mock import Mock
import httpx
import pytest
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge import configuration
from litellm.rust_bridge import bindings, configuration
# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr`
# function onto `litellm.ocr` and shadows the submodule, so import the modules
@ -259,14 +260,14 @@ def test_env_var_enables_rust_ocr(monkeypatch):
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
assert ocr_main._rust_request_override(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)
assert rust_bridge.load_rust_ocr() is bridge
assert rust_bridge._OCR.sync.load() is bridge
def test_native_bridge_loader_returns_none_when_extension_absent(monkeypatch):
@ -330,7 +331,7 @@ def test_load_rust_aocr_returns_injected_impl():
bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(aocr=bridge)
assert rust_bridge.load_rust_aocr() is bridge
assert rust_bridge._OCR.asynchronous.load() is bridge
def test_toggle_without_ocr_arg_preserves_injected_impl():
@ -341,16 +342,16 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge)
litellm.rust(False)
assert rust_bridge.load_rust_ocr() is bridge
assert rust_bridge.load_rust_aocr() is async_bridge
assert rust_bridge._OCR.sync.load() is bridge
assert rust_bridge._OCR.asynchronous.load() is async_bridge
litellm.rust(True)
assert rust_bridge.load_rust_ocr() is bridge
assert rust_bridge.load_rust_aocr() is async_bridge
assert rust_bridge._OCR.sync.load() is bridge
assert rust_bridge._OCR.asynchronous.load() is async_bridge
def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
bindings,
"get_native_bridge",
lambda: None,
)
@ -360,21 +361,21 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge)
rust_bridge.set_rust_ocr(ocr=None, aocr=None)
assert rust_bridge.load_rust_ocr() is None
assert rust_bridge.load_rust_aocr() is None
assert rust_bridge._OCR.sync.load() is None
assert rust_bridge._OCR.asynchronous.load() is None
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"),
bindings,
"get_native_bridge",
lambda: None,
)
litellm.rust(True) # no impl injected; extension isn't built in CI
assert rust_bridge.load_rust_ocr() is None
assert rust_bridge.load_rust_aocr() is None
assert rust_bridge._OCR.sync.load() is None
assert rust_bridge._OCR.asynchronous.load() is None
def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
@ -385,14 +386,14 @@ 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"),
bindings,
"get_native_bridge",
lambda: fake_module,
)
litellm.rust(True) # enabled, no impl injected -> import the extension
assert rust_bridge.load_rust_ocr() is fake_module.ocr
assert rust_bridge.load_rust_aocr() is fake_module.aocr
assert rust_bridge._OCR.sync.load() is fake_module.ocr
assert rust_bridge._OCR.asynchronous.load() is fake_module.aocr
def test_timeout_to_seconds_handles_float_timeout_and_none():
@ -408,14 +409,20 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response():
rust_bridge.set_rust_ocr(ocr=bridge)
response = rust_bridge.ocr(
prepare=lambda: rust_bridge.RustOCRRequest(
model="mistral-ocr-latest",
document=DOCUMENT,
api_key="sk-test",
api_base="https://proxy.internal",
custom_llm_provider="mistral",
extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"},
optional_params={"include_image_base64": True, "pages": [0]},
timeout=12.5,
),
model="mistral-ocr-latest",
document=DOCUMENT,
api_key="sk-test",
api_base="https://proxy.internal",
custom_llm_provider="mistral",
extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"},
optional_params={"include_image_base64": True, "pages": [0]},
timeout=12.5,
provider="mistral",
request_override=True,
eligible=True,
)
assert response == FAKE_OCR_RESPONSE
@ -443,14 +450,20 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response():
rust_bridge.set_rust_ocr(aocr=bridge)
response = await rust_bridge.aocr(
prepare=lambda: rust_bridge.RustOCRRequest(
model="mistral-ocr-maas",
document=DOCUMENT,
api_key=None,
api_base=None,
custom_llm_provider="vertex_ai",
extra_headers=None,
optional_params={"vertex_project": "project-1"},
timeout=httpx.Timeout(30.0, read=42.0),
),
model="mistral-ocr-maas",
document=DOCUMENT,
api_key=None,
api_base=None,
custom_llm_provider="vertex_ai",
extra_headers=None,
optional_params={"vertex_project": "project-1"},
timeout=httpx.Timeout(30.0, read=42.0),
provider="vertex_ai",
request_override=True,
eligible=True,
)
assert response == FAKE_OCR_RESPONSE
@ -808,23 +821,27 @@ 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.rust(False)
rust_bridge.set_rust_ocr(ocr=bridge)
def test_ocr_disabled_never_loads_or_prepares_rust(monkeypatch: pytest.MonkeyPatch):
def fail() -> None:
pytest.fail("disabled OCR must not load or prepare Rust")
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 == []
python_ocr: Final = Mock(return_value=OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr"))
litellm.rust(False)
monkeypatch.setattr(bindings, "get_native_bridge", fail)
monkeypatch.setattr(ocr_main, "_prepare_rust_ocr_call", fail)
monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", python_ocr)
response: Final = litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
assert isinstance(response, OCRResponse)
python_ocr.assert_called_once()
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."""
monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None)
litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI
rust_bridge._OCR.sync.override(None)
litellm.rust(True)
captured = {}

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

@ -76,14 +76,14 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None:
@pytest.mark.asyncio
async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(responses_websocket, "_STATE", responses_websocket._RustResponsesWebSocketState())
monkeypatch.setattr(responses_websocket, "get_native_bridge", lambda: None)
responses_websocket._RESPONSES_WEBSOCKET.override(None)
assert (
await responses_websocket.connect(
url="wss://example.test/responses",
headers={},
timeout=None,
request_override=True,
)
is None
)
@ -99,6 +99,7 @@ async def test_enabled_bridge_connects_and_adapts_socket(
url="wss://example.test/responses",
headers={"Authorization": "Bearer key"},
timeout=1.0,
request_override=True,
)
assert connection is not None

18
tests/test_litellm/rust_bridge/test_bindings.py
View file

@ -33,3 +33,21 @@ def test_binding_validates_native_attribute(
binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None)
assert binding.load() == expected
@pytest.mark.parametrize("value", (None, 3, "not callable"))
def test_callable_binding_rejects_non_callables(monkeypatch: pytest.MonkeyPatch, value: object) -> None:
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(route=value))
binding: bindings.NativeBinding[object] = bindings.NativeBinding.callable("route")
assert binding.load() is None
def test_callable_binding_resolves_callable(monkeypatch: pytest.MonkeyPatch) -> None:
native = SimpleNamespace(route=lambda: "native")
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
binding: bindings.NativeBinding[object] = bindings.NativeBinding.callable("route")
assert binding.load() is native.route

7
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 bindings, configuration
from litellm.rust_bridge import chat_completions as bridge
from litellm.rust_bridge import configuration
from litellm.types.utils import ModelResponse
RUST_RESPONSE = {
@ -54,7 +54,7 @@ class _FakeNative:
def _fake_native_bridge(monkeypatch):
"""Expose the bridge's exception classes without the compiled extension."""
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
def _hide_native_bridge(monkeypatch):
@ -63,7 +63,7 @@ def _hide_native_bridge(monkeypatch):
There is no injection seam for "the .so is absent", so the loader itself is
replaced; every other case here uses `set_rust_chat_completions`.
"""
monkeypatch.setattr(bridge, "get_native_bridge", lambda: None)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
@pytest.fixture(autouse=True)
@ -244,6 +244,7 @@ def _call_kwargs(model_response: ModelResponse) -> dict:
"extra_headers": {},
"timeout": 30.0,
"on_response": lambda _rust_response: None,
"request_override": True,
}

47
tests/test_litellm/rust_bridge/test_responses.py Normal file
View file

@ -0,0 +1,47 @@
from unittest.mock import Mock
import pytest
import litellm
from litellm.rust_bridge import responses
@pytest.fixture(autouse=True)
def reset_responses_endpoint():
responses._RESPONSES.reset()
yield
responses._RESPONSES.reset()
def test_unavailable_responses_bridge_does_not_prepare_request() -> None:
prepare = Mock(side_effect=AssertionError("unavailable Rust must not prepare a request"))
responses._RESPONSES.override(sync=None)
litellm.rust(True)
result = responses.responses(
prepare=prepare,
model="gpt-5",
provider="openai",
request_override=True,
)
assert result is None
prepare.assert_not_called()
@pytest.mark.asyncio
async def test_unavailable_async_responses_bridge_does_not_prepare_request(
) -> None:
prepare = Mock(side_effect=AssertionError("unavailable Rust must not prepare a request"))
responses._RESPONSES.override(asynchronous=None)
litellm.rust(True)
result = await responses.aresponses(
prepare=prepare,
model="gpt-5",
provider="openai",
request_override=True,
)
assert result is None
prepare.assert_not_called()

136
tests/test_litellm/rust_bridge/test_runtime.py
View file

@ -26,21 +26,56 @@ def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
def context() -> runtime.BridgeErrorContext:
return runtime.BridgeErrorContext(route="messages", provider="anthropic", model="model")
return runtime.BridgeErrorContext(provider="anthropic", model="model")
def test_invoke_tags_native_decline_before_running_fallback() -> None:
def enabled(*, request_override: bool | None = None) -> bool:
return request_override is not False
def test_disabled_route_does_not_load_or_call_rust() -> None:
bridge = runtime.RustBridge[object](
route="messages",
load=lambda: pytest.fail("disabled route must not load Rust"),
enabled=enabled,
)
value = bridge.invoke(
call=lambda _binding: pytest.fail("disabled route must not call Rust"),
fallback=lambda: "python",
adapt=str,
context=context(),
request_override=False,
)
assert value == "python"
def test_unavailable_route_hands_off_to_python() -> None:
bridge = runtime.RustBridge[object](route="messages", load=lambda: None, enabled=enabled)
value = bridge.invoke(
call=lambda _binding: pytest.fail("unavailable route must not call Rust"),
fallback=lambda: "python",
adapt=str,
context=context(),
)
assert value == "python"
def test_decline_hands_off_to_python() -> None:
calls: list[str] = []
def decline() -> object:
def decline(_binding: object) -> object:
calls.append("rust")
raise RustBridgeDeclined("unsupported")
value = runtime.invoke(
native_call=decline,
bridge = runtime.RustBridge(route="messages", load=object, enabled=enabled)
value = bridge.invoke(
call=decline,
fallback=lambda: calls.append("python") or "fallback",
adapt=str,
mode=runtime.FallbackMode.PYTHON,
context=context(),
)
@ -48,48 +83,93 @@ def test_invoke_tags_native_decline_before_running_fallback() -> None:
assert calls == ["rust", "python"]
def test_invoke_translates_upstream_without_fallback() -> None:
def fail() -> object:
def test_upstream_failure_never_hands_off() -> None:
def fail(_binding: object) -> object:
raise RustUpstreamError(429, "rate limited")
bridge = runtime.RustBridge(route="messages", load=object, enabled=enabled)
with pytest.raises(APIError, match="rate limited") as caught:
runtime.invoke(
native_call=fail,
bridge.invoke(
call=fail,
fallback=lambda: pytest.fail("fallback must not run"),
adapt=str,
mode=runtime.FallbackMode.PYTHON,
context=context(),
)
assert caught.value.status_code == 429
def test_unknown_failure_never_hands_off() -> None:
bridge = runtime.RustBridge(route="messages", load=object, enabled=enabled)
with pytest.raises(RuntimeError, match="unknown"):
bridge.invoke(
call=lambda _binding: (_ for _ in ()).throw(RuntimeError("unknown")),
fallback=lambda: pytest.fail("fallback must not run"),
adapt=str,
context=context(),
)
@pytest.mark.asyncio
async def test_ainvoke_handles_native_success() -> None:
async def native() -> int:
async def test_async_route_handles_native_success() -> None:
async def native(_binding: object) -> int:
return 3
async def fallback() -> str:
pytest.fail("fallback must not run")
assert (
await runtime.ainvoke(
native_call=native,
fallback=fallback,
adapt=str,
mode=runtime.FallbackMode.PYTHON,
context=context(),
)
== "3"
bridge = runtime.RustBridge(route="messages", load=object, enabled=enabled)
value = await bridge.ainvoke(
call=native,
fallback=fallback,
adapt=str,
context=context(),
)
assert value == "3"
def test_required_mode_rejects_unavailable_bridge() -> None:
with pytest.raises(RuntimeError, match="is unavailable"):
runtime.invoke(
native_call=None,
fallback=lambda: pytest.fail("fallback must not run"),
def test_required_route_rejects_unavailable_bridge() -> None:
bridge = runtime.RustBridge[object](route="messages", load=lambda: None, enabled=enabled)
with pytest.raises(RuntimeError, match="unavailable"):
bridge.require(
call=lambda _binding: pytest.fail("unavailable route must not call Rust"),
adapt=str,
mode=runtime.FallbackMode.RUST_REQUIRED,
context=context(),
)
def test_native_endpoint_resolves_overrides_and_resets(monkeypatch: pytest.MonkeyPatch) -> None:
def native_sync() -> str:
return "native"
async def native_async() -> str:
return "native async"
native = SimpleNamespace(
sync_route=native_sync,
async_route=native_async,
RustBridgeDeclined=RustBridgeDeclined,
RustUpstreamError=RustUpstreamError,
)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
endpoint: runtime.RustEndpoint[object, object] = runtime.RustEndpoint.native(
route="test",
sync="sync_route",
asynchronous="async_route",
enabled=enabled,
)
assert endpoint.sync.load() is native_sync
assert endpoint.asynchronous.load() is native_async
replacement = object()
endpoint.override(sync=replacement, asynchronous=None)
assert endpoint.sync.load() is replacement
assert endpoint.asynchronous.load() is None
endpoint.reset()
assert endpoint.sync.load() is native_sync
assert endpoint.asynchronous.load() is native_async

14
tests/test_litellm/test_audio_transcription_rust_bridge.py
View file

@ -4,6 +4,7 @@ import pytest
import litellm
from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch
from litellm.rust_bridge import bindings
rust_bridge = importlib.import_module("litellm.rust_bridge.transcription")
@ -77,13 +78,13 @@ async def test_enabled_async_bridge() -> None:
def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None:
rust_bridge.configure_rust_transcription(transcription=None, atranscription=None)
monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", lambda: None)
assert rust_bridge.load_rust_transcription() is None
assert rust_bridge.load_rust_atranscription() is None
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
assert rust_bridge._TRANSCRIPTION.sync.load() is None
assert rust_bridge._TRANSCRIPTION.asynchronous.load() is None
def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None)
rust_bridge._TRANSCRIPTION.sync.override(None)
with pytest.raises(RuntimeError, match="bridge is unavailable"):
BedrockAudioTranscriptionRustDispatch().audio_transcriptions(
@ -100,10 +101,7 @@ def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) ->
@pytest.mark.asyncio
async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None:
async def unavailable(**_: object) -> None:
return None
monkeypatch.setattr(rust_bridge, "atranscription", unavailable)
rust_bridge._TRANSCRIPTION.asynchronous.override(None)
with pytest.raises(RuntimeError, match="bridge is unavailable"):
await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions(