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 86

refactor(native): separate attempts from declarative dispatch

Browse source
This commit is contained in:
Yujong Lee 2026-09-05 20:19:38 -07:00 • committed by yujonglee
parent 77bf16b847
commit 3c7bb5131e
25 changed files with 998 additions and 1819 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 13427
"limit": 13426
},
"reportArgumentType": {
"limit": 2194
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 3369
"limit": 3368
},
"reportFunctionMemberAccess": {
"limit": 7
@ -90,7 +90,7 @@
"limit": 8
},
"reportReturnType": {
"limit": 180
"limit": 178
},
"reportTypedDictNotRequiredAccess": {
"limit": 22
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44136
"limit": 44025
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38271
"limit": 38269
},
"reportUnknownParameterType": {
"limit": 19584
},
"reportUnknownVariableType": {
"limit": 29814
"limit": 29813
},
"reportUnnecessaryCast": {
"limit": 110

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

@ -27,6 +27,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
from litellm.rust_bridge.dispatch import adispatch, dispatch
from litellm.types.llms.anthropic import (
ContentBlockDelta,
ContentBlockStart,
@ -453,7 +454,25 @@ class AnthropicChatCompletion(BaseLLM):
timeout=timeout,
)
return rust_chat_completions_bridge.achat_completions_or_fallback(
return adispatch(
native=lambda: rust_chat_completions_bridge.achat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
),
python=python_fallback,
route="chat_completions",
errors=rust_chat_completions_bridge.error_handling(custom_llm_provider or "", model),
)
rust_response: Final = dispatch(
native=lambda: rust_chat_completions_bridge.chat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
@ -464,19 +483,10 @@ class AnthropicChatCompletion(BaseLLM):
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
python_fallback=python_fallback,
)
rust_response: Final = rust_chat_completions_bridge.chat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
),
python=lambda: None,
route="chat_completions",
errors=rust_chat_completions_bridge.error_handling(custom_llm_provider or "", model),
)
if rust_response is not None:
return rust_response

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

@ -1,13 +1,22 @@
import base64
from typing import Final
from typing import Final, NoReturn
import httpx
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.rust_bridge import transcription as rust_transcription_bridge
from litellm.rust_bridge.dispatch import PROPAGATE, adispatch, dispatch
from litellm.types.utils import FileTypes, TranscriptionResponse
def _unavailable() -> NoReturn:
raise RuntimeError("Rust audio transcription bridge is unavailable")
async def _aunavailable() -> NoReturn:
_unavailable()
class BedrockAudioTranscriptionRustDispatch:
@staticmethod
def _audio_payload(audio_file: FileTypes) -> dict[str, object]:
@ -43,18 +52,21 @@ class BedrockAudioTranscriptionRustDispatch:
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> TranscriptionResponse:
rust_response: Final = rust_transcription_bridge.transcription(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout=timeout,
rust_response: Final = dispatch(
native=lambda: rust_transcription_bridge.transcription(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout=timeout,
),
python=_unavailable,
route="audio transcription",
errors=PROPAGATE,
)
if rust_response is None:
raise RuntimeError("Rust audio transcription bridge is unavailable")
return TranscriptionResponse(**rust_response)
async def async_audio_transcriptions(
@ -69,16 +81,19 @@ class BedrockAudioTranscriptionRustDispatch:
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> TranscriptionResponse:
rust_response: Final = await rust_transcription_bridge.atranscription(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout=timeout,
rust_response: Final = await adispatch(
native=lambda: rust_transcription_bridge.atranscription(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout=timeout,
),
python=_aunavailable,
route="audio transcription",
errors=PROPAGATE,
)
if rust_response is None:
raise RuntimeError("Rust audio transcription bridge is unavailable")
return TranscriptionResponse(**rust_response)

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

@ -18,6 +18,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
from litellm.rust_bridge.dispatch import adispatch, dispatch
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
@ -423,18 +424,20 @@ class BedrockConverseLLM(BaseAWSLLM):
additional_args=rust_logging_args,
)
if acompletion:
return rust_chat_completions_bridge.achat_completions_or_fallback(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=proxy_endpoint_url,
custom_llm_provider="bedrock",
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
python_fallback=lambda: self.async_completion(
return adispatch(
native=lambda: rust_chat_completions_bridge.achat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=proxy_endpoint_url,
custom_llm_provider="bedrock",
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
),
python=lambda: self.async_completion(
model=model,
messages=messages,
api_base=proxy_endpoint_url,
@ -452,18 +455,25 @@ class BedrockConverseLLM(BaseAWSLLM):
api_key=api_key,
skip_pre_call_logging=True,
),
route="chat_completions",
errors=rust_chat_completions_bridge.error_handling("bedrock", model),
)
rust_response: Final = rust_chat_completions_bridge.chat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=proxy_endpoint_url,
custom_llm_provider="bedrock",
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
rust_response: Final = dispatch(
native=lambda: rust_chat_completions_bridge.chat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=proxy_endpoint_url,
custom_llm_provider="bedrock",
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
),
python=lambda: None,
route="chat_completions",
errors=rust_chat_completions_bridge.error_handling("bedrock", model),
)
if rust_response is not None:
return rust_response

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

@ -2408,8 +2408,10 @@ class BaseLLMHTTPHandler:
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(
from litellm.rust_bridge.dispatch import PYTHON_ON_ERROR, adispatch, async_none
rust_response: Final = await adispatch(
native=lambda: rust_messages_bridge.amessages(
model=model,
body=upstream_body,
api_key=api_key,
@ -2417,13 +2419,11 @@ class BaseLLMHTTPHandler:
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
),
python=async_none,
route="messages",
errors=PYTHON_ON_ERROR,
)
if rust_response is None:
return None
@ -6511,11 +6511,17 @@ class BaseLLMHTTPHandler:
async def _backend_connection():
if _rust_responses_websocket_enabled(custom_llm_provider):
from litellm.rust_bridge import responses_websocket as rust_responses_websocket
from litellm.rust_bridge.dispatch import PYTHON_ON_ERROR, adispatch, async_none
rust_backend: Final = await rust_responses_websocket.connect(
url=ws_url,
headers={str(key): str(value) for key, value in headers.items()},
timeout=timeout,
rust_backend: Final = await adispatch(
native=lambda: rust_responses_websocket.connect(
url=ws_url,
headers={str(key): str(value) for key, value in headers.items()},
timeout=timeout,
),
python=async_none,
route="responses_websocket",
errors=PYTHON_ON_ERROR,
)
if rust_backend is not None:
yield rust_backend

210
litellm/ocr/main.py
View file

@ -7,8 +7,7 @@ import base64
import mimetypes
import os
import re
from collections.abc import Callable, Coroutine, Mapping
from dataclasses import dataclass
from collections.abc import Coroutine, Mapping
from io import IOBase
from typing import Any, Final, cast
@ -18,18 +17,15 @@ import litellm
from litellm._logging import verbose_logger
from litellm.constants import request_timeout
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.azure_ai.ocr.common_utils import (
is_azure_document_intelligence_model,
)
from litellm.llms.azure_ai.ocr.common_utils import is_azure_document_intelligence_model
from litellm.llms.base_llm.ocr.transformation import (
OCR_REQUEST_FORMAT_PARAM,
BaseOCRConfig,
OCRResponse,
parse_ocr_request_format,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.rust_bridge.dispatch import PROPAGATE, adispatch, dispatch
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -38,36 +34,6 @@ base_llm_http_handler = BaseLLMHTTPHandler()
#################################################
@dataclass
class _PreparedOCRRequest:
model: str
document: dict[str, Any]
api_key: str | None
api_base: str | None
custom_llm_provider: str
extra_headers: dict[str, object] | None
provider_config: BaseOCRConfig
optional_params: dict[str, object]
litellm_params: dict[str, object]
effective_timeout: float | httpx.Timeout
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",
"vertex_ai",
}
def _prepare_ocr_request(
model: str,
document: Mapping[str, object],
@ -77,7 +43,7 @@ def _prepare_ocr_request(
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
kwargs: dict[str, object],
) -> _PreparedOCRRequest:
) -> rust_ocr_bridge.PreparedOCRRequest:
litellm_logging_obj: Final = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj"))
litellm_call_id: Final = cast(str | None, kwargs.get("litellm_call_id", None))
@ -174,7 +140,7 @@ def _prepare_ocr_request(
custom_llm_provider=custom_llm_provider,
)
return _PreparedOCRRequest(
return rust_ocr_bridge.PreparedOCRRequest(
model=model,
document=document,
api_key=api_key,
@ -189,154 +155,6 @@ def _prepare_ocr_request(
)
def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native":
return False
if not prepared_request.provider_config.supports_rust_bridge():
return False
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
def _rust_bridge_optional_params(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> dict[str, object]:
optional_params: Final = dict(prepared_request.optional_params)
if prepared_request.custom_llm_provider == "vertex_ai":
vertex_project: Final = (
prepared_request.litellm_params.get("vertex_project")
or prepared_request.litellm_params.get("vertex_ai_project")
or litellm.vertex_project
or resolve_secret("VERTEXAI_PROJECT")
)
vertex_location: Final = (
prepared_request.litellm_params.get("vertex_location")
or prepared_request.litellm_params.get("vertex_ai_location")
or litellm.vertex_location
or resolve_secret("VERTEXAI_LOCATION")
or resolve_secret("VERTEX_LOCATION")
)
if vertex_project is not None:
optional_params["vertex_project"] = vertex_project
if vertex_location is not None:
optional_params["vertex_location"] = vertex_location
return optional_params
def _rust_bridge_api_base(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> str | None:
if prepared_request.api_base is not None:
return prepared_request.api_base
if prepared_request.custom_llm_provider == "azure_ai":
if is_azure_document_intelligence_model(prepared_request.model):
return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
return resolve_secret("AZURE_AI_API_BASE")
return None
def _prepare_rust_ocr_call(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> _PreparedRustOCRCall:
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 (
resolve_api_key(api_key_env_var) if api_key_env_var is not None else None
)
resolved_headers: Final = provider_config.validate_environment(
headers=prepared_request.extra_headers or {},
model=prepared_request.model,
api_key=resolved_api_key,
api_base=prepared_request.api_base,
litellm_params=prepared_request.litellm_params,
)
resolved_complete_url: Final = provider_config.get_complete_url(
api_base=prepared_request.api_base,
model=prepared_request.model,
optional_params=prepared_request.optional_params,
litellm_params=prepared_request.litellm_params,
)
rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key)
rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key)
prepared_request.litellm_logging_obj.pre_call(
input="OCR document processing",
api_key=resolved_api_key,
additional_args={
"complete_input_dict": {
"model": prepared_request.model,
"document": prepared_request.document,
**rust_optional_params,
},
"api_base": resolved_complete_url,
"headers": resolved_headers,
},
)
return _PreparedRustOCRCall(
api_key=resolved_api_key,
api_base=rust_api_base,
headers=cast(dict[str, object], resolved_headers),
optional_params=rust_optional_params,
)
def _run_rust_ocr(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
fallback: Callable[[], OCRResponse | Coroutine[object, object, OCRResponse]],
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
return rust_ocr_bridge.dispatch_ocr(
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
),
call=lambda native, prepared: native(
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_seconds=timeout_to_seconds(prepared_request.effective_timeout),
),
fallback=fallback,
adapt=OCRResponse.model_validate,
model=prepared_request.model,
provider=prepared_request.custom_llm_provider,
eligible=_rust_ocr_supported(prepared_request),
)
async def _run_rust_aocr(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
fallback: Callable[[], Coroutine[object, object, OCRResponse]],
) -> OCRResponse:
return await rust_ocr_bridge.adispatch_ocr(
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
),
call=lambda native, prepared: native(
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_seconds=timeout_to_seconds(prepared_request.effective_timeout),
),
fallback=fallback,
adapt=OCRResponse.model_validate,
model=prepared_request.model,
provider=prepared_request.custom_llm_provider,
eligible=_rust_ocr_supported(prepared_request),
)
@client
async def aocr(
model: str,
@ -453,10 +271,11 @@ async def aocr(
raise ValueError(f"Got an unexpected None response from the OCR API: {response}")
return response
return await _run_rust_aocr(
prepared_request=prepared,
resolve_api_key=get_secret_str,
fallback=python_fallback,
return await adispatch(
native=lambda: rust_ocr_bridge.aattempt_ocr(prepared_request=prepared, resolve_api_key=get_secret_str),
python=python_fallback,
route="ocr",
errors=PROPAGATE,
)
except Exception as e:
raise litellm.exception_type(
@ -714,10 +533,11 @@ def ocr(
litellm_params=prepared.litellm_params,
)
return _run_rust_ocr(
prepared_request=prepared,
resolve_api_key=get_secret_str,
fallback=python_fallback,
return dispatch(
native=lambda: rust_ocr_bridge.attempt_ocr(prepared_request=prepared, resolve_api_key=get_secret_str),
python=python_fallback,
route="ocr",
errors=PROPAGATE,
)
except Exception as e:
raise litellm.exception_type(

12
litellm/rust_bridge/bindings.py
View file

@ -67,9 +67,11 @@ def _exception_class(value: object) -> type[BaseException] | None:
return None
def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None:
declined: Final = _exception_class(_DECLINED.load())
def native_upstream_types() -> tuple[type[BaseException], ...]:
upstream: Final = _exception_class(_UPSTREAM.load())
if declined is None or upstream is None:
return None
return declined, upstream
return () if upstream is None else (upstream,)
def native_declined_types() -> tuple[type[BaseException], ...]:
declined: Final = _exception_class(_DECLINED.load())
return () if declined is None else (declined,)

141
litellm/rust_bridge/chat_completions.py
View file

@ -4,16 +4,12 @@ The Rust core owns the conversation translation, the provider call, and the
response normalization for the subset of `/chat/completions` requests it
accepts. This module only marshals inputs and hands the normalized result to
LiteLLM's existing `ModelResponse` builder.
``None`` means the provider was never called, so the caller is free to serve the
request on the Python path. A failure after the call was issued raises instead:
retrying it there would bill the customer for the same work twice.
"""
from __future__ import annotations
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, Protocol
import httpx
@ -24,19 +20,15 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
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.bindings import UNCHANGED, NativeBinding, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.dispatch import APIErrorMapping, ErrorAction, ErrorHandling
from litellm.rust_bridge.protocols import (
RustAchatCompletions,
RustChatCompletions,
RustChatCompletionsDecline,
)
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointBinding,
EndpointDispatch,
async_none,
)
from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.utils import ModelResponse
@ -94,16 +86,10 @@ def response_logger(
return log
_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,
_CHAT: Final[NativeBinding[RustChatCompletions]] = NativeBinding(lambda native: native.chat_completions)
_ACHAT: Final[NativeBinding[RustAchatCompletions]] = NativeBinding(lambda native: native.achat_completions)
_CHAT_PREFLIGHT: Final[NativeBinding[RustChatCompletionsDecline]] = NativeBinding(
lambda native: native.chat_completions_decline
)
@ -117,14 +103,14 @@ def set_rust_chat_completions(
patching module attributes."""
if not isinstance(chat_completions, Unchanged):
if chat_completions is None:
_CHAT.sync.reset()
_CHAT.reset()
else:
_CHAT.sync.override(chat_completions)
_CHAT.override(chat_completions)
if not isinstance(achat_completions, Unchanged):
if achat_completions is None:
_CHAT.asynchronous.reset()
_ACHAT.reset()
else:
_CHAT.asynchronous.override(achat_completions)
_ACHAT.override(achat_completions)
if not isinstance(decline, Unchanged):
if decline is None:
_CHAT_PREFLIGHT.reset()
@ -193,14 +179,24 @@ def rust_chat_completions_accepts(
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
return _CHAT_PREFLIGHT.accepts(
check=lambda decline: decline(
if not rust_enabled():
return False
decline: Final = _CHAT_PREFLIGHT.load()
if decline is None:
return False
try:
reason: Final = decline(
model=model,
messages=messages,
optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
),
)
)
except Exception as error: # noqa: BLE001 # capability checks perform no provider I/O
verbose_logger.debug("Native chat acceptance check failed: %s", error)
return False
if reason is not None:
verbose_logger.debug("Native chat request is ineligible: %s", reason)
return reason is None
def _build_model_response(
@ -229,14 +225,13 @@ def chat_completions(
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
) -> ModelResponse | None:
) -> DispatchResult[ModelResponse]:
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(
def call(native: RustChatCompletions, timeout_seconds: float | None) -> Mapping[str, object]:
return native(
model=model,
messages=messages,
optional_params=optional_params,
@ -245,10 +240,15 @@ def chat_completions(
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_seconds,
),
fallback=lambda: None,
)
return attempt(
load=_CHAT.load,
enabled=rust_enabled(),
eligible=True,
prepare=lambda: timeout_to_seconds(timeout),
call=call,
adapt=adapt,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -264,14 +264,13 @@ async def achat_completions(
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
) -> ModelResponse | None:
) -> DispatchResult[ModelResponse]:
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(
async def call(native: RustAchatCompletions, timeout_seconds: float | None) -> Mapping[str, object]:
return await native(
model=model,
messages=messages,
optional_params=optional_params,
@ -280,53 +279,21 @@ async def achat_completions(
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
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(
*,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
model_response: ModelResponse,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
python_fallback: Callable[[], Awaitable[object]],
) -> object:
"""Await the Rust path, falling back to the caller's own Python path when
the bridge is unavailable or the call fails.
The caller supplies the fallback, so the bridge stays free of provider
dispatch. This exists because a caller that dispatches asynchronously has
already returned a coroutine by the time a Rust failure surfaces, and so
cannot fall back on its own.
"""
def adapt(rust_response: Mapping[str, object]) -> object:
on_response(rust_response)
return _build_model_response(rust_response, model_response)
return await _CHAT.ainvoke(
return await aattempt(
load=_ACHAT.load,
enabled=rust_enabled(),
eligible=True,
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,
call=call,
adapt=adapt,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
def error_handling(provider: str, model: str) -> ErrorHandling:
return ErrorHandling(
declined=ErrorAction.SKIP,
upstream=APIErrorMapping(provider=provider, model=model),
missing_metadata=ErrorAction.SKIP,
)

137
litellm/rust_bridge/dispatch.py Normal file
View file

@ -0,0 +1,137 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from enum import Enum
from typing import Final, TypeAlias, TypeVar
from litellm._logging import verbose_logger
from litellm.exceptions import APIError
from litellm.rust_bridge.bindings import native_declined_types, native_upstream_types
from litellm.rust_bridge.runtime import DispatchResult, Handled, NativeFailed, NativeSkipped, NativeSkipReason
NativeT = TypeVar("NativeT")
PythonT = TypeVar("PythonT")
class ErrorAction(Enum):
RAISE = "raise"
SKIP = "skip"
@dataclass(frozen=True, slots=True)
class APIErrorMapping:
provider: str
model: str
FailureAction: TypeAlias = ErrorAction | APIErrorMapping
@dataclass(frozen=True, slots=True)
class ErrorHandling:
declined: FailureAction = ErrorAction.RAISE
upstream: FailureAction = ErrorAction.RAISE
unknown: FailureAction = ErrorAction.RAISE
missing_metadata: FailureAction = ErrorAction.RAISE
unexpected: FailureAction = ErrorAction.RAISE
PROPAGATE: Final = ErrorHandling()
PYTHON_ON_ERROR: Final = ErrorHandling(
declined=ErrorAction.SKIP,
upstream=ErrorAction.SKIP,
unknown=ErrorAction.SKIP,
missing_metadata=ErrorAction.SKIP,
unexpected=ErrorAction.SKIP,
)
def _handle_error(error: Exception, action: FailureAction, route: str, reason: NativeSkipReason) -> NativeSkipped:
match action:
case ErrorAction.SKIP:
return NativeSkipped(reason, str(error))
case ErrorAction.RAISE:
raise error
case APIErrorMapping(provider, model):
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 {route}: {message}",
llm_provider=provider,
model=model,
) from error
def _resolve(result: DispatchResult[NativeT], errors: ErrorHandling, route: str) -> Handled[NativeT] | NativeSkipped:
if not isinstance(result, NativeFailed):
return result
declined: Final = native_declined_types()
upstream: Final = native_upstream_types()
if not declined or not upstream:
return _handle_error(result.error, errors.missing_metadata, route, NativeSkipReason.FAILED)
if isinstance(result.error, declined):
return _handle_error(result.error, errors.declined, route, NativeSkipReason.DECLINED)
if isinstance(result.error, upstream):
return _handle_error(result.error, errors.upstream, route, NativeSkipReason.FAILED)
return _handle_error(result.error, errors.unknown, route, NativeSkipReason.FAILED)
def _log_skip(route: str, skipped: NativeSkipped) -> None:
verbose_logger.debug("Native %s skipped (%s): %s", route, skipped.reason.value, skipped.detail or "")
def dispatch(
*,
native: Callable[[], DispatchResult[NativeT]],
python: Callable[[], PythonT],
route: str,
errors: ErrorHandling,
) -> NativeT | PythonT:
try:
attempted: Final = native()
except Exception as error: # noqa: BLE001 # preserve declared handling of loading and adaptation failures
unexpected: Final = _handle_error(error, errors.unexpected, route, NativeSkipReason.FAILED)
_log_skip(route, unexpected)
return python()
result: Final = _resolve(attempted, errors, route)
match result:
case Handled(value):
return value
case NativeSkipped():
_log_skip(route, result)
return python()
async def adispatch(
*,
native: Callable[[], Awaitable[DispatchResult[NativeT]]],
python: Callable[[], Awaitable[PythonT]],
route: str,
errors: ErrorHandling,
) -> NativeT | PythonT:
try:
attempted: Final = await native()
except Exception as error: # noqa: BLE001 # preserve declared handling of loading and adaptation failures
unexpected: Final = _handle_error(error, errors.unexpected, route, NativeSkipReason.FAILED)
_log_skip(route, unexpected)
return await python()
result: Final = _resolve(attempted, errors, route)
match result:
case Handled(value):
return value
case NativeSkipped():
_log_skip(route, result)
return await python()
async def async_none() -> None:
return None

50
litellm/rust_bridge/messages.py
View file

@ -6,25 +6,13 @@ from typing import Final
import httpx
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, 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.runtime import DispatchResult, aattempt, attempt, identity
from litellm.rust_bridge.timeouts import timeout_to_seconds
_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,
)
_MESSAGES: Final[NativeBinding[RustMessages]] = NativeBinding(lambda native: native.messages)
_AMESSAGES: Final[NativeBinding[RustAmessages]] = NativeBinding(lambda native: native.amessages)
def set_rust_messages(
@ -34,22 +22,22 @@ def set_rust_messages(
) -> None:
if not isinstance(messages, Unchanged):
if messages is None:
_MESSAGES.sync.reset()
_MESSAGES.reset()
else:
_MESSAGES.sync.override(messages)
_MESSAGES.override(messages)
if not isinstance(amessages, Unchanged):
if amessages is None:
_MESSAGES.asynchronous.reset()
_AMESSAGES.reset()
else:
_MESSAGES.asynchronous.override(amessages)
_AMESSAGES.override(amessages)
def load_rust_messages() -> RustMessages | None:
return _MESSAGES.sync.load()
return _MESSAGES.load()
def load_rust_amessages() -> RustAmessages | None:
return _MESSAGES.asynchronous.load()
return _AMESSAGES.load()
def messages(
@ -61,8 +49,11 @@ def messages(
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
return _MESSAGES.invoke(
) -> DispatchResult[dict[str, object]]:
return attempt(
load=_MESSAGES.load,
enabled=True,
eligible=True,
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_messages, timeout_seconds: rust_messages(
model=model,
@ -73,9 +64,7 @@ def messages(
extra_headers=extra_headers,
timeout_seconds=timeout_seconds,
),
fallback=lambda: None,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -88,8 +77,11 @@ async def amessages(
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
return await _MESSAGES.ainvoke(
) -> DispatchResult[dict[str, object]]:
return await aattempt(
load=_AMESSAGES.load,
enabled=True,
eligible=True,
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_amessages, timeout_seconds: rust_amessages(
model=model,
@ -100,7 +92,5 @@ async def amessages(
extra_headers=extra_headers,
timeout_seconds=timeout_seconds,
),
fallback=async_none,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)

261
litellm/rust_bridge/ocr.py
View file

@ -1,31 +1,59 @@
"""Thin Python wrapper for the native Rust OCR bridge."""
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, TypeVar
from collections.abc import Callable
from dataclasses import dataclass
from typing import Final
from . import configuration as _configuration
from .bindings import UNCHANGED, Unchanged
from .protocols import RustAocr, RustOcr
from .runtime import (
BridgeErrorContext,
EndpointDispatch,
NativeErrorPolicy,
)
import httpx
from pydantic import TypeAdapter
rust_ocr_enabled = _configuration.rust_ocr_enabled
rust = _configuration.rust
ResultT = TypeVar("ResultT")
RequestT = TypeVar("RequestT")
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.azure_ai.ocr.common_utils import is_azure_document_intelligence_model
from litellm.llms.base_llm.ocr.transformation import OCR_REQUEST_FORMAT_PARAM, BaseOCRConfig, OCRResponse
from litellm.rust_bridge import configuration as _configuration
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged
from litellm.rust_bridge.protocols import RustAocr, RustOcr
from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt
from litellm.rust_bridge.timeouts import timeout_to_seconds
rust: Final = _configuration.rust
rust_ocr_enabled: Final = _configuration.rust_ocr_enabled
_OCR: Final[NativeBinding[RustOcr]] = NativeBinding(lambda native: native.ocr)
_AOCR: Final[NativeBinding[RustAocr]] = NativeBinding(lambda native: native.aocr)
_HEADERS: Final = TypeAdapter(dict[str, object])
_OCR: Final[EndpointDispatch[RustOcr, RustAocr]] = EndpointDispatch.native(
route="ocr",
sync=lambda native: native.ocr,
asynchronous=lambda native: native.aocr,
enabled=_configuration.rust_ocr_enabled,
error_policy=NativeErrorPolicy.PROPAGATE,
@dataclass(frozen=True, slots=True)
class PreparedOCRRequest:
model: str
document: dict[str, object]
api_key: str | None
api_base: str | None
custom_llm_provider: str
extra_headers: dict[str, object] | None
provider_config: BaseOCRConfig
optional_params: dict[str, object]
litellm_params: dict[str, object]
effective_timeout: float | httpx.Timeout
litellm_logging_obj: LiteLLMLoggingObj
@dataclass(frozen=True, slots=True)
class _PreparedRustOCRCall:
api_key: str | None
api_base: str | None
headers: dict[str, object]
optional_params: dict[str, object]
_RUST_OCR_PROVIDERS: Final = frozenset(
{
"mistral",
"azure_ai",
"vertex_ai",
}
)
@ -36,59 +64,168 @@ def set_rust_ocr(
) -> None:
if not isinstance(ocr, Unchanged):
if ocr is None:
_OCR.sync.reset()
_OCR.reset()
else:
_OCR.sync.override(ocr)
_OCR.override(ocr)
if not isinstance(aocr, Unchanged):
if aocr is None:
_OCR.asynchronous.reset()
_AOCR.reset()
else:
_OCR.asynchronous.override(aocr)
_AOCR.override(aocr)
def load_rust_ocr() -> RustOcr | None:
return _OCR.sync.load()
return _OCR.load()
def load_rust_aocr() -> RustAocr | None:
return _OCR.asynchronous.load()
return _AOCR.load()
def dispatch_ocr(
*,
prepare: Callable[[], RequestT],
call: Callable[[RustOcr, RequestT], Mapping[str, object]],
fallback: Callable[[], ResultT],
adapt: Callable[[Mapping[str, object]], ResultT],
model: str,
provider: str,
eligible: bool,
) -> ResultT:
return _OCR.invoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=BridgeErrorContext(provider=provider, model=model),
eligible=eligible,
def _rust_ocr_supported(prepared_request: PreparedOCRRequest) -> bool:
if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native":
return False
if not prepared_request.provider_config.supports_rust_bridge():
return False
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
def _rust_bridge_optional_params(
prepared_request: PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> dict[str, object]:
if prepared_request.custom_llm_provider != "vertex_ai":
return prepared_request.optional_params
vertex_project: Final = (
prepared_request.litellm_params.get("vertex_project")
or prepared_request.litellm_params.get("vertex_ai_project")
or litellm.vertex_project
or resolve_secret("VERTEXAI_PROJECT")
)
vertex_location: Final = (
prepared_request.litellm_params.get("vertex_location")
or prepared_request.litellm_params.get("vertex_ai_location")
or litellm.vertex_location
or resolve_secret("VERTEXAI_LOCATION")
or resolve_secret("VERTEX_LOCATION")
)
return {
**prepared_request.optional_params,
**{
name: value
for name, value in (("vertex_project", vertex_project), ("vertex_location", vertex_location))
if value is not None
},
}
def _rust_bridge_api_base(
prepared_request: PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> str | None:
if prepared_request.api_base is not None:
return prepared_request.api_base
if prepared_request.custom_llm_provider == "azure_ai":
if is_azure_document_intelligence_model(prepared_request.model):
return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
return resolve_secret("AZURE_AI_API_BASE")
return None
def _prepare_rust_ocr_call(
prepared_request: PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> _PreparedRustOCRCall:
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 (
resolve_api_key(api_key_env_var) if api_key_env_var is not None else None
)
resolved_headers: Final = _HEADERS.validate_python(
provider_config.validate_environment(
headers=prepared_request.extra_headers or {},
model=prepared_request.model,
api_key=resolved_api_key,
api_base=prepared_request.api_base,
litellm_params=prepared_request.litellm_params,
)
)
resolved_complete_url: Final = provider_config.get_complete_url(
api_base=prepared_request.api_base,
model=prepared_request.model,
optional_params=prepared_request.optional_params,
litellm_params=prepared_request.litellm_params,
)
rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key)
rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key)
prepared_request.litellm_logging_obj.pre_call(
input="OCR document processing",
api_key=resolved_api_key,
additional_args={
"complete_input_dict": {
"model": prepared_request.model,
"document": prepared_request.document,
**rust_optional_params,
},
"api_base": resolved_complete_url,
"headers": resolved_headers,
},
)
return _PreparedRustOCRCall(
api_key=resolved_api_key,
api_base=rust_api_base,
headers=resolved_headers,
optional_params=rust_optional_params,
)
async def adispatch_ocr(
*,
prepare: Callable[[], RequestT],
call: Callable[[RustAocr, RequestT], Awaitable[Mapping[str, object]]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[Mapping[str, object]], ResultT],
model: str,
provider: str,
eligible: bool,
) -> ResultT:
return await _OCR.ainvoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=BridgeErrorContext(provider=provider, model=model),
eligible=eligible,
def attempt_ocr(
prepared_request: PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> DispatchResult[OCRResponse]:
return attempt(
load=_OCR.load,
enabled=rust_ocr_enabled(),
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
),
call=lambda native, prepared: native(
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_seconds=timeout_to_seconds(prepared_request.effective_timeout),
),
adapt=OCRResponse.model_validate,
eligible=_rust_ocr_supported(prepared_request),
)
async def aattempt_ocr(
prepared_request: PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> DispatchResult[OCRResponse]:
return await aattempt(
load=_AOCR.load,
enabled=rust_ocr_enabled(),
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
),
call=lambda native, prepared: native(
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_seconds=timeout_to_seconds(prepared_request.effective_timeout),
),
adapt=OCRResponse.model_validate,
eligible=_rust_ocr_supported(prepared_request),
)

44
litellm/rust_bridge/responses_websocket.py
View file

@ -7,24 +7,17 @@ from typing import Final
import httpx
from websockets.exceptions import ConnectionClosedOK
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.protocols import (
RustResponsesWebSocket,
RustResponsesWebSocketConnection,
)
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointBinding,
async_none,
identity,
)
from litellm.rust_bridge.runtime import DispatchResult, aattempt
from litellm.rust_bridge.timeouts import timeout_to_seconds
_RESPONSES_WEBSOCKET: Final[EndpointBinding[RustResponsesWebSocketConnection]] = EndpointBinding.native(
route="responses_websocket",
select=lambda native: native.ResponsesWebSocketConnection,
enabled=rust_enabled,
_RESPONSES_WEBSOCKET: Final[NativeBinding[RustResponsesWebSocketConnection]] = NativeBinding(
lambda native: native.ResponsesWebSocketConnection,
)
@ -61,19 +54,16 @@ async def connect(
url: str,
headers: dict[str, str],
timeout: float | httpx.Timeout | None,
) -> _ConnectionAdapter | None:
try:
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 # preserve the existing WebSocket connection fallback
return None
return None if connection is None else _ConnectionAdapter(connection)
) -> DispatchResult[_ConnectionAdapter]:
return await aattempt(
load=_RESPONSES_WEBSOCKET.load,
enabled=rust_enabled(),
eligible=True,
prepare=lambda: timeout_to_seconds(timeout),
call=lambda connection_type, timeout_seconds: connection_type.connect(
url=url,
headers=headers,
timeout_seconds=timeout_seconds,
),
adapt=_ConnectionAdapter,
)

528
litellm/rust_bridge/runtime.py
View file

@ -1,39 +1,22 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from dataclasses import dataclass
from enum import Enum
from typing import Final, Generic, NoReturn, Protocol, TypeAlias, TypeVar
from litellm.exceptions import APIError
from litellm.rust_bridge.bindings import (
UNCHANGED,
NativeBinding,
Unchanged,
native_exception_types,
)
from litellm.rust_bridge.protocols import NativeModule
from typing import Final, Generic, TypeAlias, TypeVar
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 PythonFallbackReason(Enum):
NATIVE_DISABLED = "native_disabled"
NATIVE_UNAVAILABLE = "native_unavailable"
NATIVE_DECLINED = "native_declined"
class NativeErrorPolicy(Enum):
TRANSLATE = "translate"
PROPAGATE = "propagate"
class NativeSkipReason(Enum):
DISABLED = "disabled"
INELIGIBLE = "ineligible"
UNAVAILABLE = "unavailable"
DECLINED = "declined"
FAILED = "failed"
@dataclass(frozen=True, slots=True)
@ -42,470 +25,65 @@ class Handled(Generic[ResultT]):
@dataclass(frozen=True, slots=True)
class PythonFallback:
reason: PythonFallbackReason
class NativeSkipped:
reason: NativeSkipReason
detail: str | None = None
DispatchResult: TypeAlias = Handled[ResultT] | PythonFallback
@dataclass(frozen=True, slots=True)
class BridgeErrorContext:
provider: str
model: str
class NativeFailed:
error: Exception
class RustEnablement(Protocol):
def __call__(self) -> bool: ...
DispatchResult: TypeAlias = Handled[ResultT] | NativeSkipped | NativeFailed
@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)
def _select(load: Callable[[], BindingT | None], enabled: bool, eligible: bool) -> BindingT | NativeSkipped:
if not enabled:
return NativeSkipped(NativeSkipReason.DISABLED)
if not eligible:
return NativeSkipped(NativeSkipReason.INELIGIBLE)
binding: Final = load()
return NativeSkipped(NativeSkipReason.UNAVAILABLE) if binding is None else binding
@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,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> DispatchResult[ResultT]:
binding_or_fallback: Final = self._binding_or_python_fallback(
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return binding_or_fallback
preflight_result: Final = preflight() if preflight is not None else None
if preflight_result is not None:
return preflight_result
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,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> DispatchResult[ResultT]:
binding_or_fallback: Final = self._binding_or_python_fallback(
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return binding_or_fallback
preflight_result: Final = preflight() if preflight is not None else None
if preflight_result is not None:
return preflight_result
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,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
result: Final = self._attempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
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,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
result: Final = await self._aattempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
return await fallback()
def assess(
self,
*,
check: Callable[[BindingT], str | None],
) -> PythonFallback | None:
binding: Final = self._binding_or_python_fallback(eligible=True)
if isinstance(binding, PythonFallback):
return binding
reason: Final = check(binding)
return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, reason) if reason is not None else None
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,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
result: Final = self._attempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
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,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
result: Final = await self._aattempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
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)
def attempt(
*,
load: Callable[[], BindingT | None],
enabled: bool,
eligible: bool,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
) -> DispatchResult[ResultT]:
binding: Final = _select(load, enabled, eligible)
if isinstance(binding, NativeSkipped):
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
try:
value: Final = call(binding, prepare())
except Exception as error: # noqa: BLE001 # orchestration applies the endpoint's declared error policy
return NativeFailed(error)
return Handled(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,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
return self.sync.invoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
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,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
return await self.asynchronous.ainvoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
def require(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[SyncBindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
return self.sync.require(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
async def arequire(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
return await self.asynchronous.arequire(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
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 _required_reason(reason: PythonFallbackReason) -> str:
match reason:
case PythonFallbackReason.NATIVE_DISABLED:
return "is disabled"
case PythonFallbackReason.NATIVE_UNAVAILABLE:
return "is unavailable"
case PythonFallbackReason.NATIVE_DECLINED:
return "declined the request"
def always_enabled() -> bool:
return True
async def aattempt(
*,
load: Callable[[], BindingT | None],
enabled: bool,
eligible: bool,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
) -> DispatchResult[ResultT]:
binding: Final = _select(load, enabled, eligible)
if isinstance(binding, NativeSkipped):
return binding
try:
value: Final = await call(binding, prepare())
except Exception as error: # noqa: BLE001 # orchestration applies the endpoint's declared error policy
return NativeFailed(error)
return Handled(adapt(value))
def identity(value: ResultT) -> ResultT:
return value
async def async_none() -> None:
return None

50
litellm/rust_bridge/transcription.py
View file

@ -4,25 +4,13 @@ from typing import Final
import httpx
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, 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.runtime import DispatchResult, aattempt, attempt, identity
from litellm.rust_bridge.timeouts import timeout_to_seconds
_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,
)
_TRANSCRIPTION: Final[NativeBinding[RustTranscription]] = NativeBinding(lambda native: native.transcription)
_ATRANSCRIPTION: Final[NativeBinding[RustAtranscription]] = NativeBinding(lambda native: native.atranscription)
def configure_rust_transcription(
@ -32,22 +20,22 @@ def configure_rust_transcription(
) -> None:
if not isinstance(transcription, Unchanged):
if transcription is None:
_TRANSCRIPTION.sync.reset()
_TRANSCRIPTION.reset()
else:
_TRANSCRIPTION.sync.override(transcription)
_TRANSCRIPTION.override(transcription)
if not isinstance(atranscription, Unchanged):
if atranscription is None:
_TRANSCRIPTION.asynchronous.reset()
_ATRANSCRIPTION.reset()
else:
_TRANSCRIPTION.asynchronous.override(atranscription)
_ATRANSCRIPTION.override(atranscription)
def load_rust_transcription() -> RustTranscription | None:
return _TRANSCRIPTION.sync.load()
return _TRANSCRIPTION.load()
def load_rust_atranscription() -> RustAtranscription | None:
return _TRANSCRIPTION.asynchronous.load()
return _ATRANSCRIPTION.load()
def transcription(
@ -60,8 +48,11 @@ def transcription(
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
return _TRANSCRIPTION.invoke(
) -> DispatchResult[dict[str, object]]:
return attempt(
load=_TRANSCRIPTION.load,
enabled=True,
eligible=True,
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_transcription, timeout_seconds: rust_transcription(
model=model,
@ -73,9 +64,7 @@ def transcription(
optional_params=optional_params,
timeout_seconds=timeout_seconds,
),
fallback=lambda: None,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -89,8 +78,11 @@ async def atranscription(
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
return await _TRANSCRIPTION.ainvoke(
) -> DispatchResult[dict[str, object]]:
return await aattempt(
load=_ATRANSCRIPTION.load,
enabled=True,
eligible=True,
prepare=lambda: timeout_to_seconds(timeout),
call=lambda rust_atranscription, timeout_seconds: rust_atranscription(
model=model,
@ -102,7 +94,5 @@ async def atranscription(
optional_params=optional_params,
timeout_seconds=timeout_seconds,
),
fallback=async_none,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)

6
ruff-strict-budget.json
View file

@ -156,7 +156,7 @@
"limit": 215
},
"PLW0603": {
"limit": 186
"limit": 184
},
"PLW1508": {
"limit": 190
@ -231,7 +231,7 @@
"limit": 5
},
"TID251": {
"limit": 1031
"limit": 1029
},
"TRY002": {
"limit": 524
@ -246,7 +246,7 @@
"limit": 109
},
"TRY300": {
"limit": 848
"limit": 846
},
"UP028": {
"limit": 2

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

@ -9,6 +9,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.runtime import Handled, NativeSkipped, NativeSkipReason
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -133,7 +134,7 @@ def test_load_rust_amessages_returns_injected_impl():
assert rust_messages.load_rust_amessages() is bridge
def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch):
def test_messages_wrapper_reports_unavailable(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge.bindings"),
"get_native_bridge",
@ -150,7 +151,7 @@ def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch):
extra_headers={},
timeout=30.0,
)
assert result is None
assert result == NativeSkipped(NativeSkipReason.UNAVAILABLE)
def test_messages_wrapper_forwards_args_and_converts_timeout():
@ -168,7 +169,7 @@ def test_messages_wrapper_forwards_args_and_converts_timeout():
timeout=httpx.Timeout(600.0, read=42.0),
)
assert response == FAKE_MESSAGES_RESPONSE
assert response == Handled(FAKE_MESSAGES_RESPONSE)
assert bridge.calls[0] == {
"model": "claude-sonnet-4-5",
"body": REQUEST_BODY,
@ -196,7 +197,7 @@ async def test_amessages_wrapper_forwards_args():
timeout=12.5,
)
assert response == FAKE_MESSAGES_RESPONSE
assert response == Handled(FAKE_MESSAGES_RESPONSE)
assert bridge.calls[0]["model"] == "claude-sonnet-4-5"
assert bridge.calls[0]["timeout_seconds"] == 12.5
@ -244,7 +245,6 @@ async def test_gate_falls_back_to_python_when_bridge_raises():
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate()
assert response is None
assert bridge.calls == 1

6
tests/test_litellm/ocr/test_ocr_native_format.py
View file

@ -12,13 +12,13 @@ import pytest
import litellm
from litellm.llms.azure_ai.ocr.cohere_parse_transformation import AzureAICohereParseConfig
from litellm.llms.cohere.ocr.transformation import CohereParseConfig
from litellm.ocr.main import _PreparedOCRRequest, _rust_ocr_supported
from litellm.rust_bridge.ocr import PreparedOCRRequest, _rust_ocr_supported
DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"}
def _prepared(optional_params: dict[str, object]) -> _PreparedOCRRequest:
return _PreparedOCRRequest(
def _prepared(optional_params: dict[str, object]) -> PreparedOCRRequest:
return PreparedOCRRequest(
model="doc-intelligence/prebuilt-layout",
document=dict(DOCUMENT),
api_key="fake-key",

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

@ -12,6 +12,7 @@ import pytest
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge import configuration
from litellm.rust_bridge.runtime import Handled
from litellm.rust_bridge.timeouts import timeout_to_seconds
# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr`
@ -202,7 +203,7 @@ def build_prepared_request(
litellm_params: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = 12.5,
) -> Any:
return ocr_main._PreparedOCRRequest(
return rust_bridge.PreparedOCRRequest(
model=model,
document=document,
api_key=api_key,
@ -399,99 +400,13 @@ def test_timeout_to_seconds_handles_float_timeout_and_none():
assert timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0
def test_bridge_wrapper_forwards_prepared_args_and_wraps_response():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
response = rust_bridge.dispatch_ocr(
prepare=lambda: 12.5,
call=lambda native, timeout: native(
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_seconds=timeout,
),
fallback=lambda: pytest.fail("unexpected Python fallback"),
adapt=dict,
model="mistral-ocr-latest",
provider="mistral",
eligible=True,
)
assert response == FAKE_OCR_RESPONSE
call = bridge.calls[0]
assert call == {
"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_seconds": 12.5,
}
@pytest.mark.asyncio
async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response():
bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge.set_rust_ocr(aocr=bridge)
async def unexpected_fallback():
pytest.fail("unexpected Python fallback")
response = await rust_bridge.adispatch_ocr(
prepare=lambda: 42.0,
call=lambda native, timeout: native(
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_seconds=timeout,
),
fallback=unexpected_fallback,
adapt=dict,
model="mistral-ocr-maas",
provider="vertex_ai",
eligible=True,
)
assert response == FAKE_OCR_RESPONSE
assert bridge.calls[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_seconds": 42.0,
}
def test_run_rust_ocr_prepares_request_and_wraps_response():
bridge = RecordingBridge()
logging_obj = RecordingLogging()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
response = ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
response = rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
logging_obj=logging_obj,
api_base="https://proxy.internal",
@ -502,6 +417,8 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
resolve_api_key=lambda _name: None,
)
assert isinstance(response, Handled)
response = response.value
assert isinstance(response, OCRResponse)
assert response.pages[0].markdown == "hello world"
assert bridge.calls[0] == {
@ -524,8 +441,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(api_key=None, timeout=None),
resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None,
)
@ -541,8 +457,7 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver():
def _resolver(name: str) -> str | None:
raise AssertionError(f"resolver should not be called for {name}")
ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
api_key="sk-explicit",
timeout=None,
@ -563,8 +478,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
resolver_calls.append(name)
return "sk-provider-env"
ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"),
model="provider-ocr-model",
@ -583,8 +497,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
custom_llm_provider="vertex_ai",
model="mistral-ocr-maas",
@ -617,8 +530,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana
"VERTEXAI_LOCATION": "us-east5",
}.get(name)
ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
custom_llm_provider="vertex_ai",
model="mistral-ocr-maas",
@ -636,8 +548,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
custom_llm_provider="azure_ai",
model="pixtral-12b-2409",
@ -655,8 +566,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint():
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
custom_llm_provider="azure_ai",
model="doc-intelligence/prebuilt-layout",
@ -677,8 +587,7 @@ def test_run_rust_ocr_runs_pre_call_logging():
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
ocr_main._run_rust_ocr(
fallback=lambda: pytest.fail("unexpected Python fallback"),
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
logging_obj=logging_obj,
api_base="https://api.mistral.ai/v1",
@ -851,7 +760,7 @@ async def test_ocr_fallback_skips_native_preparation(
def unexpected_preparation(*_args: object, **_kwargs: object) -> None:
pytest.fail("Python fallback must not resolve native credentials or emit native pre_call")
monkeypatch.setattr(ocr_main, "_prepare_rust_ocr_call", unexpected_preparation)
monkeypatch.setattr(rust_bridge, "_prepare_rust_ocr_call", unexpected_preparation)
monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fallback)
response: Final = (

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

@ -4,6 +4,7 @@ import pytest
from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled
from litellm.rust_bridge import configuration, responses_websocket
from litellm.rust_bridge.runtime import Handled, NativeSkipped, NativeSkipReason, NativeFailed
class _FakeNativeConnection:
@ -64,18 +65,15 @@ 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:
async def test_bridge_reports_unavailable(monkeypatch: pytest.MonkeyPatch) -> None:
configuration.rust(True)
responses_websocket._RESPONSES_WEBSOCKET.override(None)
assert (
await responses_websocket.connect(
url="wss://example.test/responses",
headers={},
timeout=None,
)
is None
)
assert await responses_websocket.connect(
url="wss://example.test/responses",
headers={},
timeout=None,
) == NativeSkipped(NativeSkipReason.UNAVAILABLE)
@pytest.mark.asyncio
@ -91,7 +89,8 @@ async def test_enabled_bridge_connects_and_adapts_socket(
timeout=1.0,
)
assert connection is not None
assert isinstance(connection, Handled)
connection = connection.value
await connection.send("response.create")
assert await connection.recv() == "response.completed"
await connection.close()
@ -110,15 +109,9 @@ class _FailingNativeBridge:
@pytest.mark.asyncio
async def test_connection_failure_preserves_python_fallback() -> None:
async def test_connection_failure_is_reported_to_orchestration() -> 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
)
result = await responses_websocket.connect(url="wss://example.test/responses", headers={}, timeout=None)
assert isinstance(result, NativeFailed)
assert str(result.error) == "connection failed"

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

@ -48,7 +48,8 @@ def test_native_exception_types_reject_non_exception_classes(monkeypatch: pytest
native: Final = SimpleNamespace(RustBridgeDeclined=invalid, RustUpstreamError=RuntimeError)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
assert bindings.native_exception_types() is None
assert bindings.native_declined_types() == ()
assert bindings.native_upstream_types() == (RuntimeError,)
@pytest.mark.parametrize(
@ -59,10 +60,6 @@ def test_native_exception_types_reject_non_exception_classes(monkeypatch: pytest
"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)",
@ -85,14 +82,8 @@ def test_selectors_are_checked_by_type_checker(tmp_path: Path, expression: str,
"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"
@ -117,4 +108,6 @@ def test_selectors_are_checked_by_type_checker(tmp_path: Path, expression: str,
)
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)]
assert [(item["rule"], item["range"]["start"]["line"]) for item in diagnostics] == [
(expected_rule, len(source.read_text().splitlines()) - 1)
]

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

@ -12,6 +12,7 @@ import pytest
import litellm
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge import chat_completions as bridge
from litellm.rust_bridge.runtime import Handled, NativeSkipped, NativeFailed
from litellm.types.utils import ModelResponse
RUST_RESPONSE = {
@ -250,7 +251,8 @@ class TestSyncCall:
result = bridge.chat_completions(**_call_kwargs(model_response))
assert result is not None
assert isinstance(result, Handled)
result = result.value
assert result.choices[0].message.content == "hello from rust"
assert result.choices[0].finish_reason == "stop"
assert result.model == "claude-sonnet-4-5-20260101"
@ -266,14 +268,14 @@ class TestSyncCall:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert native.calls[0]["timeout_seconds"] == 30.0
def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch):
def test_reports_unavailable_bridge(self, monkeypatch):
_hide_native_bridge(monkeypatch)
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
assert isinstance(bridge.chat_completions(**_call_kwargs(ModelResponse())), NativeSkipped)
def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch):
def test_reports_native_decline_to_orchestration(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming")))
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
assert isinstance(bridge.chat_completions(**_call_kwargs(ModelResponse())), NativeFailed)
class TestAsyncCall:
@ -281,134 +283,18 @@ class TestAsyncCall:
async def test_builds_a_model_response(self):
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall())
result = await bridge.achat_completions(**_call_kwargs(ModelResponse()))
assert result is not None
assert isinstance(result, Handled)
result = result.value
assert result.choices[0].message.content == "hello from rust"
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
@pytest.mark.asyncio
async def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch):
async def test_reports_unavailable_bridge(self, monkeypatch):
_hide_native_bridge(monkeypatch)
assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None
assert isinstance(await bridge.achat_completions(**_call_kwargs(ModelResponse())), NativeSkipped)
@pytest.mark.asyncio
async def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch):
async def test_reports_native_decline_to_orchestration(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")))
assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None
class TestAsyncFallbackWrapper:
@pytest.mark.asyncio
async def test_returns_the_rust_response_without_running_the_fallback(self):
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall())
ran = []
async def fallback():
ran.append(True)
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result.choices[0].message.content == "hello from rust"
assert ran == []
@pytest.mark.asyncio
async def test_runs_the_fallback_when_the_core_declines(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")))
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
@pytest.mark.asyncio
async def test_runs_the_fallback_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
class TestFailureClassification:
"""A failure the provider already saw must not be retried on the Python
path: it would bill the customer for the same work twice."""
@pytest.fixture(autouse=True)
def _native_exceptions(self, monkeypatch):
_fake_native_bridge(monkeypatch)
def test_a_decline_falls_back_because_nothing_was_sent(self):
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming")))
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
def test_an_upstream_failure_is_surfaced_with_its_status(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited")))
with pytest.raises(APIError) as raised:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert raised.value.status_code == 429
assert "rate limited" in str(raised.value)
def test_a_transport_failure_with_no_response_surfaces_as_a_500(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset")))
with pytest.raises(APIError) as raised:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert raised.value.status_code == 500
def test_an_unrecognized_error_is_not_swallowed(self):
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=RuntimeError("something else")))
with pytest.raises(RuntimeError):
bridge.chat_completions(**_call_kwargs(ModelResponse()))
@pytest.mark.asyncio
async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom")))
ran = []
async def fallback():
ran.append(True)
return "python"
with pytest.raises(APIError):
await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert ran == [], "a request the provider already served must not be re-issued"
@pytest.mark.asyncio
async def test_the_async_wrapper_falls_back_on_a_decline(self):
bridge.set_rust_chat_completions(
achat_completions=_RecordingAsyncCall(error=_FakeDeclined("blank message text"))
)
async def fallback():
return "python"
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"
)
assert isinstance(await bridge.achat_completions(**_call_kwargs(ModelResponse())), NativeFailed)

205
tests/test_litellm/rust_bridge/test_dispatch.py Normal file
View file

@ -0,0 +1,205 @@
from __future__ import annotations
import asyncio
import logging
from types import SimpleNamespace
from typing import Final
import pytest
from litellm.exceptions import APIError
from litellm.rust_bridge import bindings
from litellm.rust_bridge.chat_completions import error_handling
from litellm.rust_bridge.dispatch import PROPAGATE, PYTHON_ON_ERROR, adispatch, dispatch
from litellm.rust_bridge.runtime import DispatchResult, Handled, NativeFailed, NativeSkipped, NativeSkipReason
class Declined(Exception):
pass
class Upstream(Exception):
pass
@pytest.fixture(autouse=True)
def native_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(
bindings, "get_native_bridge", lambda: SimpleNamespace(RustBridgeDeclined=Declined, RustUpstreamError=Upstream)
)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.parametrize("reason", tuple(NativeSkipReason))
async def test_shared_dispatch_calls_python_once_and_logs_skip(
asynchronous: bool, reason: NativeSkipReason, caplog: pytest.LogCaptureFixture
) -> None:
caplog.set_level(logging.DEBUG, logger="LiteLLM")
calls: Final[list[str]] = []
def native() -> DispatchResult[str]:
calls.append("native")
return NativeSkipped(reason, "diagnostic detail")
async def anative() -> DispatchResult[str]:
return native()
def python() -> str:
calls.append("python")
return "python response"
async def apython() -> str:
return python()
result: Final = (
await adispatch(native=anative, python=apython, route="test", errors=PROPAGATE)
if asynchronous
else dispatch(native=native, python=python, route="test", errors=PROPAGATE)
)
assert result == "python response"
assert calls == ["native", "python"]
assert f"Native test skipped ({reason.value}): diagnostic detail" in caplog.text
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_native_success_does_not_run_python_even_when_value_is_none(asynchronous: bool) -> None:
async def native() -> DispatchResult[None]:
return Handled(None)
def python() -> str:
pytest.fail("handled results must not run Python")
async def apython() -> str:
return python()
result: Final = (
await adispatch(native=native, python=apython, route="test", errors=PYTHON_ON_ERROR)
if asynchronous
else dispatch(native=lambda: Handled(None), python=python, route="test", errors=PYTHON_ON_ERROR)
)
assert result is None
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.parametrize("policy", ("chat", "propagate", "python"))
@pytest.mark.parametrize("kind", ("declined", "upstream", "unknown", "unexpected", "missing"))
async def test_declarations_preserve_endpoint_error_behavior(
monkeypatch: pytest.MonkeyPatch, asynchronous: bool, policy: str, kind: str
) -> None:
if kind == "missing":
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
error: Final = (
Declined("unsupported")
if kind == "declined"
else Upstream(429, "rate limited")
if kind == "upstream"
else RuntimeError("failed")
)
rules: Final = (
error_handling("anthropic", "model")
if policy == "chat"
else PYTHON_ON_ERROR
if policy == "python"
else PROPAGATE
)
calls: Final[list[str]] = []
def native() -> DispatchResult[str]:
if kind == "unexpected":
raise error
return NativeFailed(error)
async def anative() -> DispatchResult[str]:
return native()
def python() -> str:
calls.append("python")
return "python response"
async def apython() -> str:
return python()
async def run() -> str:
if asynchronous:
return await adispatch(native=anative, python=apython, route="chat_completions", errors=rules)
return dispatch(native=native, python=python, route="chat_completions", errors=rules)
if policy == "python" or (policy == "chat" and kind in ("declined", "missing")):
assert await run() == "python response"
assert calls == ["python"]
elif policy == "chat" and kind == "upstream":
with pytest.raises(APIError) as caught:
await run()
assert caught.value.status_code == 429
assert caught.value.model == "model"
assert caught.value.llm_provider == "anthropic"
assert caught.value.__cause__ is error
assert calls == []
else:
with pytest.raises(type(error)) as caught_original:
await run()
assert caught_original.value is error
assert calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_python_failure_is_never_reclassified_as_native_failure(asynchronous: bool) -> None:
error: Final = RuntimeError("Python failed")
calls: Final[list[str]] = []
async def native() -> DispatchResult[str]:
return NativeSkipped(NativeSkipReason.UNAVAILABLE)
def python() -> str:
calls.append("python")
raise error
async def apython() -> str:
return python()
async def run() -> str:
if asynchronous:
return await adispatch(native=native, python=apython, route="test", errors=PYTHON_ON_ERROR)
return dispatch(
native=lambda: NativeSkipped(NativeSkipReason.UNAVAILABLE),
python=python,
route="test",
errors=PYTHON_ON_ERROR,
)
with pytest.raises(RuntimeError) as caught:
await run()
assert caught.value is error
assert calls == ["python"]
@pytest.mark.asyncio
async def test_cancellation_does_not_run_python() -> None:
async def native() -> DispatchResult[str]:
raise asyncio.CancelledError
async def python() -> str:
pytest.fail("cancellation must not dispatch Python")
with pytest.raises(asyncio.CancelledError):
await adispatch(native=native, python=python, route="test", errors=PYTHON_ON_ERROR)
@pytest.mark.parametrize("status", (0, 401, 403, 429, 500, 503))
def test_chat_upstream_mapping_preserves_status_message_and_context(status: int) -> None:
error: Final = Upstream(status, "upstream failed")
with pytest.raises(APIError, match="upstream failed") as caught:
dispatch(
native=lambda: NativeFailed(error),
python=lambda: pytest.fail("upstream errors must not run Python"),
route="chat_completions",
errors=error_handling("anthropic", "model"),
)
assert caught.value.status_code == (status or 500)
assert caught.value.model == "model"
assert caught.value.llm_provider == "anthropic"
assert caught.value.__cause__ is error

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

@ -1,577 +1,117 @@
from __future__ import annotations
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Final
import pytest
from litellm.exceptions import APIError
from litellm.rust_bridge import bindings, runtime
from litellm.rust_bridge import runtime
class RustBridgeDeclined(Exception):
pass
class RustUpstreamError(Exception):
pass
@pytest.fixture(autouse=True)
def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(
bindings,
"get_native_bridge",
lambda: SimpleNamespace(
RustBridgeDeclined=RustBridgeDeclined,
RustUpstreamError=RustUpstreamError,
),
)
def context() -> runtime.BridgeErrorContext:
return runtime.BridgeErrorContext(provider="anthropic", model="model")
def enabled() -> bool:
return True
@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,
error_context=context(),
eligible=case.eligible,
)
assert result == "fallback"
assert tuple(events) == case.expected_events
@pytest.mark.asyncio
@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")
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"
@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,
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.asyncio
async def test_arequire_explains_unavailable_native_binding() -> None:
endpoint: Final = runtime.EndpointBinding(route="messages", load=lambda: None, enabled=enabled)
with pytest.raises(RuntimeError, match=r"^native messages endpoint is unavailable$"):
await endpoint.arequire(
prepare=lambda: pytest.fail("must not prepare"),
call=lambda _binding, _request: pytest.fail("must not invoke"),
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] = []
@pytest.mark.parametrize("state", ("disabled", "ineligible", "unavailable", "handled"))
async def test_attempt_only_prepares_selected_requests(asynchronous: bool, state: str) -> None:
events: Final[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
def test_direct_endpoint_binding_rejects_native_state_controls() -> None:
endpoint: Final = runtime.EndpointBinding(route="test", load=object, enabled=enabled)
with pytest.raises(RuntimeError, match="only native Rust bridges support binding overrides"):
endpoint.override(object())
with pytest.raises(RuntimeError, match="only native Rust bridges support binding resets"):
endpoint.reset()
@pytest.mark.parametrize(
("enabled_state", "reason", "expected"),
(
pytest.param(False, None, runtime.PythonFallbackReason.NATIVE_DISABLED, id="disabled"),
pytest.param(True, "unsupported model", runtime.PythonFallbackReason.NATIVE_DECLINED, id="declined"),
pytest.param(True, None, None, id="accepted"),
),
)
def test_assess_reports_binding_eligibility(
enabled_state: bool,
reason: str | None,
expected: runtime.PythonFallbackReason | None,
) -> None:
binding: Final = object()
checked: list[object] = []
endpoint: Final = runtime.EndpointBinding(route="test", load=lambda: binding, enabled=lambda: enabled_state)
result: Final = endpoint.assess(check=lambda value: checked.append(value) or reason)
assert (result.reason if result is not None else None) is expected
assert (result.detail if result is not None else None) == reason
assert checked == ([binding] if enabled_state else [])
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_dispatch_require_returns_adapted_native_success_without_exception_metadata(
monkeypatch: pytest.MonkeyPatch,
asynchronous: bool,
) -> None:
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
endpoint: Final = runtime.EndpointDispatch(
sync=runtime.EndpointBinding(route="test", load=object, enabled=enabled),
asynchronous=runtime.EndpointBinding(route="test", load=object, enabled=enabled),
)
async def acall(_binding: object, request: int) -> int:
return request * 2
result: Final = (
await endpoint.arequire(prepare=lambda: 3, call=acall, adapt=str, error_context=context())
if asynchronous
else endpoint.require(
prepare=lambda: 3,
call=lambda _binding, request: request * 2,
adapt=str,
error_context=context(),
)
)
assert result == "6"
@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
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.parametrize("available, accepted", ((False, False), (True, False), (True, True)))
async def test_preflight_runs_after_binding_selection_before_preparation(
asynchronous: bool, available: bool, accepted: bool
) -> None:
events: list[str] = []
def load() -> object | None:
events.append("load")
return object() if available else None
def preflight() -> runtime.PythonFallback | None:
events.append("preflight")
return None if accepted else runtime.PythonFallback(runtime.PythonFallbackReason.NATIVE_DECLINED)
def prepare() -> int:
events.append("prepare")
return 7
return 3
def call(binding: object, request: int) -> int:
events.append("native")
return request
def call(_binding: object, request: int) -> int:
events.append("call")
return request * 2
async def acall(binding: object, request: int) -> int:
return call(binding, request)
def fallback() -> str:
events.append("python")
return "3"
def adapt(value: int) -> str:
events.append("adapt")
return str(value)
async def afallback() -> str:
return fallback()
endpoint: Final = runtime.EndpointBinding(route="ocr", load=load, enabled=enabled)
result: Final = (
await endpoint.ainvoke(
prepare=prepare, call=acall, fallback=afallback, adapt=str, error_context=context(), preflight=preflight
await runtime.aattempt(
load=load,
enabled=state != "disabled",
eligible=state != "ineligible",
prepare=prepare,
call=acall,
adapt=adapt,
)
if asynchronous
else endpoint.invoke(
prepare=prepare, call=call, fallback=fallback, adapt=str, error_context=context(), preflight=preflight
else runtime.attempt(
load=load,
enabled=state != "disabled",
eligible=state != "ineligible",
prepare=prepare,
call=call,
adapt=adapt,
)
)
assert result == ("7" if available and accepted else "3")
assert events == (
["load", "preflight", "prepare", "native"]
if available and accepted
else ["load", "preflight", "python"]
if available
else ["load", "python"]
if state == "handled":
assert result == runtime.Handled("6")
assert events == ["load", "prepare", "call", "adapt"]
else:
assert result == runtime.NativeSkipped(runtime.NativeSkipReason(state))
assert events == (["load"] if state == "unavailable" else [])
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.parametrize("phase", ("prepare", "call"))
async def test_attempt_reports_failure_without_deciding_retry(asynchronous: bool, phase: str) -> None:
error: Final = RuntimeError("native failure")
def prepare() -> int:
if phase == "prepare":
raise error
return 3
def call(_binding: object, request: int) -> int:
raise error
async def acall(binding: object, request: int) -> int:
return call(binding, request)
def adapt(value: int) -> str:
pytest.fail("failed attempts cannot be adapted")
result: Final = (
await runtime.aattempt(load=object, enabled=True, eligible=True, prepare=prepare, call=acall, adapt=adapt)
if asynchronous
else runtime.attempt(load=object, enabled=True, eligible=True, prepare=prepare, call=call, adapt=adapt)
)
assert isinstance(result, runtime.NativeFailed)
assert result.error is error
def test_preflight_failure_is_not_a_native_decline() -> None:
endpoint: Final = runtime.EndpointBinding(route="ocr", load=object, enabled=enabled)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_adaptation_failure_remains_distinct_from_native_failure(asynchronous: bool) -> None:
error: Final = ValueError("invalid response")
def preflight() -> runtime.PythonFallback | None:
raise ValueError("invalid acceptance contract")
async def acall(_binding: object, request: int) -> int:
return request
with pytest.raises(ValueError, match="invalid acceptance contract"):
endpoint.invoke(
prepare=lambda: pytest.fail("must not prepare"),
call=lambda binding, request: pytest.fail("must not invoke"),
fallback=lambda: pytest.fail("must not fall back"),
adapt=str,
error_context=context(),
preflight=preflight,
)
def adapt(value: int) -> str:
raise error
async def run() -> None:
if asynchronous:
await runtime.aattempt(load=object, enabled=True, eligible=True, prepare=lambda: 3, call=acall, adapt=adapt)
else:
runtime.attempt(
load=object,
enabled=True,
eligible=True,
prepare=lambda: 3,
call=lambda binding, request: request,
adapt=adapt,
)
with pytest.raises(ValueError, match="invalid response") as caught:
await run()
assert caught.value is error

15
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.runtime import Handled
rust_bridge = importlib.import_module("litellm.rust_bridge.transcription")
@ -55,7 +56,8 @@ def test_enabled_sync_bridge_receives_audio() -> None:
optional_params={"temperature": 0},
timeout=5.0,
)
assert result == {"text": "hello"}
assert isinstance(result, Handled)
assert result.value == {"text": "hello"}
assert bridge.calls[0]["audio"] == {"data": "AQI=", "format": "wav", "filename": "audio.wav"}
@ -72,7 +74,7 @@ async def test_enabled_async_bridge() -> None:
optional_params={},
timeout=None,
)
assert result == {"text": "async"}
assert result == Handled({"text": "async"})
def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None:
@ -83,7 +85,8 @@ def test_loader_returns_none_without_native_extension(monkeypatch: pytest.Monkey
def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None)
rust_bridge.configure_rust_transcription(transcription=None)
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None)
with pytest.raises(RuntimeError, match="bridge is unavailable"):
BedrockAudioTranscriptionRustDispatch().audio_transcriptions(
@ -100,10 +103,8 @@ 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.configure_rust_transcription(atranscription=None)
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None)
with pytest.raises(RuntimeError, match="bridge is unavailable"):
await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions(

6
type-discipline-budget.json
View file

@ -1,6 +1,6 @@
{
"LIT001": {
"limit": 22173
"limit": 22165
},
"LIT002": {
"limit": 26729
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1027
"limit": 1022
},
"LIT007": {
"limit": 0
@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16426
"limit": 16419
},
"LIT011": {
"limit": 5506