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 87

refactor(native): share dispatch lifecycle across existing bridges

Browse source
This commit is contained in:
Yujong Lee 2026-09-05 12:59:50 -07:00 • committed by yujonglee
parent 217cb12623
commit 20de0ce755
20 changed files with 1566 additions and 760 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 13429
"limit": 13428
},
"reportArgumentType": {
"limit": 2198
"limit": 2196
},
"reportAssignmentType": {
"limit": 319
@ -99,7 +99,7 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44358
"limit": 44247
},
"reportUnknownLambdaType": {
"limit": 109

41
litellm/rust_bridge/bindings.py
View file

@ -1,9 +1,10 @@
from __future__ import annotations
from collections.abc import Callable
from typing import Final, Generic, TypeVar
from typing import Final, Generic, TypeVar, cast # noqa: TID251 # PyO3 module boundary
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge.protocols import NativeModule
BindingT = TypeVar("BindingT")
@ -15,12 +16,18 @@ 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."""
def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None:
self._attribute: Final = attribute
self._validate: Final = validate
def __init__(self, select: Callable[[NativeModule], BindingT]) -> None:
self._select: Final = select
self._override: BindingT | None | _Unset = _UNSET
def load(self) -> BindingT | None:
@ -29,7 +36,12 @@ class NativeBinding(Generic[BindingT]):
native: Final = get_native_bridge()
if native is None:
return None
return self._validate(getattr(native, self._attribute, None))
module: Final = cast(NativeModule, native) # cast-ok: PyO3 exports are validated individually below
try:
value: Final = self._select(module)
except AttributeError:
return None
return value if callable(value) else None
def override(self, value: BindingT | None) -> None:
self._override = value
@ -38,12 +50,19 @@ class NativeBinding(Generic[BindingT]):
self._override = _UNSET
_DECLINED: Final = NativeBinding(lambda native: native.RustBridgeDeclined)
_UPSTREAM: Final = NativeBinding(lambda native: native.RustUpstreamError)
def _exception_class(value: object) -> type[BaseException] | None:
if isinstance(value, type) and issubclass(value, BaseException):
return value
return None
def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None:
native: Final = get_native_bridge()
if native is None:
return None
declined: Final = getattr(native, "RustBridgeDeclined", None)
upstream: Final = getattr(native, "RustUpstreamError", None)
if not isinstance(declined, type) or not isinstance(upstream, type):
declined: Final = _exception_class(_DECLINED.load())
upstream: Final = _exception_class(_UPSTREAM.load())
if declined is None or upstream is None:
return None
return declined, upstream

294
litellm/rust_bridge/chat_completions.py
View file

@ -14,20 +14,29 @@ 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.protocols import (
RustAchatCompletions,
RustChatCompletions,
RustChatCompletionsDecline,
)
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointBinding,
EndpointDispatch,
async_none,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.utils import ModelResponse
@ -45,47 +54,6 @@ _LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
RUST_RESPONSE_HEADER: Final = "x-litellm-rust"
class RustChatCompletions(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout_seconds: float | None,
) -> Mapping[str, object]:
raise NotImplementedError
class RustAchatCompletions(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout_seconds: float | None,
) -> Awaitable[Mapping[str, object]]:
raise NotImplementedError
class RustChatCompletionsDecline(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
custom_llm_provider: str | None,
) -> str | None:
raise NotImplementedError
class ResponseObserver(Protocol):
"""Invoked with the payload the core returned, on success only.
@ -126,67 +94,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[EndpointDispatch[RustChatCompletions, RustAchatCompletions]] = EndpointDispatch.native(
route="chat_completions",
sync=lambda native: native.chat_completions,
asynchronous=lambda native: native.achat_completions,
enabled=rust_enabled,
)
_CHAT_PREFLIGHT: Final[EndpointBinding[RustChatCompletionsDecline]] = EndpointBinding.native(
route="chat_completions",
select=lambda native: native.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:
@ -247,81 +190,16 @@ def rust_chat_completions_accepts(
return False
if stream:
return False
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")
return False
decline: Final = _load_rust_decline()
if decline is None:
return False
try:
reason: Final = decline(
return _CHAT_PREFLIGHT.accepts(
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,11 +230,13 @@ def chat_completions(
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
) -> 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(
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_chat_completions, timeout_seconds: rust_chat_completions(
model=model,
messages=messages,
optional_params=optional_params,
@ -364,13 +244,12 @@ def chat_completions(
api_base=api_base,
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)
timeout_seconds=timeout_seconds,
),
fallback=lambda: None,
adapt=adapt,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
async def achat_completions(
@ -386,11 +265,13 @@ async def achat_completions(
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
) -> 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(
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_achat_completions, timeout_seconds: rust_achat_completions(
model=model,
messages=messages,
optional_params=optional_params,
@ -398,13 +279,12 @@ async def achat_completions(
api_base=api_base,
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)
timeout_seconds=timeout_seconds,
),
fallback=async_none,
adapt=adapt,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
async def achat_completions_or_fallback(
@ -429,18 +309,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(
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_achat_completions, timeout_seconds: 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_seconds,
),
fallback=python_fallback,
adapt=adapt,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
if response is not None:
return response
return await python_fallback()

150
litellm/rust_bridge/messages.py
View file

@ -2,90 +2,54 @@
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, Protocol, cast
from typing import Final
import httpx
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.protocols import RustAmessages, RustMessages
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointDispatch,
NativeErrorPolicy,
always_enabled,
async_none,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
class RustMessages(Protocol):
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAmessages(Protocol):
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]:
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[EndpointDispatch[RustMessages, RustAmessages]] = EndpointDispatch.native(
route="messages",
sync=lambda native: native.messages,
asynchronous=lambda native: native.amessages,
enabled=always_enabled,
error_policy=NativeErrorPolicy.PROPAGATE,
)
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
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 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))
return _MESSAGES.sync.load()
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))
return _MESSAGES.asynchronous.load()
def messages(
@ -98,17 +62,20 @@ def messages(
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | 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(
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_messages, timeout_seconds: 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_seconds,
),
fallback=lambda: None,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -122,15 +89,18 @@ async def amessages(
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | 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(
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_amessages, timeout_seconds: 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_seconds,
),
fallback=async_none,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)

119
litellm/rust_bridge/ocr.py
View file

@ -2,63 +2,36 @@
from __future__ import annotations
from collections.abc import Awaitable
from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables
from typing import Final
import httpx
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge import configuration as _configuration
from litellm.rust_bridge.protocols import RustAocr, RustOcr
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointDispatch,
NativeErrorPolicy,
async_none,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
class RustOcr(Protocol):
def __call__(
self,
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_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAocr(Protocol):
def __call__(
self,
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_seconds: float | None,
) -> Awaitable[dict[str, object]]:
raise NotImplementedError
def _as_ocr(value: object) -> RustOcr | None:
return cast(RustOcr, value) if callable(value) else None
def _as_aocr(value: object) -> RustAocr | None:
return cast(RustAocr, value) if callable(value) else None
_OCR: Final = NativeBinding("ocr", validate=_as_ocr)
_AOCR: Final = NativeBinding("aocr", validate=_as_aocr)
_OCR: Final[EndpointDispatch[RustOcr, RustAocr]] = EndpointDispatch.native(
route="ocr",
sync=lambda native: native.ocr,
asynchronous=lambda native: native.aocr,
enabled=_configuration.rust_enabled,
error_policy=NativeErrorPolicy.PROPAGATE,
)
def load_rust_ocr() -> RustOcr | None:
return _OCR.load()
return _OCR.sync.load()
def load_rust_aocr() -> RustAocr | None:
return _AOCR.load()
return _OCR.asynchronous.load()
def ocr(
@ -72,18 +45,21 @@ def ocr(
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),
return _OCR.invoke(
prepare=lambda: _timeout_to_seconds(timeout),
call=lambda rust_ocr, timeout_seconds: 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_seconds,
),
fallback=lambda: None,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -98,16 +74,19 @@ async def aocr(
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),
return await _OCR.ainvoke(
prepare=lambda: _timeout_to_seconds(timeout),
call=lambda rust_aocr, timeout_seconds: 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_seconds,
),
fallback=async_none,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)

180
litellm/rust_bridge/protocols.py Normal file
View file

@ -0,0 +1,180 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import Protocol
class RustChatCompletions(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout_seconds: float | None,
) -> Mapping[str, object]: ...
class RustAchatCompletions(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout_seconds: float | None,
) -> Awaitable[Mapping[str, object]]: ...
class RustChatCompletionsDecline(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
custom_llm_provider: str | None,
) -> str | None: ...
class RustResponsesWebSocket(Protocol):
async def send_text(self, text: str) -> None: ...
async def recv_text(self) -> str | None: ...
async def close(self) -> None: ...
class RustResponsesWebSocketConnection(Protocol):
@classmethod
async def connect(
cls,
url: str,
headers: dict[str, str],
timeout_seconds: float | None,
) -> RustResponsesWebSocket: ...
class NativeModule(Protocol):
@property
def chat_completions(self) -> RustChatCompletions: ...
@property
def achat_completions(self) -> RustAchatCompletions: ...
@property
def chat_completions_decline(self) -> RustChatCompletionsDecline: ...
@property
def ResponsesWebSocketConnection(self) -> type[RustResponsesWebSocketConnection]: ...
@property
def RustBridgeDeclined(self) -> type[BaseException]: ...
@property
def RustUpstreamError(self) -> type[BaseException]: ...
@property
def messages(self) -> RustMessages: ...
@property
def amessages(self) -> RustAmessages: ...
@property
def ocr(self) -> RustOcr: ...
@property
def aocr(self) -> RustAocr: ...
@property
def transcription(self) -> RustTranscription: ...
@property
def atranscription(self) -> RustAtranscription: ...
class RustMessages(Protocol):
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> dict[str, object]: ...
class RustAmessages(Protocol):
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]: ...
class RustOcr(Protocol):
def __call__(
self,
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_seconds: float | None,
) -> dict[str, object]: ...
class RustAocr(Protocol):
def __call__(
self,
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_seconds: float | None,
) -> Awaitable[dict[str, object]]: ...
class RustTranscription(Protocol):
def __call__(
self,
model: str,
audio: 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_seconds: float | None,
) -> dict[str, object]: ...
class RustAtranscription(Protocol):
def __call__(
self,
model: str,
audio: 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_seconds: float | None,
) -> Awaitable[dict[str, object]]: ...

95
litellm/rust_bridge/responses_websocket.py
View file

@ -2,67 +2,41 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, Protocol
from typing import Final
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.protocols import (
RustResponsesWebSocket,
RustResponsesWebSocketConnection,
)
from litellm.rust_bridge.runtime import (
AsyncEndpointDispatch,
BridgeErrorContext,
async_none,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
class RustResponsesWebSocket(Protocol):
async def send_text(self, text: str) -> None: ...
async def recv_text(self) -> str | None: ...
async def close(self) -> None: ...
class RustResponsesWebSocketConnection(Protocol):
@classmethod
async def connect(
cls,
url: str,
headers: dict[str, str],
timeout_seconds: float | None,
) -> RustResponsesWebSocket: ...
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustResponsesWebSocketState:
connection: RustResponsesWebSocketConnection | None = None
_STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState()
_RESPONSES_WEBSOCKET: Final[AsyncEndpointDispatch[RustResponsesWebSocketConnection]] = AsyncEndpointDispatch.native(
route="responses_websocket",
asynchronous=lambda native: native.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:
@ -88,15 +62,18 @@ async def connect(
headers: dict[str, str],
timeout: float | httpx.Timeout | None,
) -> _ConnectionAdapter | None:
connection_type: Final = load_rust_responses_websocket()
if connection_type is None:
return None
try:
connection: Final = await connection_type.connect(
url=url,
headers=headers,
timeout_seconds=timeout_to_seconds(timeout),
connection: Final = await _RESPONSES_WEBSOCKET.ainvoke(
prepare=lambda: timeout_to_seconds(timeout),
call=lambda connection_type, timeout_seconds: connection_type.connect(
url=url,
headers=headers,
timeout_seconds=timeout_seconds,
),
fallback=async_none,
adapt=identity,
error_context=BridgeErrorContext(provider="openai", model="responses websocket"),
)
except Exception: # noqa: BLE001 # bridge failures must fall back to Python
except Exception: # noqa: BLE001 # preserve the existing WebSocket connection fallback
return None
return _ConnectionAdapter(connection)
return None if connection is None else _ConnectionAdapter(connection)

581
litellm/rust_bridge/runtime.py
View file

@ -1,150 +1,517 @@
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,
)
from litellm.rust_bridge.protocols import NativeModule
BindingT = TypeVar("BindingT")
SelectedT = TypeVar("SelectedT")
SelectedSyncT = TypeVar("SelectedSyncT")
SelectedAsyncT = TypeVar("SelectedAsyncT")
NativeT = TypeVar("NativeT")
RequestT = TypeVar("RequestT")
ResultT = TypeVar("ResultT")
SyncBindingT = TypeVar("SyncBindingT")
AsyncBindingT = TypeVar("AsyncBindingT")
class FallbackMode(Enum):
PYTHON = "python"
RUST_REQUIRED = "rust_required"
class PythonFallbackReason(Enum):
NATIVE_DISABLED = "native_disabled"
NATIVE_UNAVAILABLE = "native_unavailable"
NATIVE_DECLINED = "native_declined"
class NativeErrorPolicy(Enum):
TRANSLATE = "translate"
PROPAGATE = "propagate"
@dataclass(frozen=True, slots=True)
class RustHandled(Generic[ResultT]):
class Handled(Generic[ResultT]):
value: 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
DispatchResult: TypeAlias = Handled[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) -> 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 EndpointBinding(Generic[BindingT]):
route: str
load: Callable[[], BindingT | None]
enabled: RustEnablement
error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE
_native_binding: NativeBinding[BindingT] | None = field(default=None, repr=False)
@staticmethod
def native(
*,
route: str,
select: Callable[[NativeModule], SelectedT],
enabled: RustEnablement,
error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE,
) -> EndpointBinding[SelectedT]:
binding: Final = NativeBinding(select)
return EndpointBinding(
route=route,
load=binding.load,
enabled=enabled,
error_policy=error_policy,
_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,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> DispatchResult[ResultT]:
binding_or_fallback: Final = self._binding_or_python_fallback(
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return binding_or_fallback
return self._attempt_call(
call=lambda: call(binding_or_fallback, prepare()),
adapt=adapt,
error_context=error_context,
)
async def _aattempt(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> DispatchResult[ResultT]:
binding_or_fallback: Final = self._binding_or_python_fallback(
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return binding_or_fallback
return await self._attempt_acall(
call=lambda: call(binding_or_fallback, prepare()),
adapt=adapt,
error_context=error_context,
)
def invoke(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], NativeT],
fallback: Callable[[], ResultT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
result: Final = self._attempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
return fallback()
async def ainvoke(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], Awaitable[NativeT]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
result: Final = await self._aattempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
return await fallback()
def accepts(
self,
*,
check: Callable[[BindingT], str | None],
eligible: bool = True,
) -> bool:
binding_or_fallback: Final = self._binding_or_python_fallback(
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 require(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
result: Final = self._attempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
self._raise_required(result)
async def arequire(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
result: Final = await self._aattempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
self._raise_required(result)
def can_attempt(
self,
*,
eligible: bool = True,
) -> bool:
return not isinstance(
self._binding_or_python_fallback(eligible=eligible),
PythonFallback,
)
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"native {self.route} endpoint {reason}{detail}")
def _binding_or_python_fallback(
self,
*,
eligible: bool,
) -> BindingT | PythonFallback:
if not eligible or not self.enabled():
return PythonFallback(PythonFallbackReason.NATIVE_DISABLED)
binding: Final = self.load()
if binding is None:
return PythonFallback(PythonFallbackReason.NATIVE_UNAVAILABLE)
return binding
def _attempt_call(
self,
*,
call: Callable[[], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
) -> DispatchResult[ResultT]:
if self.error_policy is NativeErrorPolicy.PROPAGATE:
return Handled(adapt(call()))
exceptions: Final = native_exception_types()
if exceptions is None:
try:
value_without_exceptions: Final = call()
except Exception as error: # noqa: BLE001 # preserve chat fallback when native exception classes are absent
return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error))
return Handled(adapt(value_without_exceptions))
declined, upstream = exceptions
try:
value: Final = call()
except declined as error:
return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error))
except upstream as error:
self._raise_upstream(error, error_context)
return Handled(adapt(value))
async def _attempt_acall(
self,
*,
call: Callable[[], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
) -> DispatchResult[ResultT]:
if self.error_policy is NativeErrorPolicy.PROPAGATE:
return Handled(adapt(await call()))
exceptions: Final = native_exception_types()
if exceptions is None:
try:
value_without_exceptions: Final = await call()
except Exception as error: # noqa: BLE001 # preserve chat fallback when native exception classes are absent
return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error))
return Handled(adapt(value_without_exceptions))
declined, upstream = exceptions
try:
value: Final = await call()
except declined as error:
return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error))
except upstream as error:
self._raise_upstream(error, error_context)
return Handled(adapt(value))
def _raise_upstream(self, error: BaseException, error_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=error_context.provider,
model=error_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 EndpointDispatch(Generic[SyncBindingT, AsyncBindingT]):
sync: EndpointBinding[SyncBindingT]
asynchronous: EndpointBinding[AsyncBindingT]
@staticmethod
def native(
*,
route: str,
sync: Callable[[NativeModule], SelectedSyncT],
asynchronous: Callable[[NativeModule], SelectedAsyncT],
enabled: RustEnablement,
error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE,
) -> EndpointDispatch[SelectedSyncT, SelectedAsyncT]:
return EndpointDispatch(
sync=EndpointBinding.native(route=route, select=sync, enabled=enabled, error_policy=error_policy),
asynchronous=EndpointBinding.native(
route=route,
select=asynchronous,
enabled=enabled,
error_policy=error_policy,
),
)
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,
*,
prepare: Callable[[], RequestT],
call: Callable[[SyncBindingT, RequestT], NativeT],
fallback: Callable[[], ResultT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
return self.sync.invoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=error_context,
eligible=eligible,
)
async def ainvoke(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
return await self.asynchronous.ainvoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=error_context,
eligible=eligible,
)
def require(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[SyncBindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
return self.sync.require(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
)
async def arequire(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
return await self.asynchronous.arequire(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
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))
@dataclass(frozen=True, slots=True)
class AsyncEndpointDispatch(Generic[AsyncBindingT]):
asynchronous: EndpointBinding[AsyncBindingT]
@staticmethod
def native(
*,
route: str,
asynchronous: Callable[[NativeModule], SelectedAsyncT],
enabled: RustEnablement,
) -> AsyncEndpointDispatch[SelectedAsyncT]:
return AsyncEndpointDispatch(
asynchronous=EndpointBinding.native(route=route, select=asynchronous, enabled=enabled)
)
def override(self, value: AsyncBindingT | None) -> None:
self.asynchronous.override(value)
def reset(self) -> None:
self.asynchronous.reset()
async def ainvoke(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
) -> ResultT:
return await self.asynchronous.ainvoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=error_context,
eligible=eligible,
)
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.NATIVE_DISABLED:
return "is disabled"
case PythonFallbackReason.NATIVE_UNAVAILABLE:
return "is unavailable"
case RustDeclined(reason=reason):
return f"declined the request: {reason}"
case PythonFallbackReason.NATIVE_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() -> bool:
return True
def identity(value: ResultT) -> ResultT:
return value
async def async_none() -> None:
return None

164
litellm/rust_bridge/transcription.py
View file

@ -1,99 +1,53 @@
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, Protocol, cast
from typing import Final
import httpx
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.protocols import RustAtranscription, RustTranscription
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointDispatch,
NativeErrorPolicy,
always_enabled,
async_none,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
class RustTranscription(Protocol):
def __call__(
self,
model: str,
audio: 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_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAtranscription(Protocol):
def __call__(
self,
model: str,
audio: 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_seconds: float | None,
) -> Awaitable[dict[str, object]]:
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[EndpointDispatch[RustTranscription, RustAtranscription]] = EndpointDispatch.native(
route="audio transcription",
sync=lambda native: native.transcription,
asynchronous=lambda native: native.atranscription,
enabled=always_enabled,
error_policy=NativeErrorPolicy.PROPAGATE,
)
def configure_rust_transcription(
*,
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
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 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)
)
)
return _TRANSCRIPTION.sync.load()
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)
)
)
return _TRANSCRIPTION.asynchronous.load()
def transcription(
@ -107,18 +61,21 @@ def transcription(
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),
return _TRANSCRIPTION.invoke(
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_transcription, timeout_seconds: 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_seconds,
),
fallback=lambda: None,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -133,16 +90,19 @@ async def atranscription(
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),
return await _TRANSCRIPTION.ainvoke(
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_atranscription, timeout_seconds: 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_seconds,
),
fallback=async_none,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)

4
ruff-strict-budget.json
View file

@ -156,7 +156,7 @@
"limit": 215
},
"PLW0603": {
"limit": 190
"limit": 188
},
"PLW1508": {
"limit": 190
@ -231,7 +231,7 @@
"limit": 5
},
"TID251": {
"limit": 1035
"limit": 1033
},
"TRY002": {
"limit": 524

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

@ -135,7 +135,7 @@ def test_load_rust_amessages_returns_injected_impl():
def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
importlib.import_module("litellm.rust_bridge.bindings"),
"get_native_bridge",
lambda: None,
)
@ -383,7 +383,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"),
importlib.import_module("litellm.rust_bridge.bindings"),
"get_native_bridge",
lambda: None,
)

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

@ -2467,7 +2467,7 @@ class TestRustChatCompletionsHook:
def declining_native(**_kwargs):
raise _Declined("blank message text")
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
@ -2499,7 +2499,7 @@ class TestRustChatCompletionsHook:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
async def declining_native(**_kwargs):
raise _Declined("blank message text")
@ -2559,7 +2559,7 @@ class TestRustChatCompletionsHook:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
def declining_native(**_kwargs):
raise _Declined("blank message text")

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

@ -205,7 +205,7 @@ async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch):
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
async def declining_native(**_kwargs):
raise _Declined("blank message text")
@ -282,7 +282,7 @@ async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines():
return ModelResponse()
with (
patch.object(bridge, "get_native_bridge", lambda: _FakeNative()),
patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()),
patch.object(
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
),
@ -389,7 +389,7 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines():
logging_obj = MagicMock()
with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()):
with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()):
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
@ -475,7 +475,7 @@ def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines():
logging_obj, calls = _recording_logging_obj()
with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()):
with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()):
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)

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

@ -216,13 +216,13 @@ def build_prepared_request(
@pytest.fixture(autouse=True)
def _reset_rust_flag():
"""Keep the global toggle isolated between tests."""
rust_bridge._OCR.reset()
rust_bridge._AOCR.reset()
rust_bridge._OCR.sync.reset()
rust_bridge._OCR.asynchronous.reset()
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
yield
rust_bridge._OCR.reset()
rust_bridge._AOCR.reset()
rust_bridge._OCR.sync.reset()
rust_bridge._OCR.asynchronous.reset()
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
@ -232,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._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
return bridge
@ -241,14 +241,14 @@ def fake_async_bridge():
"""Enable the async Rust path with an injected recording bridge."""
bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge._AOCR.override(bridge)
rust_bridge._OCR.asynchronous.override(bridge)
return bridge
def test_load_rust_ocr_returns_injected_impl():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
assert rust_bridge.load_rust_ocr() is bridge
@ -312,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._AOCR.override(bridge)
rust_bridge._OCR.asynchronous.override(bridge)
assert rust_bridge.load_rust_aocr() is bridge
@ -321,8 +321,8 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
bridge = RecordingBridge()
async_bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._AOCR.override(async_bridge)
rust_bridge._OCR.sync.override(bridge)
rust_bridge._OCR.asynchronous.override(async_bridge)
litellm.rust(False)
assert rust_bridge.load_rust_ocr() is bridge
@ -341,11 +341,11 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
bridge = RecordingBridge()
async_bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._AOCR.override(async_bridge)
rust_bridge._OCR.sync.override(bridge)
rust_bridge._OCR.asynchronous.override(async_bridge)
rust_bridge._OCR.override(None)
rust_bridge._AOCR.override(None)
rust_bridge._OCR.sync.override(None)
rust_bridge._OCR.asynchronous.override(None)
assert rust_bridge.load_rust_ocr() is None
assert rust_bridge.load_rust_aocr() is None
@ -392,7 +392,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response():
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
response = rust_bridge.ocr(
model="mistral-ocr-latest",
document=DOCUMENT,
@ -427,7 +427,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response():
litellm.rust(True)
rust_bridge._AOCR.override(bridge)
rust_bridge._OCR.asynchronous.override(bridge)
response = await rust_bridge.aocr(
model="mistral-ocr-maas",
document=DOCUMENT,
@ -456,7 +456,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
bridge = RecordingBridge()
logging_obj = RecordingLogging()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
response = ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -489,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._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(api_key=None, timeout=None),
@ -502,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._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
def _resolver(name: str) -> str | None:
raise AssertionError(f"resolver should not be called for {name}")
@ -522,7 +522,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
bridge = RecordingBridge()
resolver_calls = []
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
def _resolver(name):
resolver_calls.append(name)
@ -545,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._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -572,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._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
def _resolver(name: str) -> str | None:
return {
@ -596,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._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -614,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._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -635,7 +635,7 @@ def test_run_rust_ocr_runs_pre_call_logging():
logging_obj = RecordingLogging()
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._OCR.sync.override(bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
@ -723,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._OCR.override(RaisingBridge())
rust_bridge._OCR.sync.override(RaisingBridge())
with pytest.raises(CapturedException):
litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
@ -769,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._AOCR.override(RaisingAsyncBridge())
rust_bridge._OCR.asynchronous.override(RaisingAsyncBridge())
with pytest.raises(CapturedException):
await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
@ -798,7 +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._OCR.override(bridge)
rust_bridge._OCR.sync.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 == []

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

@ -65,8 +65,8 @@ 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)
configuration.rust(True)
responses_websocket._RESPONSES_WEBSOCKET.override(None)
assert (
await responses_websocket.connect(
@ -82,6 +82,7 @@ async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch)
async def test_enabled_bridge_connects_and_adapts_socket(
monkeypatch: pytest.MonkeyPatch,
) -> None:
configuration.rust(True)
responses_websocket.set_rust_responses_websocket(connection=_FakeNativeBridge)
connection = await responses_websocket.connect(
@ -94,3 +95,30 @@ async def test_enabled_bridge_connects_and_adapts_socket(
await connection.send("response.create")
assert await connection.recv() == "response.completed"
await connection.close()
class _FailingNativeBridge:
@classmethod
async def connect(
cls,
*,
url: str,
headers: dict[str, str],
timeout_seconds: float | None,
) -> _FakeNativeConnection:
raise RuntimeError("connection failed")
@pytest.mark.asyncio
async def test_connection_failure_preserves_python_fallback() -> None:
configuration.rust(True)
responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge)
assert (
await responses_websocket.connect(
url="wss://example.test/responses",
headers={},
timeout=None,
)
is None
)

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

@ -1,3 +1,7 @@
import json
import subprocess
import sys
from pathlib import Path
from types import SimpleNamespace
from typing import Final
@ -6,30 +10,111 @@ import pytest
from litellm.rust_bridge import bindings
def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None:
native = SimpleNamespace(route=lambda: "native")
def test_binding_distinguishes_disable_from_reset(monkeypatch: pytest.MonkeyPatch) -> None:
native: Final = SimpleNamespace(chat_completions=lambda: "native")
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
binding: bindings.NativeBinding[object] = bindings.NativeBinding("route", validate=lambda value: value)
assert binding.load() is native.route
binding: Final = bindings.NativeBinding(lambda module: module.chat_completions)
assert binding.load() is native.chat_completions
binding.override(None)
assert binding.load() is None
replacement = object()
binding.override(replacement)
assert binding.load() is replacement
replacement: Final = SimpleNamespace(chat_completions=lambda: "replacement")
binding.override(replacement.chat_completions)
assert binding.load() is replacement.chat_completions
binding.reset()
assert binding.load() is native.route
assert binding.load() is native.chat_completions
@pytest.mark.parametrize(("value", "expected"), ((3, 3), ("invalid", None), (None, None)))
def test_binding_validates_native_attribute(
monkeypatch: pytest.MonkeyPatch, value: object, expected: int | None
) -> None:
native: Final = SimpleNamespace(route=value)
@pytest.mark.parametrize("native", (None, SimpleNamespace(), SimpleNamespace(chat_completions=3)))
def test_missing_or_invalid_export_is_unavailable(monkeypatch: pytest.MonkeyPatch, native: object) -> None:
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None)
binding: Final = bindings.NativeBinding(lambda module: module.chat_completions)
assert binding.load() == expected
assert binding.load() is None
def test_selection_is_lazy_and_preserves_other_exports(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(bindings, "get_native_bridge", lambda: pytest.fail("must not load during construction"))
binding: Final = bindings.NativeBinding(lambda module: module.chat_completions)
native: Final = SimpleNamespace(chat_completions=lambda: "native", achat_completions=None)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
assert binding.load() is native.chat_completions
assert bindings.NativeBinding(lambda module: module.achat_completions).load() is None
@pytest.mark.parametrize("invalid", (None, str, lambda: None))
def test_native_exception_types_reject_non_exception_classes(monkeypatch: pytest.MonkeyPatch, invalid: object) -> None:
native: Final = SimpleNamespace(RustBridgeDeclined=invalid, RustUpstreamError=RuntimeError)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
assert bindings.native_exception_types() is None
@pytest.mark.parametrize(
("expression", "expected_rule"),
(
("NativeBinding(lambda native: native.chat_completion)", "reportAttributeAccessIssue"),
(
"wrong: NativeBinding[RustAchatCompletions] = NativeBinding(lambda native: native.chat_completions)",
"reportAssignmentType",
),
(
'EndpointBinding.native(route="chat", select=lambda native: native.chat_completion, enabled=always_enabled)',
"reportAttributeAccessIssue",
),
("NativeBinding(lambda native: native.ocrr)", "reportAttributeAccessIssue"),
(
"wrong: NativeBinding[RustAmessages] = NativeBinding(lambda native: native.messages)",
"reportAssignmentType",
),
(
"wrong: NativeBinding[RustAocr] = NativeBinding(lambda native: native.ocr)",
"reportAssignmentType",
),
(
"wrong: NativeBinding[RustAtranscription] = NativeBinding(lambda native: native.transcription)",
"reportAssignmentType",
),
),
)
def test_selectors_are_checked_by_type_checker(tmp_path: Path, expression: str, expected_rule: str) -> None:
source: Final = tmp_path / "binding_contract.py"
source.write_text(
"from typing_extensions import assert_type\n"
"from litellm.rust_bridge.bindings import NativeBinding\n"
"from litellm.rust_bridge.protocols import RustChatCompletions, RustAchatCompletions, "
"RustMessages, RustAmessages, RustOcr, RustAocr, RustTranscription, RustAtranscription\n"
"from litellm.rust_bridge.runtime import EndpointBinding, EndpointDispatch, always_enabled\n"
"binding = NativeBinding(lambda native: native.chat_completions)\n"
"assert_type(binding, NativeBinding[RustChatCompletions])\n"
'bridge = EndpointBinding.native(route="chat", select=lambda native: native.chat_completions, enabled=always_enabled)\n'
"assert_type(bridge, EndpointBinding[RustChatCompletions])\n"
'endpoint = EndpointDispatch.native(route="chat", sync=lambda native: native.chat_completions, '
"asynchronous=lambda native: native.achat_completions, enabled=always_enabled)\n"
"assert_type(endpoint, EndpointDispatch[RustChatCompletions, RustAchatCompletions])\n"
"assert_type(NativeBinding(lambda native: native.messages), NativeBinding[RustMessages])\n"
"assert_type(NativeBinding(lambda native: native.ocr), NativeBinding[RustOcr])\n"
"assert_type(NativeBinding(lambda native: native.transcription), NativeBinding[RustTranscription])\n"
+ expression
+ "\n"
)
config: Final = tmp_path / "pyrightconfig.json"
config.write_text(
json.dumps(
{
"include": [str(source)],
"extraPaths": [str(Path(__file__).resolve().parents[3])],
"typeCheckingMode": "basic",
}
)
)
result: Final = subprocess.run(
[sys.executable, "-m", "basedpyright", "--project", str(config), "--outputjson"],
capture_output=True,
text=True,
check=False,
)
diagnostics: Final = json.loads(result.stdout)["generalDiagnostics"]
assert result.returncode == 1, result.stdout + result.stderr
assert [(item["rule"], item["range"]["start"]["line"]) for item in diagnostics] == [(expected_rule, 13)]

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

@ -10,7 +10,7 @@ from __future__ import annotations
import pytest
import litellm
from litellm.rust_bridge import configuration
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge import chat_completions as bridge
from litellm.types.utils import ModelResponse
@ -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)
@ -393,3 +393,22 @@ class TestFailureClassification:
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
@pytest.mark.asyncio
async def test_missing_native_exception_types_preserves_python_fallback(monkeypatch):
_hide_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(
chat_completions=_RecordingCall(error=RuntimeError("connection failed")),
achat_completions=_RecordingAsyncCall(error=RuntimeError("connection failed")),
)
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
async def fallback():
return "python"
assert (
await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
== "python"
)

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

@ -1,6 +1,8 @@
from __future__ import annotations
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Final
import pytest
@ -18,78 +20,412 @@ class RustUpstreamError(Exception):
@pytest.fixture(autouse=True)
def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
native = SimpleNamespace(
RustBridgeDeclined=RustBridgeDeclined,
RustUpstreamError=RustUpstreamError,
monkeypatch.setattr(
bindings,
"get_native_bridge",
lambda: SimpleNamespace(
RustBridgeDeclined=RustBridgeDeclined,
RustUpstreamError=RustUpstreamError,
),
)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
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:
calls: list[str] = []
def enabled() -> bool:
return True
def decline() -> object:
calls.append("rust")
raise RustBridgeDeclined("unsupported")
value = runtime.invoke(
native_call=decline,
fallback=lambda: calls.append("python") or "fallback",
@dataclass(frozen=True, slots=True)
class FallbackCase:
process_enabled: bool | None = None
eligible: bool = True
binding_available: bool = True
declined: bool = False
expected_events: tuple[str, ...] = ()
FALLBACK_CASES: Final = (
pytest.param(
FallbackCase(process_enabled=False, expected_events=("python",)),
id="process-disabled",
),
pytest.param(
FallbackCase(eligible=False, expected_events=("python",)),
id="request-ineligible",
),
pytest.param(
FallbackCase(binding_available=False, expected_events=("load", "python")),
id="bridge-unavailable",
),
pytest.param(
FallbackCase(declined=True, expected_events=("load", "prepare", "rust", "python")),
id="bridge-declined",
),
)
@pytest.mark.parametrize("case", FALLBACK_CASES)
def test_invoke_falls_back_only_before_provider_success(case: FallbackCase) -> None:
events: list[str] = []
def load() -> object | None:
events.append("load")
return object() if case.binding_available else None
def call(_binding: object, _request: object) -> int:
events.append("rust")
if case.declined:
raise RustBridgeDeclined("unsupported")
return 3
bridge: Final = runtime.EndpointBinding(
route="messages", load=load, enabled=lambda: case.process_enabled is not False
)
result: Final = bridge.invoke(
prepare=lambda: events.append("prepare"),
call=call,
fallback=lambda: events.append("python") or "fallback",
adapt=str,
mode=runtime.FallbackMode.PYTHON,
context=context(),
error_context=context(),
eligible=case.eligible,
)
assert value == "fallback"
assert calls == ["rust", "python"]
def test_invoke_translates_upstream_without_fallback() -> None:
def fail() -> object:
raise RustUpstreamError(429, "rate limited")
with pytest.raises(APIError, match="rate limited") as caught:
runtime.invoke(
native_call=fail,
fallback=lambda: pytest.fail("fallback must not run"),
adapt=str,
mode=runtime.FallbackMode.PYTHON,
context=context(),
)
assert caught.value.status_code == 429
assert result == "fallback"
assert tuple(events) == case.expected_events
@pytest.mark.asyncio
async def test_ainvoke_handles_native_success() -> None:
async def native() -> int:
@pytest.mark.parametrize("case", FALLBACK_CASES)
async def test_ainvoke_matches_sync_fallback_contract(case: FallbackCase) -> None:
events: list[str] = []
def load() -> object | None:
events.append("load")
return object() if case.binding_available else None
async def call(_binding: object, _request: object) -> int:
events.append("rust")
if case.declined:
raise RustBridgeDeclined("unsupported")
return 3
async def fallback() -> str:
events.append("python")
return "fallback"
bridge: Final = runtime.EndpointBinding(
route="messages", load=load, enabled=lambda: case.process_enabled is not False
)
result: Final = await bridge.ainvoke(
prepare=lambda: events.append("prepare"),
call=call,
fallback=fallback,
adapt=str,
error_context=context(),
eligible=case.eligible,
)
assert result == "fallback"
assert tuple(events) == case.expected_events
def test_invoke_adapts_native_success_without_fallback() -> None:
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
result: Final = bridge.invoke(
prepare=lambda: 3,
call=lambda _binding, request: request * 2,
fallback=lambda: pytest.fail("fallback must not run"),
adapt=lambda value: f"adapted-{value}",
error_context=context(),
)
assert result == "adapted-6"
@pytest.mark.asyncio
async def test_ainvoke_adapts_native_success_without_fallback() -> None:
async def call(_binding: object, request: int) -> int:
return request * 2
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: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
result: Final = await bridge.ainvoke(
prepare=lambda: 3,
call=call,
fallback=fallback,
adapt=lambda value: f"adapted-{value}",
error_context=context(),
)
assert result == "adapted-6"
def test_required_mode_rejects_unavailable_bridge() -> None:
with pytest.raises(RuntimeError, match="is unavailable"):
runtime.invoke(
native_call=None,
@pytest.mark.parametrize(
("error", "expected_type", "expected_status", "expected_message"),
(
pytest.param(RustUpstreamError(401, "unauthorized"), APIError, 401, "unauthorized", id="auth"),
pytest.param(RustUpstreamError(429, "rate limited"), APIError, 429, "rate limited", id="rate-limit"),
pytest.param(RustUpstreamError(500, "failed"), APIError, 500, "failed", id="server-error"),
pytest.param(RustUpstreamError(0, "connection reset"), APIError, 500, "connection reset", id="transport"),
pytest.param(RustUpstreamError(403, "forbidden"), APIError, 403, "forbidden", id="other-status"),
),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_upstream_failure_maps_to_api_error_without_fallback(
asynchronous: bool,
error: RustUpstreamError,
expected_type: type[BaseException],
expected_status: int,
expected_message: str,
) -> None:
def fail(_binding: object, _request: object) -> object:
raise error
async def afail(binding: object, request: object) -> object:
return fail(binding, request)
async def fallback() -> str:
pytest.fail("fallback must not run")
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
async def invoke() -> None:
if asynchronous:
await bridge.ainvoke(
prepare=lambda: None, call=afail, fallback=fallback, adapt=str, error_context=context()
)
else:
bridge.invoke(
prepare=lambda: None,
call=fail,
fallback=lambda: pytest.fail("fallback must not run"),
adapt=str,
error_context=context(),
)
with pytest.raises(expected_type, match=expected_message) as caught:
await invoke()
assert type(caught.value) is expected_type
assert caught.value.status_code == expected_status
assert caught.value.llm_provider == "anthropic"
assert caught.value.model == "model"
assert caught.value.__cause__ is error
@pytest.mark.asyncio
async def test_async_upstream_failure_maps_to_api_error_without_fallback() -> None:
async def fail(_binding: object, _request: object) -> object:
raise RustUpstreamError(503, "overloaded")
async def fallback() -> object:
pytest.fail("fallback must not run")
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
with pytest.raises(APIError, match="overloaded") as caught:
await bridge.ainvoke(prepare=lambda: None, call=fail, fallback=fallback, adapt=str, error_context=context())
assert caught.value.status_code == 503
def test_unknown_failure_is_preserved_without_fallback() -> None:
error: Final = RuntimeError("unknown")
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
with pytest.raises(RuntimeError, match="unknown") as caught:
bridge.invoke(
prepare=lambda: None,
call=lambda _binding, _request: (_ for _ in ()).throw(error),
fallback=lambda: pytest.fail("fallback must not run"),
adapt=str,
mode=runtime.FallbackMode.RUST_REQUIRED,
context=context(),
error_context=context(),
)
assert caught.value is error
@pytest.mark.parametrize(
("process_enabled", "binding_available", "declined", "expected_message"),
(
pytest.param(False, True, False, "native messages endpoint is disabled", id="disabled"),
pytest.param(None, False, False, "native messages endpoint is unavailable", id="unavailable"),
pytest.param(
None,
True,
True,
"native messages endpoint declined the request: unsupported",
id="declined",
),
),
)
def test_require_explains_why_rust_did_not_handle_request(
process_enabled: bool | None,
binding_available: bool,
declined: bool,
expected_message: str,
) -> None:
def call(_binding: object, _request: object) -> object:
if declined:
raise RustBridgeDeclined("unsupported")
return object()
bridge: Final = runtime.EndpointBinding(
route="messages",
load=object if binding_available else lambda: None,
enabled=lambda: process_enabled is not False,
)
with pytest.raises(RuntimeError, match=f"^{expected_message}$"):
bridge.require(
prepare=lambda: None,
call=call,
adapt=str,
error_context=context(),
)
@pytest.mark.parametrize(
("state", "expected", "expected_events"),
(
pytest.param("disabled", False, (), id="disabled"),
pytest.param("ineligible", False, (), id="ineligible"),
pytest.param("unavailable", False, ("load",), id="unavailable"),
pytest.param("available", True, ("load",), id="available"),
),
)
def test_can_attempt_only_enabled_available_requests(
state: str,
expected: bool,
expected_events: tuple[str, ...],
) -> None:
events: list[str] = []
def load() -> object | None:
events.append("load")
return None if state == "unavailable" else object()
bridge: Final = runtime.EndpointBinding(route="messages", load=load, enabled=lambda: state != "disabled")
assert (
bridge.can_attempt(
eligible=state != "ineligible",
)
is expected
)
assert tuple(events) == expected_events
def test_native_endpoint_applies_partial_overrides_and_reset(monkeypatch: pytest.MonkeyPatch) -> None:
def native_sync() -> str:
return "native"
async def native_async() -> str:
return "native async"
def replacement_sync() -> str:
return "replacement"
monkeypatch.setattr(
bindings,
"get_native_bridge",
lambda: SimpleNamespace(chat_completions=native_sync, achat_completions=native_async),
)
endpoint: Final[runtime.EndpointDispatch[object, object]] = runtime.EndpointDispatch.native(
route="test",
sync=lambda native: native.chat_completions,
asynchronous=lambda native: native.achat_completions,
enabled=enabled,
)
assert endpoint.sync.load() is native_sync
assert endpoint.asynchronous.load() is native_async
endpoint.override(sync=replacement_sync)
assert endpoint.sync.load() is replacement_sync
assert endpoint.asynchronous.load() is native_async
endpoint.override(asynchronous=None)
assert endpoint.sync.load() is replacement_sync
assert endpoint.asynchronous.load() is None
endpoint.reset()
assert endpoint.sync.load() is native_sync
assert endpoint.asynchronous.load() is native_async
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_response_adaptation_failure_never_authorizes_fallback(asynchronous: bool) -> None:
def adapt(value: str) -> str:
assert value == "provider response"
raise RustBridgeDeclined("adapter failed after provider response")
async def native(binding: object, request: object) -> str:
return "provider response"
async def fallback() -> str:
pytest.fail("a received response must not be retried")
bridge = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
async def invoke() -> None:
if asynchronous:
await bridge.ainvoke(
prepare=lambda: None, call=native, fallback=fallback, adapt=adapt, error_context=context()
)
else:
bridge.invoke(
prepare=lambda: None,
call=lambda binding, request: "provider response",
fallback=lambda: pytest.fail("a received response must not be retried"),
adapt=adapt,
error_context=context(),
)
with pytest.raises(RustBridgeDeclined, match="adapter failed"):
await invoke()
@pytest.mark.parametrize("error", (RustBridgeDeclined("unsupported"), RustUpstreamError(429, "rate limited")))
def test_propagate_policy_preserves_native_errors(error: Exception) -> None:
def fail(_binding: object, _request: object) -> object:
raise error
bridge: Final = runtime.EndpointBinding(
route="ocr", load=object, enabled=enabled, error_policy=runtime.NativeErrorPolicy.PROPAGATE
)
with pytest.raises(type(error)) as caught:
bridge.invoke(
prepare=lambda: None,
call=fail,
fallback=lambda: pytest.fail("fallback must not run"),
adapt=str,
error_context=context(),
)
assert caught.value is error
@pytest.mark.asyncio
@pytest.mark.parametrize("error", (RustBridgeDeclined("unsupported"), RustUpstreamError(429, "rate limited")))
async def test_async_propagate_policy_preserves_native_errors(error: Exception) -> None:
async def fail(_binding: object, _request: object) -> object:
raise error
async def fallback() -> str:
pytest.fail("fallback must not run")
bridge: Final = runtime.EndpointBinding(
route="messages", load=object, enabled=enabled, error_policy=runtime.NativeErrorPolicy.PROPAGATE
)
with pytest.raises(type(error)) as caught:
await bridge.ainvoke(prepare=lambda: None, call=fail, fallback=fallback, adapt=str, error_context=context())
assert caught.value is error

2
tests/test_litellm/test_audio_transcription_rust_bridge.py
View file

@ -77,7 +77,7 @@ 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)
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None)
assert rust_bridge.load_rust_transcription() is None
assert rust_bridge.load_rust_atranscription() is None

2
type-discipline-budget.json
View file

@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1035
"limit": 1031
},
"LIT007": {
"limit": 0