From 20de0ce75570dacf3e6b17b25bd5ed67cee1a665 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 12:59:50 -0700 Subject: [PATCH] refactor(native): share dispatch lifecycle across existing bridges --- basedpyright-code-budget.json | 6 +- litellm/rust_bridge/bindings.py | 41 +- litellm/rust_bridge/chat_completions.py | 294 +++------ litellm/rust_bridge/messages.py | 150 ++--- litellm/rust_bridge/ocr.py | 119 ++-- litellm/rust_bridge/protocols.py | 180 ++++++ litellm/rust_bridge/responses_websocket.py | 95 ++- litellm/rust_bridge/runtime.py | 581 ++++++++++++++---- litellm/rust_bridge/transcription.py | 164 ++--- ruff-strict-budget.json | 4 +- .../test_rust_bridge_messages.py | 4 +- .../chat/test_anthropic_chat_handler.py | 6 +- .../chat/test_bedrock_converse_handler.py | 8 +- tests/test_litellm/ocr/test_rust_bridge.py | 56 +- .../responses/test_rust_bridge_websocket.py | 32 +- .../test_litellm/rust_bridge/test_bindings.py | 121 +++- .../rust_bridge/test_chat_completions.py | 25 +- .../test_litellm/rust_bridge/test_runtime.py | 436 +++++++++++-- .../test_audio_transcription_rust_bridge.py | 2 +- type-discipline-budget.json | 2 +- 20 files changed, 1566 insertions(+), 760 deletions(-) create mode 100644 litellm/rust_bridge/protocols.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 57ca267e504..0ed951b789f 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 13429 + "limit": 13428 }, "reportArgumentType": { - "limit": 2198 + "limit": 2196 }, "reportAssignmentType": { "limit": 319 @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44358 + "limit": 44247 }, "reportUnknownLambdaType": { "limit": 109 diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py index d16f150a2aa..e0ecdba5b58 100644 --- a/litellm/rust_bridge/bindings.py +++ b/litellm/rust_bridge/bindings.py @@ -1,9 +1,10 @@ from __future__ import annotations from collections.abc import Callable -from typing import Final, Generic, TypeVar +from typing import Final, Generic, TypeVar, cast # noqa: TID251 # PyO3 module boundary from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge.protocols import NativeModule BindingT = TypeVar("BindingT") @@ -15,12 +16,18 @@ class _Unset: _UNSET: Final = _Unset() +class Unchanged: + pass + + +UNCHANGED: Final = Unchanged() + + class NativeBinding(Generic[BindingT]): """Resolve one native attribute with an explicit, resettable test override.""" - def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None: - self._attribute: Final = attribute - self._validate: Final = validate + def __init__(self, select: Callable[[NativeModule], BindingT]) -> None: + self._select: Final = select self._override: BindingT | None | _Unset = _UNSET def load(self) -> BindingT | None: @@ -29,7 +36,12 @@ class NativeBinding(Generic[BindingT]): native: Final = get_native_bridge() if native is None: return None - return self._validate(getattr(native, self._attribute, None)) + module: Final = cast(NativeModule, native) # cast-ok: PyO3 exports are validated individually below + try: + value: Final = self._select(module) + except AttributeError: + return None + return value if callable(value) else None def override(self, value: BindingT | None) -> None: self._override = value @@ -38,12 +50,19 @@ class NativeBinding(Generic[BindingT]): self._override = _UNSET +_DECLINED: Final = NativeBinding(lambda native: native.RustBridgeDeclined) +_UPSTREAM: Final = NativeBinding(lambda native: native.RustUpstreamError) + + +def _exception_class(value: object) -> type[BaseException] | None: + if isinstance(value, type) and issubclass(value, BaseException): + return value + return None + + def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None: - native: Final = get_native_bridge() - if native is None: - return None - declined: Final = getattr(native, "RustBridgeDeclined", None) - upstream: Final = getattr(native, "RustUpstreamError", None) - if not isinstance(declined, type) or not isinstance(upstream, type): + declined: Final = _exception_class(_DECLINED.load()) + upstream: Final = _exception_class(_UPSTREAM.load()) + if declined is None or upstream is None: return None return declined, upstream diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index 674bd8847f7..ded533c5ad0 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -14,20 +14,29 @@ from __future__ import annotations import json from collections.abc import Awaitable, Callable, Mapping, Sequence -from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol import httpx from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger -from litellm.exceptions import APIError from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( convert_to_model_response_object, ) from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged from litellm.rust_bridge.configuration import rust_enabled -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge.protocols import ( + RustAchatCompletions, + RustChatCompletions, + RustChatCompletionsDecline, +) +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + EndpointBinding, + EndpointDispatch, + async_none, +) from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.utils import ModelResponse @@ -45,47 +54,6 @@ _LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) RUST_RESPONSE_HEADER: Final = "x-litellm-rust" -class RustChatCompletions(Protocol): - def __call__( - self, - model: str, - messages: Sequence[object], - optional_params: Mapping[str, object] | None, - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: Mapping[str, object] | None, - timeout_seconds: float | None, - ) -> Mapping[str, object]: - raise NotImplementedError - - -class RustAchatCompletions(Protocol): - def __call__( - self, - model: str, - messages: Sequence[object], - optional_params: Mapping[str, object] | None, - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: Mapping[str, object] | None, - timeout_seconds: float | None, - ) -> Awaitable[Mapping[str, object]]: - raise NotImplementedError - - -class RustChatCompletionsDecline(Protocol): - def __call__( - self, - model: str, - messages: Sequence[object], - optional_params: Mapping[str, object] | None, - custom_llm_provider: str | None, - ) -> str | None: - raise NotImplementedError - - class ResponseObserver(Protocol): """Invoked with the payload the core returned, on success only. @@ -126,67 +94,42 @@ def response_logger( return log -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -@dataclass(slots=True) -class _RustChatCompletionsState: - chat_completions: RustChatCompletions | None = None - achat_completions: RustAchatCompletions | None = None - decline: RustChatCompletionsDecline | None = None - - -_STATE: Final[_RustChatCompletionsState] = _RustChatCompletionsState() +_CHAT: Final[EndpointDispatch[RustChatCompletions, RustAchatCompletions]] = EndpointDispatch.native( + route="chat_completions", + sync=lambda native: native.chat_completions, + asynchronous=lambda native: native.achat_completions, + enabled=rust_enabled, +) +_CHAT_PREFLIGHT: Final[EndpointBinding[RustChatCompletionsDecline]] = EndpointBinding.native( + route="chat_completions", + select=lambda native: native.chat_completions_decline, + enabled=rust_enabled, +) def set_rust_chat_completions( *, - chat_completions: RustChatCompletions | None | _Unset = _UNSET, - achat_completions: RustAchatCompletions | None | _Unset = _UNSET, - decline: RustChatCompletionsDecline | None | _Unset = _UNSET, + chat_completions: RustChatCompletions | None | Unchanged = UNCHANGED, + achat_completions: RustAchatCompletions | None | Unchanged = UNCHANGED, + decline: RustChatCompletionsDecline | None | Unchanged = UNCHANGED, ) -> None: """Inject the native callables, so tests can supply a double instead of patching module attributes.""" - if not isinstance(chat_completions, _Unset): - _STATE.chat_completions = chat_completions - if not isinstance(achat_completions, _Unset): - _STATE.achat_completions = achat_completions - if not isinstance(decline, _Unset): - _STATE.decline = decline - - -def load_rust_chat_completions() -> RustChatCompletions | None: - if _STATE.chat_completions is not None: - return _STATE.chat_completions - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - loaded: RustChatCompletions | None = getattr(native_bridge, "chat_completions", None) - return loaded - - -def load_rust_achat_completions() -> RustAchatCompletions | None: - if _STATE.achat_completions is not None: - return _STATE.achat_completions - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - loaded: RustAchatCompletions | None = getattr(native_bridge, "achat_completions", None) - return loaded - - -def _load_rust_decline() -> RustChatCompletionsDecline | None: - if _STATE.decline is not None: - return _STATE.decline - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - loaded: RustChatCompletionsDecline | None = getattr(native_bridge, "chat_completions_decline", None) - return loaded + if not isinstance(chat_completions, Unchanged): + if chat_completions is None: + _CHAT.sync.reset() + else: + _CHAT.sync.override(chat_completions) + if not isinstance(achat_completions, Unchanged): + if achat_completions is None: + _CHAT.asynchronous.reset() + else: + _CHAT.asynchronous.override(achat_completions) + if not isinstance(decline, Unchanged): + if decline is None: + _CHAT_PREFLIGHT.reset() + else: + _CHAT_PREFLIGHT.override(decline) def _anthropic_user_id_reaches_the_body(litellm_params: Mapping[str, object] | None) -> bool: @@ -247,81 +190,16 @@ def rust_chat_completions_accepts( return False if stream: return False - if not rust_enabled(): - return False if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") return False - decline: Final = _load_rust_decline() - if decline is None: - return False - try: - reason: Final = decline( + return _CHAT_PREFLIGHT.accepts( + check=lambda decline: decline( model=model, messages=messages, optional_params=optional_params, custom_llm_provider=custom_llm_provider, - ) - except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path - verbose_logger.debug( - "Rust chat completions gate raised %s; staying on the Python path", - type(rust_error).__name__, - ) - return False - if reason is not None: - verbose_logger.debug("Rust chat completions declined (%s); using the Python path", reason) - return False - return True - - -def _rust_bridge_exceptions() -> tuple[type[BaseException], type[BaseException]] | None: - """`(declined, upstream_failed)` from the native module, or None when absent.""" - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - declined: Final = getattr(native_bridge, "RustBridgeDeclined", None) - upstream: Final = getattr(native_bridge, "RustUpstreamError", None) - if declined is None or upstream is None: - return None - return declined, upstream - - -def _reraise_or_decline( - rust_error: BaseException, - *, - model: str, - custom_llm_provider: str | None, -) -> None: - """Re-raise a failure the provider already saw, or return so the caller declines. - - A request that never reached the provider is safe to serve on the Python - path. One that did is not: the provider has already done the work, so a - second attempt bills for it twice. Those surface as an `APIError` carrying - the upstream status, which LiteLLM's exception mapping already understands. - """ - exceptions: Final = _rust_bridge_exceptions() - if exceptions is None: - verbose_logger.debug( - "Rust chat completions bridge raised %s; falling back to Python path", - type(rust_error).__name__, - ) - return - declined, upstream_failed = exceptions - if isinstance(rust_error, upstream_failed): - args: Final = rust_error.args - status: Final = args[0] if args else 0 - message: Final = args[1] if len(args) > 1 else "" - raise APIError( - status_code=int(status) or 500, - message=f"litellm rust chat completions: {message}", - llm_provider=custom_llm_provider or "", - model=model, - ) - if not isinstance(rust_error, declined): - raise rust_error - verbose_logger.debug( - "Rust chat completions declined before calling the provider (%s); using the Python path", - rust_error, + ), ) @@ -352,11 +230,13 @@ def chat_completions( timeout: float | httpx.Timeout | None, on_response: ResponseObserver, ) -> ModelResponse | None: - rust_chat_completions: Final = load_rust_chat_completions() - if rust_chat_completions is None: - return None - try: - rust_response: Final = rust_chat_completions( + def adapt(rust_response: Mapping[str, object]) -> ModelResponse: + on_response(rust_response) + return _build_model_response(rust_response, model_response) + + return _CHAT.invoke( + prepare=lambda: timeout_to_seconds(timeout), + call=lambda rust_chat_completions, timeout_seconds: rust_chat_completions( model=model, messages=messages, optional_params=optional_params, @@ -364,13 +244,12 @@ def chat_completions( api_base=api_base, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), - ) - except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw - _reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider) - return None - on_response(rust_response) - return _build_model_response(rust_response, model_response) + timeout_seconds=timeout_seconds, + ), + fallback=lambda: None, + adapt=adapt, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + ) async def achat_completions( @@ -386,11 +265,13 @@ async def achat_completions( timeout: float | httpx.Timeout | None, on_response: ResponseObserver, ) -> ModelResponse | None: - rust_achat_completions: Final = load_rust_achat_completions() - if rust_achat_completions is None: - return None - try: - rust_response: Final = await rust_achat_completions( + def adapt(rust_response: Mapping[str, object]) -> ModelResponse: + on_response(rust_response) + return _build_model_response(rust_response, model_response) + + return await _CHAT.ainvoke( + prepare=lambda: timeout_to_seconds(timeout), + call=lambda rust_achat_completions, timeout_seconds: rust_achat_completions( model=model, messages=messages, optional_params=optional_params, @@ -398,13 +279,12 @@ async def achat_completions( api_base=api_base, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), - ) - except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw - _reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider) - return None - on_response(rust_response) - return _build_model_response(rust_response, model_response) + timeout_seconds=timeout_seconds, + ), + fallback=async_none, + adapt=adapt, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + ) async def achat_completions_or_fallback( @@ -429,18 +309,24 @@ async def achat_completions_or_fallback( already returned a coroutine by the time a Rust failure surfaces, and so cannot fall back on its own. """ - response: Final = await achat_completions( - model=model, - messages=messages, - optional_params=optional_params, - model_response=model_response, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout=timeout, - on_response=on_response, + + def adapt(rust_response: Mapping[str, object]) -> object: + on_response(rust_response) + return _build_model_response(rust_response, model_response) + + return await _CHAT.ainvoke( + prepare=lambda: timeout_to_seconds(timeout), + call=lambda rust_achat_completions, timeout_seconds: rust_achat_completions( + model=model, + messages=messages, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_seconds, + ), + fallback=python_fallback, + adapt=adapt, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) - if response is not None: - return response - return await python_fallback() diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 40d0ddf622b..580b46d7c38 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -2,90 +2,54 @@ from __future__ import annotations -from collections.abc import Awaitable -from dataclasses import dataclass -from typing import Final, Protocol, cast +from typing import Final import httpx +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.protocols import RustAmessages, RustMessages +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + EndpointDispatch, + NativeErrorPolicy, + always_enabled, + async_none, + identity, +) from litellm.rust_bridge.timeouts import timeout_to_seconds - -class RustMessages(Protocol): - def __call__( - self, - model: str, - body: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - timeout_seconds: float | None, - ) -> dict[str, object]: - raise NotImplementedError - - -class RustAmessages(Protocol): - def __call__( - self, - model: str, - body: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - timeout_seconds: float | None, - ) -> Awaitable[dict[str, object]]: - raise NotImplementedError - - -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -@dataclass(slots=True) -class _RustMessagesState: - messages: RustMessages | None = None - amessages: RustAmessages | None = None - - -_STATE: Final[_RustMessagesState] = _RustMessagesState() +_MESSAGES: Final[EndpointDispatch[RustMessages, RustAmessages]] = EndpointDispatch.native( + route="messages", + sync=lambda native: native.messages, + asynchronous=lambda native: native.amessages, + enabled=always_enabled, + error_policy=NativeErrorPolicy.PROPAGATE, +) def set_rust_messages( *, - messages: RustMessages | None | _Unset = _UNSET, - amessages: RustAmessages | None | _Unset = _UNSET, + messages: RustMessages | None | Unchanged = UNCHANGED, + amessages: RustAmessages | None | Unchanged = UNCHANGED, ) -> None: - if not isinstance(messages, _Unset): - _STATE.messages = messages - if not isinstance(amessages, _Unset): - _STATE.amessages = amessages + if not isinstance(messages, Unchanged): + if messages is None: + _MESSAGES.sync.reset() + else: + _MESSAGES.sync.override(messages) + if not isinstance(amessages, Unchanged): + if amessages is None: + _MESSAGES.asynchronous.reset() + else: + _MESSAGES.asynchronous.override(amessages) def load_rust_messages() -> RustMessages | None: - if _STATE.messages is not None: - return _STATE.messages - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustMessages, getattr(native_bridge, "messages", None)) + return _MESSAGES.sync.load() def load_rust_amessages() -> RustAmessages | None: - if _STATE.amessages is not None: - return _STATE.amessages - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustAmessages, getattr(native_bridge, "amessages", None)) + return _MESSAGES.asynchronous.load() def messages( @@ -98,17 +62,20 @@ def messages( extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: - rust_messages: Final = load_rust_messages() - if rust_messages is None: - return None - return rust_messages( - model=model, - body=body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), + return _MESSAGES.invoke( + prepare=lambda: timeout_to_seconds(timeout), + call=lambda rust_messages, timeout_seconds: rust_messages( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_seconds, + ), + fallback=lambda: None, + adapt=identity, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -122,15 +89,18 @@ async def amessages( extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: - rust_amessages: Final = load_rust_amessages() - if rust_amessages is None: - return None - return await rust_amessages( - model=model, - body=body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), + return await _MESSAGES.ainvoke( + prepare=lambda: timeout_to_seconds(timeout), + call=lambda rust_amessages, timeout_seconds: rust_amessages( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_seconds, + ), + fallback=async_none, + adapt=identity, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index b7fdb5a98ef..734e48fd3ed 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,63 +2,36 @@ from __future__ import annotations -from collections.abc import Awaitable -from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables +from typing import Final import httpx -from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge import configuration as _configuration +from litellm.rust_bridge.protocols import RustAocr, RustOcr +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + EndpointDispatch, + NativeErrorPolicy, + async_none, + identity, +) from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds - -class RustOcr(Protocol): - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: - raise NotImplementedError - - -class RustAocr(Protocol): - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> Awaitable[dict[str, object]]: - raise NotImplementedError - - -def _as_ocr(value: object) -> RustOcr | None: - return cast(RustOcr, value) if callable(value) else None - - -def _as_aocr(value: object) -> RustAocr | None: - return cast(RustAocr, value) if callable(value) else None - - -_OCR: Final = NativeBinding("ocr", validate=_as_ocr) -_AOCR: Final = NativeBinding("aocr", validate=_as_aocr) +_OCR: Final[EndpointDispatch[RustOcr, RustAocr]] = EndpointDispatch.native( + route="ocr", + sync=lambda native: native.ocr, + asynchronous=lambda native: native.aocr, + enabled=_configuration.rust_enabled, + error_policy=NativeErrorPolicy.PROPAGATE, +) def load_rust_ocr() -> RustOcr | None: - return _OCR.load() + return _OCR.sync.load() def load_rust_aocr() -> RustAocr | None: - return _AOCR.load() + return _OCR.asynchronous.load() def ocr( @@ -72,18 +45,21 @@ def ocr( optional_params: dict[str, object], timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: - rust_ocr: Final = load_rust_ocr() - if rust_ocr is None: - return None - return rust_ocr( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=_timeout_to_seconds(timeout), + return _OCR.invoke( + prepare=lambda: _timeout_to_seconds(timeout), + call=lambda rust_ocr, timeout_seconds: rust_ocr( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ), + fallback=lambda: None, + adapt=identity, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -98,16 +74,19 @@ async def aocr( optional_params: dict[str, object], timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: - rust_aocr: Final = load_rust_aocr() - if rust_aocr is None: - return None - return await rust_aocr( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=_timeout_to_seconds(timeout), + return await _OCR.ainvoke( + prepare=lambda: _timeout_to_seconds(timeout), + call=lambda rust_aocr, timeout_seconds: rust_aocr( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ), + fallback=async_none, + adapt=identity, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) diff --git a/litellm/rust_bridge/protocols.py b/litellm/rust_bridge/protocols.py new file mode 100644 index 00000000000..b08fcbf81aa --- /dev/null +++ b/litellm/rust_bridge/protocols.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Mapping, Sequence +from typing import Protocol + + +class RustChatCompletions(Protocol): + def __call__( + self, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object] | None, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: Mapping[str, object] | None, + timeout_seconds: float | None, + ) -> Mapping[str, object]: ... + + +class RustAchatCompletions(Protocol): + def __call__( + self, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object] | None, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: Mapping[str, object] | None, + timeout_seconds: float | None, + ) -> Awaitable[Mapping[str, object]]: ... + + +class RustChatCompletionsDecline(Protocol): + def __call__( + self, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object] | None, + custom_llm_provider: str | None, + ) -> str | None: ... + + +class RustResponsesWebSocket(Protocol): + async def send_text(self, text: str) -> None: ... + + async def recv_text(self) -> str | None: ... + + async def close(self) -> None: ... + + +class RustResponsesWebSocketConnection(Protocol): + @classmethod + async def connect( + cls, + url: str, + headers: dict[str, str], + timeout_seconds: float | None, + ) -> RustResponsesWebSocket: ... + + +class NativeModule(Protocol): + @property + def chat_completions(self) -> RustChatCompletions: ... + + @property + def achat_completions(self) -> RustAchatCompletions: ... + + @property + def chat_completions_decline(self) -> RustChatCompletionsDecline: ... + + @property + def ResponsesWebSocketConnection(self) -> type[RustResponsesWebSocketConnection]: ... + + @property + def RustBridgeDeclined(self) -> type[BaseException]: ... + + @property + def RustUpstreamError(self) -> type[BaseException]: ... + + @property + def messages(self) -> RustMessages: ... + + @property + def amessages(self) -> RustAmessages: ... + + @property + def ocr(self) -> RustOcr: ... + + @property + def aocr(self) -> RustAocr: ... + + @property + def transcription(self) -> RustTranscription: ... + + @property + def atranscription(self) -> RustAtranscription: ... + + +class RustMessages(Protocol): + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> dict[str, object]: ... + + +class RustAmessages(Protocol): + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> Awaitable[dict[str, object]]: ... + + +class RustOcr(Protocol): + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: ... + + +class RustAocr(Protocol): + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> Awaitable[dict[str, object]]: ... + + +class RustTranscription(Protocol): + def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: ... + + +class RustAtranscription(Protocol): + def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> Awaitable[dict[str, object]]: ... diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 0634867af1c..8d0b7c263d0 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -2,67 +2,41 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import Final, Protocol +from typing import Final import httpx from websockets.exceptions import ConnectionClosedOK -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.configuration import rust_enabled +from litellm.rust_bridge.protocols import ( + RustResponsesWebSocket, + RustResponsesWebSocketConnection, +) +from litellm.rust_bridge.runtime import ( + AsyncEndpointDispatch, + BridgeErrorContext, + async_none, + identity, +) from litellm.rust_bridge.timeouts import timeout_to_seconds - -class RustResponsesWebSocket(Protocol): - async def send_text(self, text: str) -> None: ... - - async def recv_text(self) -> str | None: ... - - async def close(self) -> None: ... - - -class RustResponsesWebSocketConnection(Protocol): - @classmethod - async def connect( - cls, - url: str, - headers: dict[str, str], - timeout_seconds: float | None, - ) -> RustResponsesWebSocket: ... - - -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -@dataclass(slots=True) -class _RustResponsesWebSocketState: - connection: RustResponsesWebSocketConnection | None = None - - -_STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() +_RESPONSES_WEBSOCKET: Final[AsyncEndpointDispatch[RustResponsesWebSocketConnection]] = AsyncEndpointDispatch.native( + route="responses_websocket", + asynchronous=lambda native: native.ResponsesWebSocketConnection, + enabled=rust_enabled, +) def set_rust_responses_websocket( *, - connection: RustResponsesWebSocketConnection | None | _Unset = _UNSET, + connection: RustResponsesWebSocketConnection | None | Unchanged = UNCHANGED, ) -> None: - if not isinstance(connection, _Unset): - _STATE.connection = connection - - -def load_rust_responses_websocket() -> RustResponsesWebSocketConnection | None: - if _STATE.connection is not None: - return _STATE.connection - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - connection_type: Final[RustResponsesWebSocketConnection | None] = getattr( - native_bridge, "ResponsesWebSocketConnection", None - ) - return connection_type + if not isinstance(connection, Unchanged): + if connection is None: + _RESPONSES_WEBSOCKET.reset() + else: + _RESPONSES_WEBSOCKET.override(connection) class _ConnectionAdapter: @@ -88,15 +62,18 @@ async def connect( headers: dict[str, str], timeout: float | httpx.Timeout | None, ) -> _ConnectionAdapter | None: - connection_type: Final = load_rust_responses_websocket() - if connection_type is None: - return None try: - connection: Final = await connection_type.connect( - url=url, - headers=headers, - timeout_seconds=timeout_to_seconds(timeout), + connection: Final = await _RESPONSES_WEBSOCKET.ainvoke( + prepare=lambda: timeout_to_seconds(timeout), + call=lambda connection_type, timeout_seconds: connection_type.connect( + url=url, + headers=headers, + timeout_seconds=timeout_seconds, + ), + fallback=async_none, + adapt=identity, + error_context=BridgeErrorContext(provider="openai", model="responses websocket"), ) - except Exception: # noqa: BLE001 # bridge failures must fall back to Python + except Exception: # noqa: BLE001 # preserve the existing WebSocket connection fallback return None - return _ConnectionAdapter(connection) + return None if connection is None else _ConnectionAdapter(connection) diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index d411673439f..eef349b22b5 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -1,150 +1,517 @@ from __future__ import annotations from collections.abc import Awaitable, Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum -from typing import Final, Generic, NoReturn, TypeAlias, TypeVar +from typing import Final, Generic, NoReturn, Protocol, TypeAlias, TypeVar from litellm.exceptions import APIError -from litellm.rust_bridge.bindings import native_exception_types +from litellm.rust_bridge.bindings import ( + UNCHANGED, + NativeBinding, + Unchanged, + native_exception_types, +) +from litellm.rust_bridge.protocols import NativeModule +BindingT = TypeVar("BindingT") +SelectedT = TypeVar("SelectedT") +SelectedSyncT = TypeVar("SelectedSyncT") +SelectedAsyncT = TypeVar("SelectedAsyncT") NativeT = TypeVar("NativeT") +RequestT = TypeVar("RequestT") ResultT = TypeVar("ResultT") +SyncBindingT = TypeVar("SyncBindingT") +AsyncBindingT = TypeVar("AsyncBindingT") -class FallbackMode(Enum): - PYTHON = "python" - RUST_REQUIRED = "rust_required" +class PythonFallbackReason(Enum): + NATIVE_DISABLED = "native_disabled" + NATIVE_UNAVAILABLE = "native_unavailable" + NATIVE_DECLINED = "native_declined" + + +class NativeErrorPolicy(Enum): + TRANSLATE = "translate" + PROPAGATE = "propagate" @dataclass(frozen=True, slots=True) -class RustHandled(Generic[ResultT]): +class Handled(Generic[ResultT]): value: ResultT @dataclass(frozen=True, slots=True) -class RustDeclined: - reason: str +class PythonFallback: + reason: PythonFallbackReason + detail: str | None = None -@dataclass(frozen=True, slots=True) -class RustUnavailable: - pass - - -RustAttempt: TypeAlias = RustHandled[ResultT] | RustDeclined | RustUnavailable +DispatchResult: TypeAlias = Handled[ResultT] | PythonFallback @dataclass(frozen=True, slots=True) class BridgeErrorContext: - route: str provider: str model: str -def invoke( - *, - native_call: Callable[[], NativeT] | None, - fallback: Callable[[], ResultT], - adapt: Callable[[NativeT], ResultT], - mode: FallbackMode, - context: BridgeErrorContext, -) -> ResultT: - result: Final = attempt(native_call=native_call, adapt=adapt, context=context) - if isinstance(result, RustHandled): - return result.value - if mode is FallbackMode.PYTHON: - return fallback() - _raise_required(result, context) +class RustEnablement(Protocol): + def __call__(self) -> bool: ... -async def ainvoke( - *, - native_call: Callable[[], Awaitable[NativeT]] | None, - fallback: Callable[[], Awaitable[ResultT]], - adapt: Callable[[NativeT], ResultT], - mode: FallbackMode, - context: BridgeErrorContext, -) -> ResultT: - result: Final = await aattempt(native_call=native_call, adapt=adapt, context=context) - if isinstance(result, RustHandled): - return result.value - if mode is FallbackMode.PYTHON: - return await fallback() - _raise_required(result, context) +@dataclass(frozen=True, slots=True) +class EndpointBinding(Generic[BindingT]): + route: str + load: Callable[[], BindingT | None] + enabled: RustEnablement + error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE + _native_binding: NativeBinding[BindingT] | None = field(default=None, repr=False) + + @staticmethod + def native( + *, + route: str, + select: Callable[[NativeModule], SelectedT], + enabled: RustEnablement, + error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE, + ) -> EndpointBinding[SelectedT]: + binding: Final = NativeBinding(select) + return EndpointBinding( + route=route, + load=binding.load, + enabled=enabled, + error_policy=error_policy, + _native_binding=binding, + ) + + def override(self, value: BindingT | None) -> None: + if self._native_binding is None: + raise RuntimeError("only native Rust bridges support binding overrides") + self._native_binding.override(value) + + def reset(self) -> None: + if self._native_binding is None: + raise RuntimeError("only native Rust bridges support binding resets") + self._native_binding.reset() + + def _attempt( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[BindingT, RequestT], NativeT], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> DispatchResult[ResultT]: + binding_or_fallback: Final = self._binding_or_python_fallback( + eligible=eligible, + ) + if isinstance(binding_or_fallback, PythonFallback): + return binding_or_fallback + return self._attempt_call( + call=lambda: call(binding_or_fallback, prepare()), + adapt=adapt, + error_context=error_context, + ) + + async def _aattempt( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[BindingT, RequestT], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> DispatchResult[ResultT]: + binding_or_fallback: Final = self._binding_or_python_fallback( + eligible=eligible, + ) + if isinstance(binding_or_fallback, PythonFallback): + return binding_or_fallback + return await self._attempt_acall( + call=lambda: call(binding_or_fallback, prepare()), + adapt=adapt, + error_context=error_context, + ) + + def invoke( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[BindingT, RequestT], NativeT], + fallback: Callable[[], ResultT], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + result: Final = self._attempt( + prepare=prepare, + call=call, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) + match result: + case Handled(value=value): + return value + case PythonFallback(): + return fallback() + + async def ainvoke( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[BindingT, RequestT], Awaitable[NativeT]], + fallback: Callable[[], Awaitable[ResultT]], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + result: Final = await self._aattempt( + prepare=prepare, + call=call, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) + match result: + case Handled(value=value): + return value + case PythonFallback(): + return await fallback() + + def accepts( + self, + *, + check: Callable[[BindingT], str | None], + eligible: bool = True, + ) -> bool: + binding_or_fallback: Final = self._binding_or_python_fallback( + eligible=eligible, + ) + if isinstance(binding_or_fallback, PythonFallback): + return False + try: + reason: Final = check(binding_or_fallback) + except Exception: # noqa: BLE001 # preflight performs no provider I/O, so Python handoff is safe + return False + return reason is None + + def require( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[BindingT, RequestT], NativeT], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + result: Final = self._attempt( + prepare=prepare, + call=call, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) + match result: + case Handled(value=value): + return value + case PythonFallback(): + self._raise_required(result) + + async def arequire( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[BindingT, RequestT], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + result: Final = await self._aattempt( + prepare=prepare, + call=call, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) + match result: + case Handled(value=value): + return value + case PythonFallback(): + self._raise_required(result) + + def can_attempt( + self, + *, + eligible: bool = True, + ) -> bool: + return not isinstance( + self._binding_or_python_fallback(eligible=eligible), + PythonFallback, + ) + + def _raise_required(self, fallback: PythonFallback) -> NoReturn: + detail: Final = f": {fallback.detail}" if fallback.detail else "" + reason: Final = _required_reason(fallback.reason) + raise RuntimeError(f"native {self.route} endpoint {reason}{detail}") + + def _binding_or_python_fallback( + self, + *, + eligible: bool, + ) -> BindingT | PythonFallback: + if not eligible or not self.enabled(): + return PythonFallback(PythonFallbackReason.NATIVE_DISABLED) + binding: Final = self.load() + if binding is None: + return PythonFallback(PythonFallbackReason.NATIVE_UNAVAILABLE) + return binding + + def _attempt_call( + self, + *, + call: Callable[[], NativeT], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + ) -> DispatchResult[ResultT]: + if self.error_policy is NativeErrorPolicy.PROPAGATE: + return Handled(adapt(call())) + exceptions: Final = native_exception_types() + if exceptions is None: + try: + value_without_exceptions: Final = call() + except Exception as error: # noqa: BLE001 # preserve chat fallback when native exception classes are absent + return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error)) + return Handled(adapt(value_without_exceptions)) + declined, upstream = exceptions + try: + value: Final = call() + except declined as error: + return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error)) + except upstream as error: + self._raise_upstream(error, error_context) + return Handled(adapt(value)) + + async def _attempt_acall( + self, + *, + call: Callable[[], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + ) -> DispatchResult[ResultT]: + if self.error_policy is NativeErrorPolicy.PROPAGATE: + return Handled(adapt(await call())) + exceptions: Final = native_exception_types() + if exceptions is None: + try: + value_without_exceptions: Final = await call() + except Exception as error: # noqa: BLE001 # preserve chat fallback when native exception classes are absent + return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error)) + return Handled(adapt(value_without_exceptions)) + declined, upstream = exceptions + try: + value: Final = await call() + except declined as error: + return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error)) + except upstream as error: + self._raise_upstream(error, error_context) + return Handled(adapt(value)) + + def _raise_upstream(self, error: BaseException, error_context: BridgeErrorContext) -> NoReturn: + args: Final[tuple[object, ...]] = error.args + attribute_status: Final = getattr(error, "status_code", None) + attribute_message: Final = getattr(error, "message", None) + status_value: Final = attribute_status if isinstance(attribute_status, int) else (args[0] if args else 0) + message_value: Final = ( + attribute_message if isinstance(attribute_message, str) else (args[1] if len(args) > 1 else str(error)) + ) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + raise APIError( + status_code=status or 500, + message=f"litellm rust {self.route}: {message}", + llm_provider=error_context.provider, + model=error_context.model, + ) from error -def attempt( - *, - native_call: Callable[[], NativeT] | None, - adapt: Callable[[NativeT], ResultT], - context: BridgeErrorContext, -) -> RustAttempt[ResultT]: - if native_call is None: - return RustUnavailable() - exceptions: Final = native_exception_types() - if exceptions is None: - return RustHandled(adapt(native_call())) - declined, upstream = exceptions - try: - value: Final = native_call() - except declined as error: - return RustDeclined(reason=_decline_reason(error)) - except upstream as error: - _raise_upstream(error, context) - return RustHandled(adapt(value)) +@dataclass(frozen=True, slots=True) +class EndpointDispatch(Generic[SyncBindingT, AsyncBindingT]): + sync: EndpointBinding[SyncBindingT] + asynchronous: EndpointBinding[AsyncBindingT] + + @staticmethod + def native( + *, + route: str, + sync: Callable[[NativeModule], SelectedSyncT], + asynchronous: Callable[[NativeModule], SelectedAsyncT], + enabled: RustEnablement, + error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE, + ) -> EndpointDispatch[SelectedSyncT, SelectedAsyncT]: + return EndpointDispatch( + sync=EndpointBinding.native(route=route, select=sync, enabled=enabled, error_policy=error_policy), + asynchronous=EndpointBinding.native( + route=route, + select=asynchronous, + enabled=enabled, + error_policy=error_policy, + ), + ) + + def override( + self, + *, + sync: SyncBindingT | None | Unchanged = UNCHANGED, + asynchronous: AsyncBindingT | None | Unchanged = UNCHANGED, + ) -> None: + if not isinstance(sync, Unchanged): + self.sync.override(sync) + if not isinstance(asynchronous, Unchanged): + self.asynchronous.override(asynchronous) + + def reset(self) -> None: + self.sync.reset() + self.asynchronous.reset() + + def invoke( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[SyncBindingT, RequestT], NativeT], + fallback: Callable[[], ResultT], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + return self.sync.invoke( + prepare=prepare, + call=call, + fallback=fallback, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) + + async def ainvoke( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]], + fallback: Callable[[], Awaitable[ResultT]], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + return await self.asynchronous.ainvoke( + prepare=prepare, + call=call, + fallback=fallback, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) + + def require( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[SyncBindingT, RequestT], NativeT], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + return self.sync.require( + prepare=prepare, + call=call, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) + + async def arequire( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + return await self.asynchronous.arequire( + prepare=prepare, + call=call, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) -async def aattempt( - *, - native_call: Callable[[], Awaitable[NativeT]] | None, - adapt: Callable[[NativeT], ResultT], - context: BridgeErrorContext, -) -> RustAttempt[ResultT]: - if native_call is None: - return RustUnavailable() - exceptions: Final = native_exception_types() - if exceptions is None: - return RustHandled(adapt(await native_call())) - declined, upstream = exceptions - try: - value: Final = await native_call() - except declined as error: - return RustDeclined(reason=_decline_reason(error)) - except upstream as error: - _raise_upstream(error, context) - return RustHandled(adapt(value)) +@dataclass(frozen=True, slots=True) +class AsyncEndpointDispatch(Generic[AsyncBindingT]): + asynchronous: EndpointBinding[AsyncBindingT] + + @staticmethod + def native( + *, + route: str, + asynchronous: Callable[[NativeModule], SelectedAsyncT], + enabled: RustEnablement, + ) -> AsyncEndpointDispatch[SelectedAsyncT]: + return AsyncEndpointDispatch( + asynchronous=EndpointBinding.native(route=route, select=asynchronous, enabled=enabled) + ) + + def override(self, value: AsyncBindingT | None) -> None: + self.asynchronous.override(value) + + def reset(self) -> None: + self.asynchronous.reset() + + async def ainvoke( + self, + *, + prepare: Callable[[], RequestT], + call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]], + fallback: Callable[[], Awaitable[ResultT]], + adapt: Callable[[NativeT], ResultT], + error_context: BridgeErrorContext, + eligible: bool = True, + ) -> ResultT: + return await self.asynchronous.ainvoke( + prepare=prepare, + call=call, + fallback=fallback, + adapt=adapt, + error_context=error_context, + eligible=eligible, + ) -def _decline_reason(error: BaseException) -> str: +def _error_message(error: BaseException) -> str: reason: Final[object] = error.args[0] if error.args else str(error) return reason if isinstance(reason, str) else str(reason) -def _raise_required( - result: RustDeclined | RustUnavailable, - context: BridgeErrorContext, -) -> NoReturn: - raise RuntimeError(f"Rust {context.route} bridge {_required_reason(result)}") - - -def _required_reason(result: RustDeclined | RustUnavailable) -> str: - match result: - case RustUnavailable(): +def _required_reason(reason: PythonFallbackReason) -> str: + match reason: + case PythonFallbackReason.NATIVE_DISABLED: + return "is disabled" + case PythonFallbackReason.NATIVE_UNAVAILABLE: return "is unavailable" - case RustDeclined(reason=reason): - return f"declined the request: {reason}" + case PythonFallbackReason.NATIVE_DECLINED: + return "declined the request" -def _raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoReturn: - args: Final[tuple[object, ...]] = error.args - status_value: Final = args[0] if args else 0 - message_value: Final = args[1] if len(args) > 1 else str(error) - status: Final = status_value if isinstance(status_value, int) else 0 - message: Final = message_value if isinstance(message_value, str) else str(message_value) - raise APIError( - status_code=status or 500, - message=f"litellm rust {context.route}: {message}", - llm_provider=context.provider, - model=context.model, - ) from error +def always_enabled() -> bool: + return True + + +def identity(value: ResultT) -> ResultT: + return value + + +async def async_none() -> None: + return None diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 6c81786accd..b93ab5e625a 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -1,99 +1,53 @@ from __future__ import annotations -from collections.abc import Awaitable -from dataclasses import dataclass -from typing import Final, Protocol, cast +from typing import Final import httpx +from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.protocols import RustAtranscription, RustTranscription +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + EndpointDispatch, + NativeErrorPolicy, + always_enabled, + async_none, + identity, +) from litellm.rust_bridge.timeouts import timeout_to_seconds - -class RustTranscription(Protocol): - def __call__( - self, - model: str, - audio: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: - raise NotImplementedError - - -class RustAtranscription(Protocol): - def __call__( - self, - model: str, - audio: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> Awaitable[dict[str, object]]: - raise NotImplementedError - - -class _Unset: - pass - - -_UNSET: Final[_Unset] = _Unset() - - -@dataclass -class _RustTranscriptionState: - transcription: RustTranscription | None = None - atranscription: RustAtranscription | None = None - - -_STATE: Final = _RustTranscriptionState() +_TRANSCRIPTION: Final[EndpointDispatch[RustTranscription, RustAtranscription]] = EndpointDispatch.native( + route="audio transcription", + sync=lambda native: native.transcription, + asynchronous=lambda native: native.atranscription, + enabled=always_enabled, + error_policy=NativeErrorPolicy.PROPAGATE, +) def configure_rust_transcription( *, - transcription: RustTranscription | None | _Unset = _UNSET, - atranscription: RustAtranscription | None | _Unset = _UNSET, + transcription: RustTranscription | None | Unchanged = UNCHANGED, + atranscription: RustAtranscription | None | Unchanged = UNCHANGED, ) -> None: - if not isinstance(transcription, _Unset): - _STATE.transcription = transcription - if not isinstance(atranscription, _Unset): - _STATE.atranscription = atranscription + if not isinstance(transcription, Unchanged): + if transcription is None: + _TRANSCRIPTION.sync.reset() + else: + _TRANSCRIPTION.sync.override(transcription) + if not isinstance(atranscription, Unchanged): + if atranscription is None: + _TRANSCRIPTION.asynchronous.reset() + else: + _TRANSCRIPTION.asynchronous.override(atranscription) def load_rust_transcription() -> RustTranscription | None: - if _STATE.transcription is not None: - return _STATE.transcription - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - return ( - None - if native_bridge is None - else cast( # cast-ok: native extension protocol is runtime-defined - RustTranscription, getattr(native_bridge, "transcription", None) - ) - ) + return _TRANSCRIPTION.sync.load() def load_rust_atranscription() -> RustAtranscription | None: - if _STATE.atranscription is not None: - return _STATE.atranscription - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - return ( - None - if native_bridge is None - else cast( # cast-ok: native extension protocol is runtime-defined - RustAtranscription, getattr(native_bridge, "atranscription", None) - ) - ) + return _TRANSCRIPTION.asynchronous.load() def transcription( @@ -107,18 +61,21 @@ def transcription( optional_params: dict[str, object], timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: - rust_transcription: Final = load_rust_transcription() - if rust_transcription is None: - return None - return rust_transcription( - model=model, - audio=audio, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=timeout_to_seconds(timeout), + return _TRANSCRIPTION.invoke( + prepare=lambda: timeout_to_seconds(timeout), + call=lambda rust_transcription, timeout_seconds: rust_transcription( + model=model, + audio=audio, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ), + fallback=lambda: None, + adapt=identity, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -133,16 +90,19 @@ async def atranscription( optional_params: dict[str, object], timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: - rust_atranscription: Final = load_rust_atranscription() - if rust_atranscription is None: - return None - return await rust_atranscription( - model=model, - audio=audio, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=timeout_to_seconds(timeout), + return await _TRANSCRIPTION.ainvoke( + prepare=lambda: timeout_to_seconds(timeout), + call=lambda rust_atranscription, timeout_seconds: rust_atranscription( + model=model, + audio=audio, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ), + fallback=async_none, + adapt=identity, + error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc314..55699214431 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -156,7 +156,7 @@ "limit": 215 }, "PLW0603": { - "limit": 190 + "limit": 188 }, "PLW1508": { "limit": 190 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1035 + "limit": 1033 }, "TRY002": { "limit": 524 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 a30474245c6..e8126600467 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -135,7 +135,7 @@ def test_load_rust_amessages_returns_injected_impl(): def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + importlib.import_module("litellm.rust_bridge.bindings"), "get_native_bridge", lambda: None, ) @@ -383,7 +383,7 @@ async def test_fake_stream_wraps_rust_response_as_anthropic_sse(): @pytest.mark.asyncio async def test_gate_falls_back_when_bridge_unavailable(monkeypatch): monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + importlib.import_module("litellm.rust_bridge.bindings"), "get_native_bridge", lambda: None, ) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index b4b173b20c3..da70f422f44 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2467,7 +2467,7 @@ class TestRustChatCompletionsHook: def declining_native(**_kwargs): raise _Declined("blank message text") - monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()) bridge.set_rust_chat_completions( decline=lambda **_kwargs: None, chat_completions=declining_native ) @@ -2499,7 +2499,7 @@ class TestRustChatCompletionsHook: RustBridgeDeclined = _Declined RustUpstreamError = type("_Upstream", (Exception,), {}) - monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()) async def declining_native(**_kwargs): raise _Declined("blank message text") @@ -2559,7 +2559,7 @@ class TestRustChatCompletionsHook: RustBridgeDeclined = _Declined RustUpstreamError = type("_Upstream", (Exception,), {}) - monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()) def declining_native(**_kwargs): raise _Declined("blank message text") diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index f34b8eb1fb9..becd6ecb832 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -205,7 +205,7 @@ async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch): RustBridgeDeclined = _Declined RustUpstreamError = type("_Upstream", (Exception,), {}) - monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()) async def declining_native(**_kwargs): raise _Declined("blank message text") @@ -282,7 +282,7 @@ async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines(): return ModelResponse() with ( - patch.object(bridge, "get_native_bridge", lambda: _FakeNative()), + patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()), patch.object( BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS ), @@ -389,7 +389,7 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines(): logging_obj = MagicMock() - with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()): + with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()): bridge.set_rust_chat_completions( decline=lambda **_kwargs: None, chat_completions=declining_native ) @@ -475,7 +475,7 @@ def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines(): logging_obj, calls = _recording_logging_obj() - with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()): + with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()): bridge.set_rust_chat_completions( decline=lambda **_kwargs: None, chat_completions=declining_native ) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index d28184b96a6..28e7acb8a65 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -216,13 +216,13 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge._OCR.reset() - rust_bridge._AOCR.reset() + rust_bridge._OCR.sync.reset() + rust_bridge._OCR.asynchronous.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge._OCR.reset() - rust_bridge._AOCR.reset() + rust_bridge._OCR.sync.reset() + rust_bridge._OCR.asynchronous.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -232,7 +232,7 @@ def fake_bridge(): """Enable the Rust path with an injected recording bridge (no native wheel).""" bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) return bridge @@ -241,14 +241,14 @@ def fake_async_bridge(): """Enable the async Rust path with an injected recording bridge.""" bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge._AOCR.override(bridge) + rust_bridge._OCR.asynchronous.override(bridge) return bridge def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) assert rust_bridge.load_rust_ocr() is bridge @@ -312,7 +312,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch): def test_load_rust_aocr_returns_injected_impl(): bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge._AOCR.override(bridge) + rust_bridge._OCR.asynchronous.override(bridge) assert rust_bridge.load_rust_aocr() is bridge @@ -321,8 +321,8 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) - rust_bridge._AOCR.override(async_bridge) + rust_bridge._OCR.sync.override(bridge) + rust_bridge._OCR.asynchronous.override(async_bridge) litellm.rust(False) assert rust_bridge.load_rust_ocr() is bridge @@ -341,11 +341,11 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch): bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) - rust_bridge._AOCR.override(async_bridge) + rust_bridge._OCR.sync.override(bridge) + rust_bridge._OCR.asynchronous.override(async_bridge) - rust_bridge._OCR.override(None) - rust_bridge._AOCR.override(None) + rust_bridge._OCR.sync.override(None) + rust_bridge._OCR.asynchronous.override(None) assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None @@ -392,7 +392,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) response = rust_bridge.ocr( model="mistral-ocr-latest", document=DOCUMENT, @@ -427,7 +427,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): litellm.rust(True) - rust_bridge._AOCR.override(bridge) + rust_bridge._OCR.asynchronous.override(bridge) response = await rust_bridge.aocr( model="mistral-ocr-maas", document=DOCUMENT, @@ -456,7 +456,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) response = ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -489,7 +489,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request(api_key=None, timeout=None), @@ -502,7 +502,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) def _resolver(name: str) -> str | None: raise AssertionError(f"resolver should not be called for {name}") @@ -522,7 +522,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): bridge = RecordingBridge() resolver_calls = [] litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) def _resolver(name): resolver_calls.append(name) @@ -545,7 +545,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -572,7 +572,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) def _resolver(name: str) -> str | None: return { @@ -596,7 +596,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -614,7 +614,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -635,7 +635,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): logging_obj = RecordingLogging() bridge = RecordingBridge() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -723,7 +723,7 @@ def test_ocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge._OCR.override(RaisingBridge()) + rust_bridge._OCR.sync.override(RaisingBridge()) with pytest.raises(CapturedException): litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -769,7 +769,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge._AOCR.override(RaisingAsyncBridge()) + rust_bridge._OCR.asynchronous.override(RaisingAsyncBridge()) with pytest.raises(CapturedException): await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -798,7 +798,7 @@ def test_ocr_does_not_route_to_rust_when_disabled(): """With the flag off, the bridge must not be consulted even if an impl exists.""" bridge = RecordingBridge() litellm.rust(False) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.sync.override(bridge) # The impl stays available for injection, but the disabled flag gates usage, # so ocr() never reaches the Rust path (asserted via the enabled-path test). assert bridge.calls == [] diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 74d96bda336..6e24d6def7a 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -65,8 +65,8 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: @pytest.mark.asyncio async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(responses_websocket, "_STATE", responses_websocket._RustResponsesWebSocketState()) - monkeypatch.setattr(responses_websocket, "get_native_bridge", lambda: None) + configuration.rust(True) + responses_websocket._RESPONSES_WEBSOCKET.override(None) assert ( await responses_websocket.connect( @@ -82,6 +82,7 @@ async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) async def test_enabled_bridge_connects_and_adapts_socket( monkeypatch: pytest.MonkeyPatch, ) -> None: + configuration.rust(True) responses_websocket.set_rust_responses_websocket(connection=_FakeNativeBridge) connection = await responses_websocket.connect( @@ -94,3 +95,30 @@ async def test_enabled_bridge_connects_and_adapts_socket( await connection.send("response.create") assert await connection.recv() == "response.completed" await connection.close() + + +class _FailingNativeBridge: + @classmethod + async def connect( + cls, + *, + url: str, + headers: dict[str, str], + timeout_seconds: float | None, + ) -> _FakeNativeConnection: + raise RuntimeError("connection failed") + + +@pytest.mark.asyncio +async def test_connection_failure_preserves_python_fallback() -> None: + configuration.rust(True) + responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge) + + assert ( + await responses_websocket.connect( + url="wss://example.test/responses", + headers={}, + timeout=None, + ) + is None + ) diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py index 88036a5a556..f9d7f6dbd4e 100644 --- a/tests/test_litellm/rust_bridge/test_bindings.py +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -1,3 +1,7 @@ +import json +import subprocess +import sys +from pathlib import Path from types import SimpleNamespace from typing import Final @@ -6,30 +10,111 @@ import pytest from litellm.rust_bridge import bindings -def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None: - native = SimpleNamespace(route=lambda: "native") +def test_binding_distinguishes_disable_from_reset(monkeypatch: pytest.MonkeyPatch) -> None: + native: Final = SimpleNamespace(chat_completions=lambda: "native") monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) - binding: bindings.NativeBinding[object] = bindings.NativeBinding("route", validate=lambda value: value) - - assert binding.load() is native.route + binding: Final = bindings.NativeBinding(lambda module: module.chat_completions) + assert binding.load() is native.chat_completions binding.override(None) assert binding.load() is None - - replacement = object() - binding.override(replacement) - assert binding.load() is replacement - + replacement: Final = SimpleNamespace(chat_completions=lambda: "replacement") + binding.override(replacement.chat_completions) + assert binding.load() is replacement.chat_completions binding.reset() - assert binding.load() is native.route + assert binding.load() is native.chat_completions -@pytest.mark.parametrize(("value", "expected"), ((3, 3), ("invalid", None), (None, None))) -def test_binding_validates_native_attribute( - monkeypatch: pytest.MonkeyPatch, value: object, expected: int | None -) -> None: - native: Final = SimpleNamespace(route=value) +@pytest.mark.parametrize("native", (None, SimpleNamespace(), SimpleNamespace(chat_completions=3))) +def test_missing_or_invalid_export_is_unavailable(monkeypatch: pytest.MonkeyPatch, native: object) -> None: monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) - binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None) + binding: Final = bindings.NativeBinding(lambda module: module.chat_completions) - assert binding.load() == expected + assert binding.load() is None + + +def test_selection_is_lazy_and_preserves_other_exports(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bindings, "get_native_bridge", lambda: pytest.fail("must not load during construction")) + binding: Final = bindings.NativeBinding(lambda module: module.chat_completions) + native: Final = SimpleNamespace(chat_completions=lambda: "native", achat_completions=None) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + + assert binding.load() is native.chat_completions + assert bindings.NativeBinding(lambda module: module.achat_completions).load() is None + + +@pytest.mark.parametrize("invalid", (None, str, lambda: None)) +def test_native_exception_types_reject_non_exception_classes(monkeypatch: pytest.MonkeyPatch, invalid: object) -> None: + native: Final = SimpleNamespace(RustBridgeDeclined=invalid, RustUpstreamError=RuntimeError) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + + assert bindings.native_exception_types() is None + + +@pytest.mark.parametrize( + ("expression", "expected_rule"), + ( + ("NativeBinding(lambda native: native.chat_completion)", "reportAttributeAccessIssue"), + ( + "wrong: NativeBinding[RustAchatCompletions] = NativeBinding(lambda native: native.chat_completions)", + "reportAssignmentType", + ), + ( + 'EndpointBinding.native(route="chat", select=lambda native: native.chat_completion, enabled=always_enabled)', + "reportAttributeAccessIssue", + ), + ("NativeBinding(lambda native: native.ocrr)", "reportAttributeAccessIssue"), + ( + "wrong: NativeBinding[RustAmessages] = NativeBinding(lambda native: native.messages)", + "reportAssignmentType", + ), + ( + "wrong: NativeBinding[RustAocr] = NativeBinding(lambda native: native.ocr)", + "reportAssignmentType", + ), + ( + "wrong: NativeBinding[RustAtranscription] = NativeBinding(lambda native: native.transcription)", + "reportAssignmentType", + ), + ), +) +def test_selectors_are_checked_by_type_checker(tmp_path: Path, expression: str, expected_rule: str) -> None: + source: Final = tmp_path / "binding_contract.py" + source.write_text( + "from typing_extensions import assert_type\n" + "from litellm.rust_bridge.bindings import NativeBinding\n" + "from litellm.rust_bridge.protocols import RustChatCompletions, RustAchatCompletions, " + "RustMessages, RustAmessages, RustOcr, RustAocr, RustTranscription, RustAtranscription\n" + "from litellm.rust_bridge.runtime import EndpointBinding, EndpointDispatch, always_enabled\n" + "binding = NativeBinding(lambda native: native.chat_completions)\n" + "assert_type(binding, NativeBinding[RustChatCompletions])\n" + 'bridge = EndpointBinding.native(route="chat", select=lambda native: native.chat_completions, enabled=always_enabled)\n' + "assert_type(bridge, EndpointBinding[RustChatCompletions])\n" + 'endpoint = EndpointDispatch.native(route="chat", sync=lambda native: native.chat_completions, ' + "asynchronous=lambda native: native.achat_completions, enabled=always_enabled)\n" + "assert_type(endpoint, EndpointDispatch[RustChatCompletions, RustAchatCompletions])\n" + "assert_type(NativeBinding(lambda native: native.messages), NativeBinding[RustMessages])\n" + "assert_type(NativeBinding(lambda native: native.ocr), NativeBinding[RustOcr])\n" + "assert_type(NativeBinding(lambda native: native.transcription), NativeBinding[RustTranscription])\n" + + expression + + "\n" + ) + config: Final = tmp_path / "pyrightconfig.json" + config.write_text( + json.dumps( + { + "include": [str(source)], + "extraPaths": [str(Path(__file__).resolve().parents[3])], + "typeCheckingMode": "basic", + } + ) + ) + result: Final = subprocess.run( + [sys.executable, "-m", "basedpyright", "--project", str(config), "--outputjson"], + capture_output=True, + text=True, + check=False, + ) + diagnostics: Final = json.loads(result.stdout)["generalDiagnostics"] + assert result.returncode == 1, result.stdout + result.stderr + assert [(item["rule"], item["range"]["start"]["line"]) for item in diagnostics] == [(expected_rule, 13)] diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index b2fd2e6dcc0..49cb1822ccc 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -10,7 +10,7 @@ from __future__ import annotations import pytest import litellm -from litellm.rust_bridge import configuration +from litellm.rust_bridge import bindings, configuration from litellm.rust_bridge import chat_completions as bridge from litellm.types.utils import ModelResponse @@ -54,7 +54,7 @@ class _FakeNative: def _fake_native_bridge(monkeypatch): """Expose the bridge's exception classes without the compiled extension.""" - monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) def _hide_native_bridge(monkeypatch): @@ -63,7 +63,7 @@ def _hide_native_bridge(monkeypatch): There is no injection seam for "the .so is absent", so the loader itself is replaced; every other case here uses `set_rust_chat_completions`. """ - monkeypatch.setattr(bridge, "get_native_bridge", lambda: None) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) @pytest.fixture(autouse=True) @@ -393,3 +393,22 @@ class TestFailureClassification: result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result == "python" + + +@pytest.mark.asyncio +async def test_missing_native_exception_types_preserves_python_fallback(monkeypatch): + _hide_native_bridge(monkeypatch) + bridge.set_rust_chat_completions( + chat_completions=_RecordingCall(error=RuntimeError("connection failed")), + achat_completions=_RecordingAsyncCall(error=RuntimeError("connection failed")), + ) + + assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None + + async def fallback(): + return "python" + + assert ( + await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) + == "python" + ) diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index b0fa510069b..fff0f18c312 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -1,6 +1,8 @@ from __future__ import annotations +from dataclasses import dataclass from types import SimpleNamespace +from typing import Final import pytest @@ -18,78 +20,412 @@ class RustUpstreamError(Exception): @pytest.fixture(autouse=True) def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None: - native = SimpleNamespace( - RustBridgeDeclined=RustBridgeDeclined, - RustUpstreamError=RustUpstreamError, + monkeypatch.setattr( + bindings, + "get_native_bridge", + lambda: SimpleNamespace( + RustBridgeDeclined=RustBridgeDeclined, + RustUpstreamError=RustUpstreamError, + ), ) - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) def context() -> runtime.BridgeErrorContext: - return runtime.BridgeErrorContext(route="messages", provider="anthropic", model="model") + return runtime.BridgeErrorContext(provider="anthropic", model="model") -def test_invoke_tags_native_decline_before_running_fallback() -> None: - calls: list[str] = [] +def enabled() -> bool: + return True - def decline() -> object: - calls.append("rust") - raise RustBridgeDeclined("unsupported") - value = runtime.invoke( - native_call=decline, - fallback=lambda: calls.append("python") or "fallback", +@dataclass(frozen=True, slots=True) +class FallbackCase: + process_enabled: bool | None = None + eligible: bool = True + binding_available: bool = True + declined: bool = False + expected_events: tuple[str, ...] = () + + +FALLBACK_CASES: Final = ( + pytest.param( + FallbackCase(process_enabled=False, expected_events=("python",)), + id="process-disabled", + ), + pytest.param( + FallbackCase(eligible=False, expected_events=("python",)), + id="request-ineligible", + ), + pytest.param( + FallbackCase(binding_available=False, expected_events=("load", "python")), + id="bridge-unavailable", + ), + pytest.param( + FallbackCase(declined=True, expected_events=("load", "prepare", "rust", "python")), + id="bridge-declined", + ), +) + + +@pytest.mark.parametrize("case", FALLBACK_CASES) +def test_invoke_falls_back_only_before_provider_success(case: FallbackCase) -> None: + events: list[str] = [] + + def load() -> object | None: + events.append("load") + return object() if case.binding_available else None + + def call(_binding: object, _request: object) -> int: + events.append("rust") + if case.declined: + raise RustBridgeDeclined("unsupported") + return 3 + + bridge: Final = runtime.EndpointBinding( + route="messages", load=load, enabled=lambda: case.process_enabled is not False + ) + result: Final = bridge.invoke( + prepare=lambda: events.append("prepare"), + call=call, + fallback=lambda: events.append("python") or "fallback", adapt=str, - mode=runtime.FallbackMode.PYTHON, - context=context(), + error_context=context(), + eligible=case.eligible, ) - assert value == "fallback" - assert calls == ["rust", "python"] - - -def test_invoke_translates_upstream_without_fallback() -> None: - def fail() -> object: - raise RustUpstreamError(429, "rate limited") - - with pytest.raises(APIError, match="rate limited") as caught: - runtime.invoke( - native_call=fail, - fallback=lambda: pytest.fail("fallback must not run"), - adapt=str, - mode=runtime.FallbackMode.PYTHON, - context=context(), - ) - - assert caught.value.status_code == 429 + assert result == "fallback" + assert tuple(events) == case.expected_events @pytest.mark.asyncio -async def test_ainvoke_handles_native_success() -> None: - async def native() -> int: +@pytest.mark.parametrize("case", FALLBACK_CASES) +async def test_ainvoke_matches_sync_fallback_contract(case: FallbackCase) -> None: + events: list[str] = [] + + def load() -> object | None: + events.append("load") + return object() if case.binding_available else None + + async def call(_binding: object, _request: object) -> int: + events.append("rust") + if case.declined: + raise RustBridgeDeclined("unsupported") return 3 + async def fallback() -> str: + events.append("python") + return "fallback" + + bridge: Final = runtime.EndpointBinding( + route="messages", load=load, enabled=lambda: case.process_enabled is not False + ) + result: Final = await bridge.ainvoke( + prepare=lambda: events.append("prepare"), + call=call, + fallback=fallback, + adapt=str, + error_context=context(), + eligible=case.eligible, + ) + + assert result == "fallback" + assert tuple(events) == case.expected_events + + +def test_invoke_adapts_native_success_without_fallback() -> None: + bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) + + result: Final = bridge.invoke( + prepare=lambda: 3, + call=lambda _binding, request: request * 2, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=lambda value: f"adapted-{value}", + error_context=context(), + ) + + assert result == "adapted-6" + + +@pytest.mark.asyncio +async def test_ainvoke_adapts_native_success_without_fallback() -> None: + async def call(_binding: object, request: int) -> int: + return request * 2 + async def fallback() -> str: pytest.fail("fallback must not run") - assert ( - await runtime.ainvoke( - native_call=native, - fallback=fallback, - adapt=str, - mode=runtime.FallbackMode.PYTHON, - context=context(), - ) - == "3" + bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) + result: Final = await bridge.ainvoke( + prepare=lambda: 3, + call=call, + fallback=fallback, + adapt=lambda value: f"adapted-{value}", + error_context=context(), ) + assert result == "adapted-6" -def test_required_mode_rejects_unavailable_bridge() -> None: - with pytest.raises(RuntimeError, match="is unavailable"): - runtime.invoke( - native_call=None, + +@pytest.mark.parametrize( + ("error", "expected_type", "expected_status", "expected_message"), + ( + pytest.param(RustUpstreamError(401, "unauthorized"), APIError, 401, "unauthorized", id="auth"), + pytest.param(RustUpstreamError(429, "rate limited"), APIError, 429, "rate limited", id="rate-limit"), + pytest.param(RustUpstreamError(500, "failed"), APIError, 500, "failed", id="server-error"), + pytest.param(RustUpstreamError(0, "connection reset"), APIError, 500, "connection reset", id="transport"), + pytest.param(RustUpstreamError(403, "forbidden"), APIError, 403, "forbidden", id="other-status"), + ), +) +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_upstream_failure_maps_to_api_error_without_fallback( + asynchronous: bool, + error: RustUpstreamError, + expected_type: type[BaseException], + expected_status: int, + expected_message: str, +) -> None: + def fail(_binding: object, _request: object) -> object: + raise error + + async def afail(binding: object, request: object) -> object: + return fail(binding, request) + + async def fallback() -> str: + pytest.fail("fallback must not run") + + bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) + + async def invoke() -> None: + if asynchronous: + await bridge.ainvoke( + prepare=lambda: None, call=afail, fallback=fallback, adapt=str, error_context=context() + ) + else: + bridge.invoke( + prepare=lambda: None, + call=fail, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + error_context=context(), + ) + + with pytest.raises(expected_type, match=expected_message) as caught: + await invoke() + + assert type(caught.value) is expected_type + assert caught.value.status_code == expected_status + assert caught.value.llm_provider == "anthropic" + assert caught.value.model == "model" + assert caught.value.__cause__ is error + + +@pytest.mark.asyncio +async def test_async_upstream_failure_maps_to_api_error_without_fallback() -> None: + async def fail(_binding: object, _request: object) -> object: + raise RustUpstreamError(503, "overloaded") + + async def fallback() -> object: + pytest.fail("fallback must not run") + + bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) + + with pytest.raises(APIError, match="overloaded") as caught: + await bridge.ainvoke(prepare=lambda: None, call=fail, fallback=fallback, adapt=str, error_context=context()) + + assert caught.value.status_code == 503 + + +def test_unknown_failure_is_preserved_without_fallback() -> None: + error: Final = RuntimeError("unknown") + bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) + + with pytest.raises(RuntimeError, match="unknown") as caught: + bridge.invoke( + prepare=lambda: None, + call=lambda _binding, _request: (_ for _ in ()).throw(error), fallback=lambda: pytest.fail("fallback must not run"), adapt=str, - mode=runtime.FallbackMode.RUST_REQUIRED, - context=context(), + error_context=context(), ) + + assert caught.value is error + + +@pytest.mark.parametrize( + ("process_enabled", "binding_available", "declined", "expected_message"), + ( + pytest.param(False, True, False, "native messages endpoint is disabled", id="disabled"), + pytest.param(None, False, False, "native messages endpoint is unavailable", id="unavailable"), + pytest.param( + None, + True, + True, + "native messages endpoint declined the request: unsupported", + id="declined", + ), + ), +) +def test_require_explains_why_rust_did_not_handle_request( + process_enabled: bool | None, + binding_available: bool, + declined: bool, + expected_message: str, +) -> None: + def call(_binding: object, _request: object) -> object: + if declined: + raise RustBridgeDeclined("unsupported") + return object() + + bridge: Final = runtime.EndpointBinding( + route="messages", + load=object if binding_available else lambda: None, + enabled=lambda: process_enabled is not False, + ) + + with pytest.raises(RuntimeError, match=f"^{expected_message}$"): + bridge.require( + prepare=lambda: None, + call=call, + adapt=str, + error_context=context(), + ) + + +@pytest.mark.parametrize( + ("state", "expected", "expected_events"), + ( + pytest.param("disabled", False, (), id="disabled"), + pytest.param("ineligible", False, (), id="ineligible"), + pytest.param("unavailable", False, ("load",), id="unavailable"), + pytest.param("available", True, ("load",), id="available"), + ), +) +def test_can_attempt_only_enabled_available_requests( + state: str, + expected: bool, + expected_events: tuple[str, ...], +) -> None: + events: list[str] = [] + + def load() -> object | None: + events.append("load") + return None if state == "unavailable" else object() + + bridge: Final = runtime.EndpointBinding(route="messages", load=load, enabled=lambda: state != "disabled") + + assert ( + bridge.can_attempt( + eligible=state != "ineligible", + ) + is expected + ) + assert tuple(events) == expected_events + + +def test_native_endpoint_applies_partial_overrides_and_reset(monkeypatch: pytest.MonkeyPatch) -> None: + def native_sync() -> str: + return "native" + + async def native_async() -> str: + return "native async" + + def replacement_sync() -> str: + return "replacement" + + monkeypatch.setattr( + bindings, + "get_native_bridge", + lambda: SimpleNamespace(chat_completions=native_sync, achat_completions=native_async), + ) + endpoint: Final[runtime.EndpointDispatch[object, object]] = runtime.EndpointDispatch.native( + route="test", + sync=lambda native: native.chat_completions, + asynchronous=lambda native: native.achat_completions, + enabled=enabled, + ) + + assert endpoint.sync.load() is native_sync + assert endpoint.asynchronous.load() is native_async + endpoint.override(sync=replacement_sync) + assert endpoint.sync.load() is replacement_sync + assert endpoint.asynchronous.load() is native_async + endpoint.override(asynchronous=None) + assert endpoint.sync.load() is replacement_sync + assert endpoint.asynchronous.load() is None + endpoint.reset() + assert endpoint.sync.load() is native_sync + assert endpoint.asynchronous.load() is native_async + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_response_adaptation_failure_never_authorizes_fallback(asynchronous: bool) -> None: + def adapt(value: str) -> str: + assert value == "provider response" + raise RustBridgeDeclined("adapter failed after provider response") + + async def native(binding: object, request: object) -> str: + return "provider response" + + async def fallback() -> str: + pytest.fail("a received response must not be retried") + + bridge = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) + + async def invoke() -> None: + if asynchronous: + await bridge.ainvoke( + prepare=lambda: None, call=native, fallback=fallback, adapt=adapt, error_context=context() + ) + else: + bridge.invoke( + prepare=lambda: None, + call=lambda binding, request: "provider response", + fallback=lambda: pytest.fail("a received response must not be retried"), + adapt=adapt, + error_context=context(), + ) + + with pytest.raises(RustBridgeDeclined, match="adapter failed"): + await invoke() + + +@pytest.mark.parametrize("error", (RustBridgeDeclined("unsupported"), RustUpstreamError(429, "rate limited"))) +def test_propagate_policy_preserves_native_errors(error: Exception) -> None: + def fail(_binding: object, _request: object) -> object: + raise error + + bridge: Final = runtime.EndpointBinding( + route="ocr", load=object, enabled=enabled, error_policy=runtime.NativeErrorPolicy.PROPAGATE + ) + + with pytest.raises(type(error)) as caught: + bridge.invoke( + prepare=lambda: None, + call=fail, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + error_context=context(), + ) + + assert caught.value is error + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", (RustBridgeDeclined("unsupported"), RustUpstreamError(429, "rate limited"))) +async def test_async_propagate_policy_preserves_native_errors(error: Exception) -> None: + async def fail(_binding: object, _request: object) -> object: + raise error + + async def fallback() -> str: + pytest.fail("fallback must not run") + + bridge: Final = runtime.EndpointBinding( + route="messages", load=object, enabled=enabled, error_policy=runtime.NativeErrorPolicy.PROPAGATE + ) + + with pytest.raises(type(error)) as caught: + await bridge.ainvoke(prepare=lambda: None, call=fail, fallback=fallback, adapt=str, error_context=context()) + + assert caught.value is error diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index 112464bda22..1b0fcfacd1d 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -77,7 +77,7 @@ async def test_enabled_async_bridge() -> None: def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None: rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) - monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", lambda: None) + monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None) assert rust_bridge.load_rust_transcription() is None assert rust_bridge.load_rust_atranscription() is None diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e7186dfe186..fb968f6ffcd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1035 + "limit": 1031 }, "LIT007": { "limit": 0