diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index c82be07a5c5..bfe2e86eb7c 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -465,6 +465,7 @@ class AnthropicChatCompletion(BaseLLM): timeout=timeout, on_response=log_rust_post_call, python_fallback=python_fallback, + request_override=rust_chat_completions_bridge.rust_request_override(litellm_params), ) rust_response: Final = rust_chat_completions_bridge.chat_completions( model=model, @@ -477,6 +478,7 @@ class AnthropicChatCompletion(BaseLLM): extra_headers=headers, timeout=timeout, on_response=log_rust_post_call, + request_override=rust_chat_completions_bridge.rust_request_override(litellm_params), ) if rust_response is not None: return rust_response diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index b1f8c957ff4..2e0a228ccee 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -53,8 +53,6 @@ class BedrockAudioTranscriptionRustDispatch: optional_params=optional_params, timeout=timeout, ) - if rust_response is None: - raise RuntimeError("Rust audio transcription bridge is unavailable") return TranscriptionResponse(**rust_response) async def async_audio_transcriptions( @@ -79,6 +77,4 @@ class BedrockAudioTranscriptionRustDispatch: optional_params=optional_params, timeout=timeout, ) - if rust_response is None: - raise RuntimeError("Rust audio transcription bridge is unavailable") return TranscriptionResponse(**rust_response) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index a75124325ae..12c74be3e73 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -452,6 +452,7 @@ class BedrockConverseLLM(BaseAWSLLM): api_key=api_key, skip_pre_call_logging=True, ), + request_override=rust_chat_completions_bridge.rust_request_override(litellm_params), ) rust_response: Final = rust_chat_completions_bridge.chat_completions( model=model, @@ -464,6 +465,7 @@ class BedrockConverseLLM(BaseAWSLLM): extra_headers=headers, timeout=timeout, on_response=log_rust_post_call, + request_override=rust_chat_completions_bridge.rust_request_override(litellm_params), ) if rust_response is not None: return rust_response diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f281c249c72..c87d5028bf1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2401,34 +2401,24 @@ class BaseLLMHTTPHandler: ) -> AnthropicMessagesResponse | None: if custom_llm_provider not in ("azure_ai", "anthropic"): return None - from litellm.rust_bridge.configuration import rust_enabled - raw_request_override: Final = litellm_params.get("rust") request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - if not rust_enabled(request_override=request_override): - return None if has_agentic_hook: return None from litellm.rust_bridge import messages as rust_messages_bridge upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"} - try: - rust_response: Final = await rust_messages_bridge.amessages( - model=model, - body=upstream_body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=headers, - timeout=timeout, - ) - except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path - verbose_logger.debug( - "Rust Anthropic messages bridge raised %s; falling back to Python path", - type(rust_error).__name__, - ) - return None + rust_response: Final = await rust_messages_bridge.amessages( + model=model, + body=upstream_body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + request_override=request_override, + ) if rust_response is None: return None @@ -6517,10 +6507,12 @@ class BaseLLMHTTPHandler: if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params): from litellm.rust_bridge import responses_websocket as rust_responses_websocket + raw_rust_override: Final = litellm_params.get("rust") rust_backend: Final = await rust_responses_websocket.connect( url=ws_url, headers={str(key): str(value) for key, value in headers.items()}, timeout=timeout, + request_override=raw_rust_override if isinstance(raw_rust_override, bool) else None, ) if rust_backend is not None: yield rust_backend diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b260ec6e06f..eb74e5f2c69 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -52,14 +52,6 @@ class _PreparedOCRRequest: litellm_logging_obj: LiteLLMLoggingObj -@dataclass -class _PreparedRustOCRCall: - api_key: str | None - api_base: str | None - headers: dict[str, object] - optional_params: dict[str, object] - - _RUST_OCR_PROVIDERS: Final = { "mistral", "azure_ai", @@ -194,10 +186,9 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS -def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool: +def _rust_request_override(prepared_request: _PreparedOCRRequest) -> bool | None: raw_request_override: Final = prepared_request.litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override) + return raw_request_override if isinstance(raw_request_override, bool) else None def _rust_bridge_optional_params( @@ -242,7 +233,7 @@ def _rust_bridge_api_base( def _prepare_rust_ocr_call( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> _PreparedRustOCRCall: +) -> rust_ocr_bridge.RustOCRRequest: provider_config: Final = prepared_request.provider_config api_key_env_var: Final = provider_config.get_api_key_env_var() resolved_api_key: Final = prepared_request.api_key or ( @@ -276,11 +267,15 @@ def _prepare_rust_ocr_call( "headers": resolved_headers, }, ) - return _PreparedRustOCRCall( + return rust_ocr_bridge.RustOCRRequest( + model=prepared_request.model, + document=prepared_request.document, api_key=resolved_api_key, api_base=rust_api_base, - headers=cast(dict[str, object], resolved_headers), + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=cast(dict[str, object], resolved_headers), optional_params=rust_optional_params, + timeout=prepared_request.effective_timeout, ) @@ -288,21 +283,15 @@ def _run_rust_ocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], ) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_ocr() is None: - return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) rust_response: Final = rust_ocr_bridge.ocr( + prepare=lambda: _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ), model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, + provider=prepared_request.custom_llm_provider, + request_override=_rust_request_override(prepared_request), + eligible=_rust_ocr_supported(prepared_request), ) if rust_response is None: return None @@ -313,21 +302,15 @@ async def _run_rust_aocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], ) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_aocr() is None: - return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) rust_response: Final = await rust_ocr_bridge.aocr( + prepare=lambda: _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ), model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, + provider=prepared_request.custom_llm_provider, + request_override=_rust_request_override(prepared_request), + eligible=_rust_ocr_supported(prepared_request), ) if rust_response is None: return None @@ -428,17 +411,14 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): - from litellm.secret_managers.main import get_secret_str + from litellm.secret_managers.main import get_secret_str - rust_response: Final = await _run_rust_aocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - if rust_response is None: - verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response + rust_response: Final = await _run_rust_aocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + if rust_response is not None: + return rust_response response = base_llm_http_handler.ocr( model=prepared.model, @@ -700,17 +680,14 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): - from litellm.secret_managers.main import get_secret_str + from litellm.secret_managers.main import get_secret_str - rust_response: Final = _run_rust_ocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - if rust_response is None: - verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response + rust_response: Final = _run_rust_ocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + if rust_response is not None: + return rust_response response: Final = base_llm_http_handler.ocr( model=prepared.model, diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py index d16f150a2aa..86d2ef624e2 100644 --- a/litellm/rust_bridge/bindings.py +++ b/litellm/rust_bridge/bindings.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable -from typing import Final, Generic, TypeVar +from typing import Final, Generic, TypeVar, cast from litellm.rust_bridge.loader import get_native_bridge @@ -15,6 +15,13 @@ class _Unset: _UNSET: Final = _Unset() +class Unchanged: + pass + + +UNCHANGED: Final = Unchanged() + + class NativeBinding(Generic[BindingT]): """Resolve one native attribute with an explicit, resettable test override.""" @@ -23,6 +30,13 @@ class NativeBinding(Generic[BindingT]): self._validate: Final = validate self._override: BindingT | None | _Unset = _UNSET + @classmethod + def callable(cls, attribute: str) -> NativeBinding[BindingT]: + return cls( + attribute, + validate=lambda value: cast(BindingT, value) if callable(value) else None, + ) + def load(self) -> BindingT | None: if not isinstance(self._override, _Unset): return self._override diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index c599667ab17..3a78ef30aaa 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -14,20 +14,24 @@ from __future__ import annotations import json from collections.abc import Awaitable, Callable, Mapping, Sequence -from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol import httpx from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger -from litellm.exceptions import APIError from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( convert_to_model_response_object, ) from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged from litellm.rust_bridge.configuration import rust_enabled -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + RustBridge, + RustEndpoint, + async_none, +) from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.utils import ModelResponse @@ -126,67 +130,42 @@ def response_logger( return log -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -@dataclass(slots=True) -class _RustChatCompletionsState: - chat_completions: RustChatCompletions | None = None - achat_completions: RustAchatCompletions | None = None - decline: RustChatCompletionsDecline | None = None - - -_STATE: Final[_RustChatCompletionsState] = _RustChatCompletionsState() +_CHAT: Final[RustEndpoint[RustChatCompletions, RustAchatCompletions]] = RustEndpoint.native( + route="chat completions", + sync="chat_completions", + asynchronous="achat_completions", + enabled=rust_enabled, +) +_CHAT_PREFLIGHT: Final[RustBridge[RustChatCompletionsDecline]] = RustBridge.native( + route="chat completions preflight", + attribute="chat_completions_decline", + enabled=rust_enabled, +) def set_rust_chat_completions( *, - chat_completions: RustChatCompletions | None | _Unset = _UNSET, - achat_completions: RustAchatCompletions | None | _Unset = _UNSET, - decline: RustChatCompletionsDecline | None | _Unset = _UNSET, + chat_completions: RustChatCompletions | None | Unchanged = UNCHANGED, + achat_completions: RustAchatCompletions | None | Unchanged = UNCHANGED, + decline: RustChatCompletionsDecline | None | Unchanged = UNCHANGED, ) -> None: """Inject the native callables, so tests can supply a double instead of patching module attributes.""" - if not isinstance(chat_completions, _Unset): - _STATE.chat_completions = chat_completions - if not isinstance(achat_completions, _Unset): - _STATE.achat_completions = achat_completions - if not isinstance(decline, _Unset): - _STATE.decline = decline - - -def load_rust_chat_completions() -> RustChatCompletions | None: - if _STATE.chat_completions is not None: - return _STATE.chat_completions - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - loaded: RustChatCompletions | None = getattr(native_bridge, "chat_completions", None) - return loaded - - -def load_rust_achat_completions() -> RustAchatCompletions | None: - if _STATE.achat_completions is not None: - return _STATE.achat_completions - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - loaded: RustAchatCompletions | None = getattr(native_bridge, "achat_completions", None) - return loaded - - -def _load_rust_decline() -> RustChatCompletionsDecline | None: - if _STATE.decline is not None: - return _STATE.decline - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - loaded: RustChatCompletionsDecline | None = getattr(native_bridge, "chat_completions_decline", None) - return loaded + if not isinstance(chat_completions, Unchanged): + if chat_completions is None: + _CHAT.sync.reset() + else: + _CHAT.sync.override(chat_completions) + if not isinstance(achat_completions, Unchanged): + if achat_completions is None: + _CHAT.asynchronous.reset() + else: + _CHAT.asynchronous.override(achat_completions) + if not isinstance(decline, Unchanged): + if decline is None: + _CHAT_PREFLIGHT.reset() + else: + _CHAT_PREFLIGHT.override(decline) def _anthropic_user_id_reaches_the_body(litellm_params: Mapping[str, object] | None) -> bool: @@ -227,6 +206,11 @@ def _litellm_metadata_reaches_the_provider( return False +def rust_request_override(litellm_params: Mapping[str, object] | None) -> bool | None: + raw_request_override: Final = litellm_params.get("rust") if litellm_params is not None else None + return raw_request_override if isinstance(raw_request_override, bool) else None + + def rust_chat_completions_accepts( *, model: str, @@ -247,82 +231,17 @@ def rust_chat_completions_accepts( return False if stream: return False - request_override: Final = litellm_params.get("rust") if litellm_params is not None else None - if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None): - return False if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") return False - decline: Final = _load_rust_decline() - if decline is None: - return False - try: - reason: Final = decline( + return _CHAT_PREFLIGHT.accepts( + request_override=rust_request_override(litellm_params), + check=lambda decline: decline( model=model, messages=messages, optional_params=optional_params, custom_llm_provider=custom_llm_provider, - ) - except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path - verbose_logger.debug( - "Rust chat completions gate raised %s; staying on the Python path", - type(rust_error).__name__, - ) - return False - if reason is not None: - verbose_logger.debug("Rust chat completions declined (%s); using the Python path", reason) - return False - return True - - -def _rust_bridge_exceptions() -> tuple[type[BaseException], type[BaseException]] | None: - """`(declined, upstream_failed)` from the native module, or None when absent.""" - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - declined: Final = getattr(native_bridge, "RustBridgeDeclined", None) - upstream: Final = getattr(native_bridge, "RustUpstreamError", None) - if declined is None or upstream is None: - return None - return declined, upstream - - -def _reraise_or_decline( - rust_error: BaseException, - *, - model: str, - custom_llm_provider: str | None, -) -> None: - """Re-raise a failure the provider already saw, or return so the caller declines. - - A request that never reached the provider is safe to serve on the Python - path. One that did is not: the provider has already done the work, so a - second attempt bills for it twice. Those surface as an `APIError` carrying - the upstream status, which LiteLLM's exception mapping already understands. - """ - exceptions: Final = _rust_bridge_exceptions() - if exceptions is None: - verbose_logger.debug( - "Rust chat completions bridge raised %s; falling back to Python path", - type(rust_error).__name__, - ) - return - declined, upstream_failed = exceptions - if isinstance(rust_error, upstream_failed): - args: Final = rust_error.args - status: Final = args[0] if args else 0 - message: Final = args[1] if len(args) > 1 else "" - raise APIError( - status_code=int(status) or 500, - message=f"litellm rust chat completions: {message}", - llm_provider=custom_llm_provider or "", - model=model, - ) - if not isinstance(rust_error, declined): - raise rust_error - verbose_logger.debug( - "Rust chat completions declined before calling the provider (%s); using the Python path", - rust_error, + ), ) @@ -352,12 +271,14 @@ def chat_completions( extra_headers: Mapping[str, object] | None, timeout: float | httpx.Timeout | None, on_response: ResponseObserver, + request_override: bool | None, ) -> ModelResponse | None: - rust_chat_completions: Final = load_rust_chat_completions() - if rust_chat_completions is None: - return None - try: - rust_response: Final = rust_chat_completions( + def adapt(rust_response: Mapping[str, object]) -> ModelResponse: + on_response(rust_response) + return _build_model_response(rust_response, model_response) + + return _CHAT.invoke( + call=lambda rust_chat_completions: rust_chat_completions( model=model, messages=messages, optional_params=optional_params, @@ -366,12 +287,12 @@ def chat_completions( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, timeout_seconds=timeout_to_seconds(timeout), - ) - except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw - _reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider) - return None - on_response(rust_response) - return _build_model_response(rust_response, model_response) + ), + fallback=lambda: None, + adapt=adapt, + context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + request_override=request_override, + ) async def achat_completions( @@ -386,12 +307,14 @@ async def achat_completions( extra_headers: Mapping[str, object] | None, timeout: float | httpx.Timeout | None, on_response: ResponseObserver, + request_override: bool | None, ) -> ModelResponse | None: - rust_achat_completions: Final = load_rust_achat_completions() - if rust_achat_completions is None: - return None - try: - rust_response: Final = await rust_achat_completions( + def adapt(rust_response: Mapping[str, object]) -> ModelResponse: + on_response(rust_response) + return _build_model_response(rust_response, model_response) + + return await _CHAT.ainvoke( + call=lambda rust_achat_completions: rust_achat_completions( model=model, messages=messages, optional_params=optional_params, @@ -400,12 +323,12 @@ async def achat_completions( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, timeout_seconds=timeout_to_seconds(timeout), - ) - except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw - _reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider) - return None - on_response(rust_response) - return _build_model_response(rust_response, model_response) + ), + fallback=async_none, + adapt=adapt, + context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + request_override=request_override, + ) async def achat_completions_or_fallback( @@ -421,6 +344,7 @@ async def achat_completions_or_fallback( timeout: float | httpx.Timeout | None, on_response: ResponseObserver, python_fallback: Callable[[], Awaitable[object]], + request_override: bool | None, ) -> object: """Await the Rust path, falling back to the caller's own Python path when the bridge is unavailable or the call fails. @@ -430,18 +354,24 @@ async def achat_completions_or_fallback( already returned a coroutine by the time a Rust failure surfaces, and so cannot fall back on its own. """ - response: Final = await achat_completions( - model=model, - messages=messages, - optional_params=optional_params, - model_response=model_response, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout=timeout, - on_response=on_response, + + def adapt(rust_response: Mapping[str, object]) -> object: + on_response(rust_response) + return _build_model_response(rust_response, model_response) + + return await _CHAT.ainvoke( + call=lambda rust_achat_completions: rust_achat_completions( + model=model, + messages=messages, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ), + fallback=python_fallback, + adapt=adapt, + context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + request_override=request_override, ) - if response is not None: - return response - return await python_fallback() diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 40d0ddf622b..311462bb3a2 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -3,11 +3,18 @@ from __future__ import annotations from collections.abc import Awaitable -from dataclasses import dataclass -from typing import Final, Protocol, cast +from typing import Final, Protocol import httpx +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.configuration import rust_enabled +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + RustEndpoint, + async_none, + identity, +) from litellm.rust_bridge.timeouts import timeout_to_seconds @@ -39,53 +46,29 @@ class RustAmessages(Protocol): raise NotImplementedError -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -@dataclass(slots=True) -class _RustMessagesState: - messages: RustMessages | None = None - amessages: RustAmessages | None = None - - -_STATE: Final[_RustMessagesState] = _RustMessagesState() +_MESSAGES: Final[RustEndpoint[RustMessages, RustAmessages]] = RustEndpoint.native( + route="messages", + sync="messages", + asynchronous="amessages", + enabled=rust_enabled, +) def set_rust_messages( *, - messages: RustMessages | None | _Unset = _UNSET, - amessages: RustAmessages | None | _Unset = _UNSET, + messages: RustMessages | None | Unchanged = UNCHANGED, + amessages: RustAmessages | None | Unchanged = UNCHANGED, ) -> None: - if not isinstance(messages, _Unset): - _STATE.messages = messages - if not isinstance(amessages, _Unset): - _STATE.amessages = amessages - - -def load_rust_messages() -> RustMessages | None: - if _STATE.messages is not None: - return _STATE.messages - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustMessages, getattr(native_bridge, "messages", None)) - - -def load_rust_amessages() -> RustAmessages | None: - if _STATE.amessages is not None: - return _STATE.amessages - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustAmessages, getattr(native_bridge, "amessages", None)) + if not isinstance(messages, Unchanged): + if messages is None: + _MESSAGES.sync.reset() + else: + _MESSAGES.sync.override(messages) + if not isinstance(amessages, Unchanged): + if amessages is None: + _MESSAGES.asynchronous.reset() + else: + _MESSAGES.asynchronous.override(amessages) def messages( @@ -97,18 +80,22 @@ def messages( custom_llm_provider: str | None, extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, + request_override: bool | None = None, ) -> dict[str, object] | None: - rust_messages: Final = load_rust_messages() - if rust_messages is None: - return None - return rust_messages( - model=model, - body=body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), + return _MESSAGES.invoke( + call=lambda rust_messages: rust_messages( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ), + fallback=lambda: None, + adapt=identity, + context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + request_override=request_override, ) @@ -121,16 +108,20 @@ async def amessages( custom_llm_provider: str | None, extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, + request_override: bool | None = None, ) -> dict[str, object] | None: - rust_amessages: Final = load_rust_amessages() - if rust_amessages is None: - return None - return await rust_amessages( - model=model, - body=body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), + return await _MESSAGES.ainvoke( + call=lambda rust_amessages: rust_amessages( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ), + fallback=async_none, + adapt=identity, + context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + request_override=request_override, ) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 86038438f57..a530b99cdd7 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,18 +2,38 @@ from __future__ import annotations -from collections.abc import Awaitable -from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from typing import Final, Protocol import httpx from litellm.rust_bridge import configuration as _configuration +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + RustEndpoint, + async_none, + identity, +) from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds rust_ocr_enabled = _configuration.rust_ocr_enabled rust = _configuration.rust +@dataclass(frozen=True, slots=True) +class RustOCRRequest: + model: str + document: dict[str, object] + api_key: str | None + api_base: str | None + custom_llm_provider: str | None + extra_headers: dict[str, object] | None + optional_params: dict[str, object] + timeout: float | httpx.Timeout | None + + class RustOcr(Protocol): def __call__( self, @@ -44,98 +64,88 @@ class RustAocr(Protocol): raise NotImplementedError -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -_rust_ocr_impl: RustOcr | None = None -_rust_aocr_impl: RustAocr | None = None +_OCR: Final[RustEndpoint[RustOcr, RustAocr]] = RustEndpoint.native( + route="ocr", + sync="ocr", + asynchronous="aocr", + enabled=_configuration.rust_ocr_enabled, +) def set_rust_ocr( *, - ocr: RustOcr | None | _Unset = _UNSET, - aocr: RustAocr | None | _Unset = _UNSET, + ocr: RustOcr | None | Unchanged = UNCHANGED, + aocr: RustAocr | None | Unchanged = UNCHANGED, ) -> None: - global _rust_ocr_impl, _rust_aocr_impl - if not isinstance(ocr, _Unset): - _rust_ocr_impl = ocr - if not isinstance(aocr, _Unset): - _rust_aocr_impl = aocr - - -def load_rust_ocr() -> RustOcr | None: - if _rust_ocr_impl is not None: - return _rust_ocr_impl - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustOcr, native_bridge.ocr) - - -def load_rust_aocr() -> RustAocr | None: - if _rust_aocr_impl is not None: - return _rust_aocr_impl - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustAocr, getattr(native_bridge, "aocr", None)) + if not isinstance(ocr, Unchanged): + if ocr is None: + _OCR.sync.reset() + else: + _OCR.sync.override(ocr) + if not isinstance(aocr, Unchanged): + if aocr is None: + _OCR.asynchronous.reset() + else: + _OCR.asynchronous.override(aocr) def ocr( *, + prepare: Callable[[], RustOCRRequest], model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: - rust_ocr: Final = load_rust_ocr() - if rust_ocr is None: - return None - return rust_ocr( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=_timeout_to_seconds(timeout), + provider: str, + request_override: bool | None, + eligible: bool, +) -> Mapping[str, object] | None: + return _OCR.invoke( + call=lambda rust_ocr: _call_ocr(rust_ocr, prepare()), + fallback=lambda: None, + adapt=identity, + context=BridgeErrorContext(provider=provider, model=model), + request_override=request_override, + eligible=eligible, ) async def aocr( *, + prepare: Callable[[], RustOCRRequest], model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: - rust_aocr: Final = load_rust_aocr() - if rust_aocr is None: - return None - return await rust_aocr( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=_timeout_to_seconds(timeout), + provider: str, + request_override: bool | None, + eligible: bool, +) -> Mapping[str, object] | None: + return await _OCR.ainvoke( + call=lambda rust_aocr: _call_aocr(rust_aocr, prepare()), + fallback=async_none, + adapt=identity, + context=BridgeErrorContext(provider=provider, model=model), + request_override=request_override, + eligible=eligible, + ) + + +def _call_ocr(rust_ocr: RustOcr, request: RustOCRRequest) -> Mapping[str, object]: + return rust_ocr( + model=request.model, + document=request.document, + api_key=request.api_key, + api_base=request.api_base, + custom_llm_provider=request.custom_llm_provider, + extra_headers=request.extra_headers, + optional_params=request.optional_params, + timeout_seconds=_timeout_to_seconds(request.timeout), + ) + + +def _call_aocr(rust_aocr: RustAocr, request: RustOCRRequest) -> Awaitable[Mapping[str, object]]: + return rust_aocr( + model=request.model, + document=request.document, + api_key=request.api_key, + api_base=request.api_base, + custom_llm_provider=request.custom_llm_provider, + extra_headers=request.extra_headers, + optional_params=request.optional_params, + timeout_seconds=_timeout_to_seconds(request.timeout), ) diff --git a/litellm/rust_bridge/responses.py b/litellm/rust_bridge/responses.py new file mode 100644 index 00000000000..fdd1aa92c48 --- /dev/null +++ b/litellm/rust_bridge/responses.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from typing import Final, Protocol + +from litellm.rust_bridge.configuration import rust_enabled +from litellm.rust_bridge.runtime import BridgeErrorContext, RustEndpoint, async_none, identity + + +class RustResponses(Protocol): + def __call__(self, request: Mapping[str, object]) -> Mapping[str, object]: ... + + +class RustAresponses(Protocol): + def __call__(self, request: Mapping[str, object]) -> Awaitable[Mapping[str, object]]: ... + + +_RESPONSES: Final[RustEndpoint[RustResponses, RustAresponses]] = RustEndpoint.native( + route="responses", + sync="responses", + asynchronous="aresponses", + enabled=rust_enabled, +) + + +def responses( + *, + prepare: Callable[[], Mapping[str, object]], + model: str, + provider: str, + request_override: bool | None = None, +) -> Mapping[str, object] | None: + return _RESPONSES.invoke( + call=lambda rust_responses: rust_responses(prepare()), + fallback=lambda: None, + adapt=identity, + context=BridgeErrorContext(provider=provider, model=model), + request_override=request_override, + ) + + +async def aresponses( + *, + prepare: Callable[[], Mapping[str, object]], + model: str, + provider: str, + request_override: bool | None = None, +) -> Mapping[str, object] | None: + return await _RESPONSES.ainvoke( + call=lambda rust_aresponses: rust_aresponses(prepare()), + fallback=async_none, + adapt=identity, + context=BridgeErrorContext(provider=provider, model=model), + request_override=request_override, + ) diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 0634867af1c..10a516af19c 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -2,13 +2,18 @@ from __future__ import annotations -from dataclasses import dataclass from typing import Final, Protocol import httpx from websockets.exceptions import ConnectionClosedOK -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.configuration import rust_enabled +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + RustBridge, + async_none, +) from litellm.rust_bridge.timeouts import timeout_to_seconds @@ -30,39 +35,22 @@ class RustResponsesWebSocketConnection(Protocol): ) -> RustResponsesWebSocket: ... -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -@dataclass(slots=True) -class _RustResponsesWebSocketState: - connection: RustResponsesWebSocketConnection | None = None - - -_STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() +_RESPONSES_WEBSOCKET: Final[RustBridge[RustResponsesWebSocketConnection]] = RustBridge.native( + route="responses websocket", + attribute="ResponsesWebSocketConnection", + enabled=rust_enabled, +) def set_rust_responses_websocket( *, - connection: RustResponsesWebSocketConnection | None | _Unset = _UNSET, + connection: RustResponsesWebSocketConnection | None | Unchanged = UNCHANGED, ) -> None: - if not isinstance(connection, _Unset): - _STATE.connection = connection - - -def load_rust_responses_websocket() -> RustResponsesWebSocketConnection | None: - if _STATE.connection is not None: - return _STATE.connection - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - connection_type: Final[RustResponsesWebSocketConnection | None] = getattr( - native_bridge, "ResponsesWebSocketConnection", None - ) - return connection_type + if not isinstance(connection, Unchanged): + if connection is None: + _RESPONSES_WEBSOCKET.reset() + else: + _RESPONSES_WEBSOCKET.override(connection) class _ConnectionAdapter: @@ -87,16 +75,16 @@ async def connect( url: str, headers: dict[str, str], timeout: float | httpx.Timeout | None, + request_override: bool | None = None, ) -> _ConnectionAdapter | None: - connection_type: Final = load_rust_responses_websocket() - if connection_type is None: - return None - try: - connection: Final = await connection_type.connect( + return await _RESPONSES_WEBSOCKET.ainvoke( + call=lambda connection_type: connection_type.connect( url=url, headers=headers, timeout_seconds=timeout_to_seconds(timeout), - ) - except Exception: # noqa: BLE001 # bridge failures must fall back to Python - return None - return _ConnectionAdapter(connection) + ), + fallback=async_none, + adapt=_ConnectionAdapter, + context=BridgeErrorContext(provider="openai", model="responses websocket"), + request_override=request_override, + ) diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index 00f06c046a2..0aab3f6b0ec 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -1,20 +1,29 @@ from __future__ import annotations from collections.abc import Awaitable, Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum -from typing import Final, Generic, NoReturn, TypeAlias, TypeVar +from typing import Final, Generic, NoReturn, Protocol, TypeAlias, TypeVar from litellm.exceptions import APIError -from litellm.rust_bridge.bindings import native_exception_types +from litellm.rust_bridge.bindings import ( + UNCHANGED, + NativeBinding, + Unchanged, + native_exception_types, +) +BindingT = TypeVar("BindingT") NativeT = TypeVar("NativeT") ResultT = TypeVar("ResultT") +SyncBindingT = TypeVar("SyncBindingT") +AsyncBindingT = TypeVar("AsyncBindingT") -class FallbackMode(Enum): - PYTHON = "python" - RUST_REQUIRED = "rust_required" +class PythonFallbackReason(Enum): + RUST_DISABLED = "rust_disabled" + RUST_UNAVAILABLE = "rust_unavailable" + RUST_DECLINED = "rust_declined" @dataclass(frozen=True, slots=True) @@ -23,153 +32,410 @@ class RustHandled(Generic[ResultT]): @dataclass(frozen=True, slots=True) -class RustDeclined: - reason: str +class PythonFallback: + reason: PythonFallbackReason + detail: str | None = None -@dataclass(frozen=True, slots=True) -class RustUnavailable: - pass - - -RustAttempt: TypeAlias = RustHandled[ResultT] | RustDeclined | RustUnavailable +RustDispatchResult: TypeAlias = RustHandled[ResultT] | PythonFallback @dataclass(frozen=True, slots=True) class BridgeErrorContext: - route: str provider: str model: str -def invoke( - *, - native_call: Callable[[], NativeT] | None, - fallback: Callable[[], ResultT], - adapt: Callable[[NativeT], ResultT], - mode: FallbackMode, - context: BridgeErrorContext, -) -> ResultT: - result: Final = attempt(native_call=native_call, adapt=adapt, context=context) - if isinstance(result, RustHandled): - return result.value - if mode is FallbackMode.PYTHON: - return fallback() - _raise_required(result, context) +class RustEnablement(Protocol): + def __call__(self, *, request_override: bool | None = None) -> bool: ... -async def ainvoke( - *, - native_call: Callable[[], Awaitable[NativeT]] | None, - fallback: Callable[[], Awaitable[ResultT]], - adapt: Callable[[NativeT], ResultT], - mode: FallbackMode, - context: BridgeErrorContext, -) -> ResultT: - result: Final = await aattempt(native_call=native_call, adapt=adapt, context=context) - if isinstance(result, RustHandled): - return result.value - if mode is FallbackMode.PYTHON: - return await fallback() - _raise_required(result, context) +@dataclass(frozen=True, slots=True) +class RustBridge(Generic[BindingT]): + route: str + load: Callable[[], BindingT | None] + enabled: RustEnablement + _native_binding: NativeBinding[BindingT] | None = field(default=None, repr=False) + + @classmethod + def native( + cls, + *, + route: str, + attribute: str, + enabled: RustEnablement, + ) -> RustBridge[BindingT]: + binding: NativeBinding[BindingT] = NativeBinding.callable(attribute) + return cls( + route=route, + load=binding.load, + enabled=enabled, + _native_binding=binding, + ) + + def override(self, value: BindingT | None) -> None: + if self._native_binding is None: + raise RuntimeError("only native Rust bridges support binding overrides") + self._native_binding.override(value) + + def reset(self) -> None: + if self._native_binding is None: + raise RuntimeError("only native Rust bridges support binding resets") + self._native_binding.reset() + + def _attempt( + self, + *, + call: Callable[[BindingT], NativeT], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> RustDispatchResult[ResultT]: + binding_or_fallback: Final = self._binding_or_python_fallback( + request_override=request_override, + eligible=eligible, + ) + if isinstance(binding_or_fallback, PythonFallback): + return binding_or_fallback + return self._attempt_call( + call=lambda: call(binding_or_fallback), + adapt=adapt, + context=context, + ) + + async def _aattempt( + self, + *, + call: Callable[[BindingT], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> RustDispatchResult[ResultT]: + binding_or_fallback: Final = self._binding_or_python_fallback( + request_override=request_override, + eligible=eligible, + ) + if isinstance(binding_or_fallback, PythonFallback): + return binding_or_fallback + return await self._attempt_acall( + call=lambda: call(binding_or_fallback), + adapt=adapt, + context=context, + ) + + def invoke( + self, + *, + call: Callable[[BindingT], NativeT], + fallback: Callable[[], ResultT], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> ResultT: + result: Final = self._attempt( + call=call, + adapt=adapt, + context=context, + request_override=request_override, + eligible=eligible, + ) + match result: + case RustHandled(value=value): + return value + case PythonFallback(): + return fallback() + + async def ainvoke( + self, + *, + call: Callable[[BindingT], Awaitable[NativeT]], + fallback: Callable[[], Awaitable[ResultT]], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> ResultT: + result: Final = await self._aattempt( + call=call, + adapt=adapt, + context=context, + request_override=request_override, + eligible=eligible, + ) + match result: + case RustHandled(value=value): + return value + case PythonFallback(): + return await fallback() + + def require( + self, + *, + call: Callable[[BindingT], NativeT], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> ResultT: + result: Final = self._attempt( + call=call, + adapt=adapt, + context=context, + request_override=request_override, + eligible=eligible, + ) + match result: + case RustHandled(value=value): + return value + case PythonFallback(): + self._raise_required(result) + + async def arequire( + self, + *, + call: Callable[[BindingT], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> ResultT: + result: Final = await self._aattempt( + call=call, + adapt=adapt, + context=context, + request_override=request_override, + eligible=eligible, + ) + match result: + case RustHandled(value=value): + return value + case PythonFallback(): + self._raise_required(result) + + def accepts( + self, + *, + check: Callable[[BindingT], str | None], + request_override: bool | None = None, + eligible: bool = True, + ) -> bool: + binding_or_fallback: Final = self._binding_or_python_fallback( + request_override=request_override, + eligible=eligible, + ) + if isinstance(binding_or_fallback, PythonFallback): + return False + try: + reason: Final = check(binding_or_fallback) + except Exception: # noqa: BLE001 # preflight performs no provider I/O, so Python handoff is safe + return False + return reason is None + + def _binding_or_python_fallback( + self, + *, + request_override: bool | None, + eligible: bool, + ) -> BindingT | PythonFallback: + if not eligible or not self.enabled(request_override=request_override): + return PythonFallback(PythonFallbackReason.RUST_DISABLED) + binding: Final = self.load() + if binding is None: + return PythonFallback(PythonFallbackReason.RUST_UNAVAILABLE) + return binding + + def _attempt_call( + self, + *, + call: Callable[[], NativeT], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + ) -> RustDispatchResult[ResultT]: + exceptions: Final = native_exception_types() + if exceptions is None: + return RustHandled(adapt(call())) + declined, upstream = exceptions + try: + value: Final = call() + except declined as error: + return PythonFallback(PythonFallbackReason.RUST_DECLINED, _error_message(error)) + except upstream as error: + self._raise_upstream(error, context) + return RustHandled(adapt(value)) + + async def _attempt_acall( + self, + *, + call: Callable[[], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + ) -> RustDispatchResult[ResultT]: + exceptions: Final = native_exception_types() + if exceptions is None: + return RustHandled(adapt(await call())) + declined, upstream = exceptions + try: + value: Final = await call() + except declined as error: + return PythonFallback(PythonFallbackReason.RUST_DECLINED, _error_message(error)) + except upstream as error: + self._raise_upstream(error, context) + return RustHandled(adapt(value)) + + def _raise_required(self, fallback: PythonFallback) -> NoReturn: + detail: Final = f": {fallback.detail}" if fallback.detail else "" + reason: Final = _required_reason(fallback.reason) + raise RuntimeError(f"Rust {self.route} bridge {reason}{detail}") + + def _raise_upstream(self, error: BaseException, context: BridgeErrorContext) -> NoReturn: + args: Final[tuple[object, ...]] = error.args + attribute_status: Final = getattr(error, "status_code", None) + attribute_message: Final = getattr(error, "message", None) + status_value: Final = attribute_status if isinstance(attribute_status, int) else (args[0] if args else 0) + message_value: Final = ( + attribute_message if isinstance(attribute_message, str) else (args[1] if len(args) > 1 else str(error)) + ) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + raise APIError( + status_code=status or 500, + message=f"litellm rust {self.route}: {message}", + llm_provider=context.provider, + model=context.model, + ) from error -def attempt( - *, - native_call: Callable[[], NativeT] | None, - adapt: Callable[[NativeT], ResultT], - context: BridgeErrorContext, -) -> RustAttempt[ResultT]: - if native_call is None: - return RustUnavailable() - exceptions: Final = native_exception_types() - if exceptions is None: - return RustHandled(adapt(native_call())) - declined, upstream = exceptions - try: - value: Final = native_call() - except declined as error: - return RustDeclined(reason=_decline_reason(error)) - except upstream as error: - _raise_upstream(error, context) - return RustHandled(adapt(value)) +@dataclass(frozen=True, slots=True) +class RustEndpoint(Generic[SyncBindingT, AsyncBindingT]): + sync: RustBridge[SyncBindingT] + asynchronous: RustBridge[AsyncBindingT] + + @classmethod + def native( + cls, + *, + route: str, + sync: str, + asynchronous: str, + enabled: RustEnablement, + ) -> RustEndpoint[SyncBindingT, AsyncBindingT]: + return cls( + sync=RustBridge.native(route=route, attribute=sync, enabled=enabled), + asynchronous=RustBridge.native( + route=route, + attribute=asynchronous, + enabled=enabled, + ), + ) + + def override( + self, + *, + sync: SyncBindingT | None | Unchanged = UNCHANGED, + asynchronous: AsyncBindingT | None | Unchanged = UNCHANGED, + ) -> None: + if not isinstance(sync, Unchanged): + self.sync.override(sync) + if not isinstance(asynchronous, Unchanged): + self.asynchronous.override(asynchronous) + + def reset(self) -> None: + self.sync.reset() + self.asynchronous.reset() + + def invoke( + self, + *, + call: Callable[[SyncBindingT], NativeT], + fallback: Callable[[], ResultT], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> ResultT: + return self.sync.invoke( + call=call, + fallback=fallback, + adapt=adapt, + context=context, + request_override=request_override, + eligible=eligible, + ) + + async def ainvoke( + self, + *, + call: Callable[[AsyncBindingT], Awaitable[NativeT]], + fallback: Callable[[], Awaitable[ResultT]], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> ResultT: + return await self.asynchronous.ainvoke( + call=call, + fallback=fallback, + adapt=adapt, + context=context, + request_override=request_override, + eligible=eligible, + ) + + def require( + self, + *, + call: Callable[[SyncBindingT], NativeT], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> ResultT: + return self.sync.require( + call=call, + adapt=adapt, + context=context, + request_override=request_override, + eligible=eligible, + ) + + async def arequire( + self, + *, + call: Callable[[AsyncBindingT], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, + request_override: bool | None = None, + eligible: bool = True, + ) -> ResultT: + return await self.asynchronous.arequire( + call=call, + adapt=adapt, + context=context, + request_override=request_override, + eligible=eligible, + ) -async def aattempt( - *, - native_call: Callable[[], Awaitable[NativeT]] | None, - adapt: Callable[[NativeT], ResultT], - context: BridgeErrorContext, -) -> RustAttempt[ResultT]: - if native_call is None: - return RustUnavailable() - exceptions: Final = native_exception_types() - if exceptions is None: - return RustHandled(adapt(await native_call())) - declined, upstream = exceptions - try: - value: Final = await native_call() - except declined as error: - return RustDeclined(reason=_decline_reason(error)) - except upstream as error: - _raise_upstream(error, context) - return RustHandled(adapt(value)) - - -def call(operation: Callable[[], ResultT], context: BridgeErrorContext) -> ResultT: - exceptions: Final = native_exception_types() - if exceptions is None: - return operation() - upstream: Final = exceptions[1] - try: - return operation() - except upstream as error: - _raise_upstream(error, context) - - -async def acall(operation: Callable[[], Awaitable[ResultT]], context: BridgeErrorContext) -> ResultT: - exceptions: Final = native_exception_types() - if exceptions is None: - return await operation() - upstream: Final = exceptions[1] - try: - return await operation() - except upstream as error: - _raise_upstream(error, context) - - -def _decline_reason(error: BaseException) -> str: +def _error_message(error: BaseException) -> str: reason: Final[object] = error.args[0] if error.args else str(error) return reason if isinstance(reason, str) else str(reason) -def _raise_required( - result: RustDeclined | RustUnavailable, - context: BridgeErrorContext, -) -> NoReturn: - raise RuntimeError(f"Rust {context.route} bridge {_required_reason(result)}") - - -def _required_reason(result: RustDeclined | RustUnavailable) -> str: - match result: - case RustUnavailable(): +def _required_reason(reason: PythonFallbackReason) -> str: + match reason: + case PythonFallbackReason.RUST_DISABLED: + return "is disabled" + case PythonFallbackReason.RUST_UNAVAILABLE: return "is unavailable" - case RustDeclined(reason=reason): - return f"declined the request: {reason}" + case PythonFallbackReason.RUST_DECLINED: + return "declined the request" -def _raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoReturn: - args: Final[tuple[object, ...]] = error.args - status_value: Final = args[0] if args else 0 - message_value: Final = args[1] if len(args) > 1 else str(error) - status: Final = status_value if isinstance(status_value, int) else 0 - message: Final = message_value if isinstance(message_value, str) else str(message_value) - raise APIError( - status_code=status or 500, - message=f"litellm rust {context.route}: {message}", - llm_provider=context.provider, - model=context.model, - ) from error +def always_enabled(*, request_override: bool | None = None) -> bool: + return True def identity(value: ResultT) -> ResultT: diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 3d71f6f8a50..9e0d340a325 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -1,11 +1,17 @@ from __future__ import annotations from collections.abc import Awaitable -from dataclasses import dataclass -from typing import Final, Protocol, cast +from typing import Final, Protocol import httpx +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + RustEndpoint, + always_enabled, + identity, +) from litellm.rust_bridge.timeouts import timeout_to_seconds @@ -39,62 +45,30 @@ class RustAtranscription(Protocol): raise NotImplementedError -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -@dataclass -class _RustTranscriptionState: - transcription: RustTranscription | None = None - atranscription: RustAtranscription | None = None - - -_STATE: Final = _RustTranscriptionState() +_TRANSCRIPTION: Final[RustEndpoint[RustTranscription, RustAtranscription]] = RustEndpoint.native( + route="audio transcription", + sync="transcription", + asynchronous="atranscription", + enabled=always_enabled, +) def configure_rust_transcription( enabled: bool = True, *, - transcription: RustTranscription | None | _Unset = _UNSET, - atranscription: RustAtranscription | None | _Unset = _UNSET, + transcription: RustTranscription | None | Unchanged = UNCHANGED, + atranscription: RustAtranscription | None | Unchanged = UNCHANGED, ) -> None: - if not isinstance(transcription, _Unset): - _STATE.transcription = transcription - if not isinstance(atranscription, _Unset): - _STATE.atranscription = atranscription - - -def load_rust_transcription() -> RustTranscription | None: - if _STATE.transcription is not None: - return _STATE.transcription - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - return ( - None - if native_bridge is None - else cast( # cast-ok: native extension protocol is runtime-defined - RustTranscription, getattr(native_bridge, "transcription", None) - ) - ) - - -def load_rust_atranscription() -> RustAtranscription | None: - if _STATE.atranscription is not None: - return _STATE.atranscription - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - return ( - None - if native_bridge is None - else cast( # cast-ok: native extension protocol is runtime-defined - RustAtranscription, getattr(native_bridge, "atranscription", None) - ) - ) + if not isinstance(transcription, Unchanged): + if transcription is None: + _TRANSCRIPTION.sync.reset() + else: + _TRANSCRIPTION.sync.override(transcription) + if not isinstance(atranscription, Unchanged): + if atranscription is None: + _TRANSCRIPTION.asynchronous.reset() + else: + _TRANSCRIPTION.asynchronous.override(atranscription) def transcription( @@ -107,19 +81,20 @@ def transcription( extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: - rust_transcription: Final = load_rust_transcription() - if rust_transcription is None: - return None - return rust_transcription( - model=model, - audio=audio, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=timeout_to_seconds(timeout), +) -> dict[str, object]: + return _TRANSCRIPTION.require( + call=lambda rust_transcription: rust_transcription( + model=model, + audio=audio, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_to_seconds(timeout), + ), + adapt=identity, + context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -133,17 +108,18 @@ async def atranscription( extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: - rust_atranscription: Final = load_rust_atranscription() - if rust_atranscription is None: - return None - return await rust_atranscription( - model=model, - audio=audio, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=timeout_to_seconds(timeout), +) -> dict[str, object]: + return await _TRANSCRIPTION.arequire( + call=lambda rust_atranscription: rust_atranscription( + model=model, + audio=audio, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_to_seconds(timeout), + ), + adapt=identity, + context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index b2cf253d164..e0de5e45997 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -8,7 +8,7 @@ import pytest import litellm from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.rust_bridge import configuration +from litellm.rust_bridge import bindings, configuration from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -123,7 +123,7 @@ def test_load_rust_messages_returns_injected_impl(): bridge = RecordingMessages() litellm.rust(True) rust_messages.set_rust_messages(messages=bridge) - assert rust_messages.load_rust_messages() is bridge + assert rust_messages._MESSAGES.sync.load() is bridge def test_bare_rust_still_toggles_ocr(): @@ -140,17 +140,16 @@ def test_load_rust_amessages_returns_injected_impl(): bridge = RecordingAsyncMessages() litellm.rust(True) rust_messages.set_rust_messages(amessages=bridge) - assert rust_messages.load_rust_amessages() is bridge + assert rust_messages._MESSAGES.asynchronous.load() is bridge def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + bindings, "get_native_bridge", lambda: None, ) litellm.rust(True) - assert rust_messages.load_rust_messages() is None result = rust_messages.messages( model="claude", body=REQUEST_BODY, @@ -248,14 +247,14 @@ async def test_gate_invokes_rust_and_marks_response_header(): @pytest.mark.asyncio -async def test_gate_falls_back_to_python_when_bridge_raises(): +async def test_gate_does_not_handoff_on_unknown_bridge_failure(): bridge = RaisingAsyncMessages() litellm.rust(True) rust_messages.set_rust_messages(amessages=bridge) - response = await _gate() + with pytest.raises(RuntimeError, match="upstream request failed"): + await _gate() - assert response is None assert bridge.calls == 1 @@ -405,7 +404,7 @@ async def test_fake_stream_wraps_rust_response_as_anthropic_sse(): @pytest.mark.asyncio async def test_gate_falls_back_when_bridge_unavailable(monkeypatch): monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + bindings, "get_native_bridge", lambda: None, ) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 1c2e07e0d24..60abb7d0f87 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -3,14 +3,15 @@ import builtins import importlib import types -from typing import Any +from typing import Any, Final +from unittest.mock import Mock import httpx import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.rust_bridge import configuration +from litellm.rust_bridge import bindings, configuration # `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` # function onto `litellm.ocr` and shadows the submodule, so import the modules @@ -259,14 +260,14 @@ def test_env_var_enables_rust_ocr(monkeypatch): def test_explicit_false_overrides_process_enable(): litellm.rust(True) - assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False + assert ocr_main._rust_request_override(build_prepared_request(litellm_params={"rust": False})) is False def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - assert rust_bridge.load_rust_ocr() is bridge + assert rust_bridge._OCR.sync.load() is bridge def test_native_bridge_loader_returns_none_when_extension_absent(monkeypatch): @@ -330,7 +331,7 @@ def test_load_rust_aocr_returns_injected_impl(): bridge = RecordingAsyncBridge() litellm.rust(True) rust_bridge.set_rust_ocr(aocr=bridge) - assert rust_bridge.load_rust_aocr() is bridge + assert rust_bridge._OCR.asynchronous.load() is bridge def test_toggle_without_ocr_arg_preserves_injected_impl(): @@ -341,16 +342,16 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) litellm.rust(False) - assert rust_bridge.load_rust_ocr() is bridge - assert rust_bridge.load_rust_aocr() is async_bridge + assert rust_bridge._OCR.sync.load() is bridge + assert rust_bridge._OCR.asynchronous.load() is async_bridge litellm.rust(True) - assert rust_bridge.load_rust_ocr() is bridge - assert rust_bridge.load_rust_aocr() is async_bridge + assert rust_bridge._OCR.sync.load() is bridge + assert rust_bridge._OCR.asynchronous.load() is async_bridge def test_explicit_ocr_none_clears_injected_impl(monkeypatch): monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + bindings, "get_native_bridge", lambda: None, ) @@ -360,21 +361,21 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch): rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) rust_bridge.set_rust_ocr(ocr=None, aocr=None) - assert rust_bridge.load_rust_ocr() is None - assert rust_bridge.load_rust_aocr() is None + assert rust_bridge._OCR.sync.load() is None + assert rust_bridge._OCR.asynchronous.load() is None def test_load_rust_ocr_none_when_extension_absent(monkeypatch): """With no injected impl and no compiled wheel, the loader returns None so the caller degrades to the Python path instead of raising ImportError.""" monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + bindings, "get_native_bridge", lambda: None, ) litellm.rust(True) # no impl injected; extension isn't built in CI - assert rust_bridge.load_rust_ocr() is None - assert rust_bridge.load_rust_aocr() is None + assert rust_bridge._OCR.sync.load() is None + assert rust_bridge._OCR.asynchronous.load() is None def test_load_rust_ocr_uses_compiled_extension(monkeypatch): @@ -385,14 +386,14 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + bindings, "get_native_bridge", lambda: fake_module, ) litellm.rust(True) # enabled, no impl injected -> import the extension - assert rust_bridge.load_rust_ocr() is fake_module.ocr - assert rust_bridge.load_rust_aocr() is fake_module.aocr + assert rust_bridge._OCR.sync.load() is fake_module.ocr + assert rust_bridge._OCR.asynchronous.load() is fake_module.aocr def test_timeout_to_seconds_handles_float_timeout_and_none(): @@ -408,14 +409,20 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): rust_bridge.set_rust_ocr(ocr=bridge) response = rust_bridge.ocr( + prepare=lambda: rust_bridge.RustOCRRequest( + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base="https://proxy.internal", + custom_llm_provider="mistral", + extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True, "pages": [0]}, + timeout=12.5, + ), model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://proxy.internal", - custom_llm_provider="mistral", - extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True, "pages": [0]}, - timeout=12.5, + provider="mistral", + request_override=True, + eligible=True, ) assert response == FAKE_OCR_RESPONSE @@ -443,14 +450,20 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): rust_bridge.set_rust_ocr(aocr=bridge) response = await rust_bridge.aocr( + prepare=lambda: rust_bridge.RustOCRRequest( + model="mistral-ocr-maas", + document=DOCUMENT, + api_key=None, + api_base=None, + custom_llm_provider="vertex_ai", + extra_headers=None, + optional_params={"vertex_project": "project-1"}, + timeout=httpx.Timeout(30.0, read=42.0), + ), model="mistral-ocr-maas", - document=DOCUMENT, - api_key=None, - api_base=None, - custom_llm_provider="vertex_ai", - extra_headers=None, - optional_params={"vertex_project": "project-1"}, - timeout=httpx.Timeout(30.0, read=42.0), + provider="vertex_ai", + request_override=True, + eligible=True, ) assert response == FAKE_OCR_RESPONSE @@ -808,23 +821,27 @@ def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) -def test_ocr_does_not_route_to_rust_when_disabled(): - """With the flag off, the bridge must not be consulted even if an impl exists.""" - bridge = RecordingBridge() - litellm.rust(False) - rust_bridge.set_rust_ocr(ocr=bridge) +def test_ocr_disabled_never_loads_or_prepares_rust(monkeypatch: pytest.MonkeyPatch): + def fail() -> None: + pytest.fail("disabled OCR must not load or prepare Rust") - assert rust_bridge.rust_ocr_enabled() is False - # The impl stays available for injection, but the disabled flag gates usage, - # so ocr() never reaches the Rust path (asserted via the enabled-path test). - assert bridge.calls == [] + python_ocr: Final = Mock(return_value=OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr")) + litellm.rust(False) + monkeypatch.setattr(bindings, "get_native_bridge", fail) + monkeypatch.setattr(ocr_main, "_prepare_rust_ocr_call", fail) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", python_ocr) + + response: Final = litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert isinstance(response, OCRResponse) + python_ocr.assert_called_once() def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): """Rust enabled but no bridge available (no injected impl, no compiled wheel): ocr() must degrade to the Python HTTP handler instead of raising.""" - monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) - litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI + rust_bridge._OCR.sync.override(None) + litellm.rust(True) captured = {} diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 4b446368dbe..b5e1f340ca9 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -76,14 +76,14 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: @pytest.mark.asyncio async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(responses_websocket, "_STATE", responses_websocket._RustResponsesWebSocketState()) - monkeypatch.setattr(responses_websocket, "get_native_bridge", lambda: None) + responses_websocket._RESPONSES_WEBSOCKET.override(None) assert ( await responses_websocket.connect( url="wss://example.test/responses", headers={}, timeout=None, + request_override=True, ) is None ) @@ -99,6 +99,7 @@ async def test_enabled_bridge_connects_and_adapts_socket( url="wss://example.test/responses", headers={"Authorization": "Bearer key"}, timeout=1.0, + request_override=True, ) assert connection is not None diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py index 88036a5a556..41d283ce59d 100644 --- a/tests/test_litellm/rust_bridge/test_bindings.py +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -33,3 +33,21 @@ def test_binding_validates_native_attribute( binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None) assert binding.load() == expected + + +@pytest.mark.parametrize("value", (None, 3, "not callable")) +def test_callable_binding_rejects_non_callables(monkeypatch: pytest.MonkeyPatch, value: object) -> None: + monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(route=value)) + + binding: bindings.NativeBinding[object] = bindings.NativeBinding.callable("route") + + assert binding.load() is None + + +def test_callable_binding_resolves_callable(monkeypatch: pytest.MonkeyPatch) -> None: + native = SimpleNamespace(route=lambda: "native") + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + + binding: bindings.NativeBinding[object] = bindings.NativeBinding.callable("route") + + assert binding.load() is native.route diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 0489f4ff017..5cbe2a5b44f 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -10,8 +10,8 @@ from __future__ import annotations import pytest import litellm +from litellm.rust_bridge import bindings, configuration from litellm.rust_bridge import chat_completions as bridge -from litellm.rust_bridge import configuration from litellm.types.utils import ModelResponse RUST_RESPONSE = { @@ -54,7 +54,7 @@ class _FakeNative: def _fake_native_bridge(monkeypatch): """Expose the bridge's exception classes without the compiled extension.""" - monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) def _hide_native_bridge(monkeypatch): @@ -63,7 +63,7 @@ def _hide_native_bridge(monkeypatch): There is no injection seam for "the .so is absent", so the loader itself is replaced; every other case here uses `set_rust_chat_completions`. """ - monkeypatch.setattr(bridge, "get_native_bridge", lambda: None) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) @pytest.fixture(autouse=True) @@ -244,6 +244,7 @@ def _call_kwargs(model_response: ModelResponse) -> dict: "extra_headers": {}, "timeout": 30.0, "on_response": lambda _rust_response: None, + "request_override": True, } diff --git a/tests/test_litellm/rust_bridge/test_responses.py b/tests/test_litellm/rust_bridge/test_responses.py new file mode 100644 index 00000000000..9aea36e5509 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_responses.py @@ -0,0 +1,47 @@ +from unittest.mock import Mock + +import pytest + +import litellm +from litellm.rust_bridge import responses + + +@pytest.fixture(autouse=True) +def reset_responses_endpoint(): + responses._RESPONSES.reset() + yield + responses._RESPONSES.reset() + + +def test_unavailable_responses_bridge_does_not_prepare_request() -> None: + prepare = Mock(side_effect=AssertionError("unavailable Rust must not prepare a request")) + responses._RESPONSES.override(sync=None) + litellm.rust(True) + + result = responses.responses( + prepare=prepare, + model="gpt-5", + provider="openai", + request_override=True, + ) + + assert result is None + prepare.assert_not_called() + + +@pytest.mark.asyncio +async def test_unavailable_async_responses_bridge_does_not_prepare_request( +) -> None: + prepare = Mock(side_effect=AssertionError("unavailable Rust must not prepare a request")) + responses._RESPONSES.override(asynchronous=None) + litellm.rust(True) + + result = await responses.aresponses( + prepare=prepare, + model="gpt-5", + provider="openai", + request_override=True, + ) + + assert result is None + prepare.assert_not_called() diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index b0fa510069b..1003b112337 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -26,21 +26,56 @@ def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None: def context() -> runtime.BridgeErrorContext: - return runtime.BridgeErrorContext(route="messages", provider="anthropic", model="model") + return runtime.BridgeErrorContext(provider="anthropic", model="model") -def test_invoke_tags_native_decline_before_running_fallback() -> None: +def enabled(*, request_override: bool | None = None) -> bool: + return request_override is not False + + +def test_disabled_route_does_not_load_or_call_rust() -> None: + bridge = runtime.RustBridge[object]( + route="messages", + load=lambda: pytest.fail("disabled route must not load Rust"), + enabled=enabled, + ) + + value = bridge.invoke( + call=lambda _binding: pytest.fail("disabled route must not call Rust"), + fallback=lambda: "python", + adapt=str, + context=context(), + request_override=False, + ) + + assert value == "python" + + +def test_unavailable_route_hands_off_to_python() -> None: + bridge = runtime.RustBridge[object](route="messages", load=lambda: None, enabled=enabled) + + value = bridge.invoke( + call=lambda _binding: pytest.fail("unavailable route must not call Rust"), + fallback=lambda: "python", + adapt=str, + context=context(), + ) + + assert value == "python" + + +def test_decline_hands_off_to_python() -> None: calls: list[str] = [] - def decline() -> object: + def decline(_binding: object) -> object: calls.append("rust") raise RustBridgeDeclined("unsupported") - value = runtime.invoke( - native_call=decline, + bridge = runtime.RustBridge(route="messages", load=object, enabled=enabled) + value = bridge.invoke( + call=decline, fallback=lambda: calls.append("python") or "fallback", adapt=str, - mode=runtime.FallbackMode.PYTHON, context=context(), ) @@ -48,48 +83,93 @@ def test_invoke_tags_native_decline_before_running_fallback() -> None: assert calls == ["rust", "python"] -def test_invoke_translates_upstream_without_fallback() -> None: - def fail() -> object: +def test_upstream_failure_never_hands_off() -> None: + def fail(_binding: object) -> object: raise RustUpstreamError(429, "rate limited") + bridge = runtime.RustBridge(route="messages", load=object, enabled=enabled) with pytest.raises(APIError, match="rate limited") as caught: - runtime.invoke( - native_call=fail, + bridge.invoke( + call=fail, fallback=lambda: pytest.fail("fallback must not run"), adapt=str, - mode=runtime.FallbackMode.PYTHON, context=context(), ) assert caught.value.status_code == 429 +def test_unknown_failure_never_hands_off() -> None: + bridge = runtime.RustBridge(route="messages", load=object, enabled=enabled) + + with pytest.raises(RuntimeError, match="unknown"): + bridge.invoke( + call=lambda _binding: (_ for _ in ()).throw(RuntimeError("unknown")), + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + context=context(), + ) + + @pytest.mark.asyncio -async def test_ainvoke_handles_native_success() -> None: - async def native() -> int: +async def test_async_route_handles_native_success() -> None: + async def native(_binding: object) -> int: return 3 async def fallback() -> str: pytest.fail("fallback must not run") - assert ( - await runtime.ainvoke( - native_call=native, - fallback=fallback, - adapt=str, - mode=runtime.FallbackMode.PYTHON, - context=context(), - ) - == "3" + bridge = runtime.RustBridge(route="messages", load=object, enabled=enabled) + value = await bridge.ainvoke( + call=native, + fallback=fallback, + adapt=str, + context=context(), ) + assert value == "3" -def test_required_mode_rejects_unavailable_bridge() -> None: - with pytest.raises(RuntimeError, match="is unavailable"): - runtime.invoke( - native_call=None, - fallback=lambda: pytest.fail("fallback must not run"), + +def test_required_route_rejects_unavailable_bridge() -> None: + bridge = runtime.RustBridge[object](route="messages", load=lambda: None, enabled=enabled) + + with pytest.raises(RuntimeError, match="unavailable"): + bridge.require( + call=lambda _binding: pytest.fail("unavailable route must not call Rust"), adapt=str, - mode=runtime.FallbackMode.RUST_REQUIRED, context=context(), ) + + +def test_native_endpoint_resolves_overrides_and_resets(monkeypatch: pytest.MonkeyPatch) -> None: + def native_sync() -> str: + return "native" + + async def native_async() -> str: + return "native async" + + native = SimpleNamespace( + sync_route=native_sync, + async_route=native_async, + RustBridgeDeclined=RustBridgeDeclined, + RustUpstreamError=RustUpstreamError, + ) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + endpoint: runtime.RustEndpoint[object, object] = runtime.RustEndpoint.native( + route="test", + sync="sync_route", + asynchronous="async_route", + enabled=enabled, + ) + + assert endpoint.sync.load() is native_sync + assert endpoint.asynchronous.load() is native_async + + replacement = object() + endpoint.override(sync=replacement, asynchronous=None) + assert endpoint.sync.load() is replacement + assert endpoint.asynchronous.load() is None + + endpoint.reset() + assert endpoint.sync.load() is native_sync + assert endpoint.asynchronous.load() is native_async diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index bbeb6c38f78..da7d1ed8533 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -4,6 +4,7 @@ import pytest import litellm from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch +from litellm.rust_bridge import bindings rust_bridge = importlib.import_module("litellm.rust_bridge.transcription") @@ -77,13 +78,13 @@ async def test_enabled_async_bridge() -> None: def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None: rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) - monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", lambda: None) - assert rust_bridge.load_rust_transcription() is None - assert rust_bridge.load_rust_atranscription() is None + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + assert rust_bridge._TRANSCRIPTION.sync.load() is None + assert rust_bridge._TRANSCRIPTION.asynchronous.load() is None def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None) + rust_bridge._TRANSCRIPTION.sync.override(None) with pytest.raises(RuntimeError, match="bridge is unavailable"): BedrockAudioTranscriptionRustDispatch().audio_transcriptions( @@ -100,10 +101,7 @@ def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> @pytest.mark.asyncio async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: - async def unavailable(**_: object) -> None: - return None - - monkeypatch.setattr(rust_bridge, "atranscription", unavailable) + rust_bridge._TRANSCRIPTION.asynchronous.override(None) with pytest.raises(RuntimeError, match="bridge is unavailable"): await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions(