diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 2f7f22032be..1f3308459ea 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 13427 + "limit": 13426 }, "reportArgumentType": { "limit": 2194 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 3369 + "limit": 3368 }, "reportFunctionMemberAccess": { "limit": 7 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 180 + "limit": 178 }, "reportTypedDictNotRequiredAccess": { "limit": 22 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44136 + "limit": 44025 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38269 }, "reportUnknownParameterType": { "limit": 19584 }, "reportUnknownVariableType": { - "limit": 29814 + "limit": 29813 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index c82be07a5c5..318bb043270 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -27,6 +27,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts +from litellm.rust_bridge.dispatch import adispatch, dispatch from litellm.types.llms.anthropic import ( ContentBlockDelta, ContentBlockStart, @@ -453,7 +454,25 @@ class AnthropicChatCompletion(BaseLLM): timeout=timeout, ) - return rust_chat_completions_bridge.achat_completions_or_fallback( + return adispatch( + native=lambda: rust_chat_completions_bridge.achat_completions( + model=model, + messages=messages, + optional_params=rust_optional_params, + model_response=model_response, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + on_response=log_rust_post_call, + ), + python=python_fallback, + route="chat_completions", + errors=rust_chat_completions_bridge.error_handling(custom_llm_provider or "", model), + ) + rust_response: Final = dispatch( + native=lambda: rust_chat_completions_bridge.chat_completions( model=model, messages=messages, optional_params=rust_optional_params, @@ -464,19 +483,10 @@ class AnthropicChatCompletion(BaseLLM): extra_headers=headers, timeout=timeout, on_response=log_rust_post_call, - python_fallback=python_fallback, - ) - rust_response: Final = rust_chat_completions_bridge.chat_completions( - model=model, - messages=messages, - optional_params=rust_optional_params, - model_response=model_response, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=headers, - timeout=timeout, - on_response=log_rust_post_call, + ), + python=lambda: None, + route="chat_completions", + errors=rust_chat_completions_bridge.error_handling(custom_llm_provider or "", model), ) if rust_response is not None: return rust_response diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index b1f8c957ff4..a1452602523 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -1,13 +1,22 @@ import base64 -from typing import Final +from typing import Final, NoReturn import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.rust_bridge import transcription as rust_transcription_bridge +from litellm.rust_bridge.dispatch import PROPAGATE, adispatch, dispatch from litellm.types.utils import FileTypes, TranscriptionResponse +def _unavailable() -> NoReturn: + raise RuntimeError("Rust audio transcription bridge is unavailable") + + +async def _aunavailable() -> NoReturn: + _unavailable() + + class BedrockAudioTranscriptionRustDispatch: @staticmethod def _audio_payload(audio_file: FileTypes) -> dict[str, object]: @@ -43,18 +52,21 @@ class BedrockAudioTranscriptionRustDispatch: optional_params: dict[str, object], timeout: float | httpx.Timeout | None, ) -> TranscriptionResponse: - rust_response: Final = rust_transcription_bridge.transcription( - model=model, - audio=self._audio_payload(audio_file), - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout=timeout, + rust_response: Final = dispatch( + native=lambda: rust_transcription_bridge.transcription( + model=model, + audio=self._audio_payload(audio_file), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ), + python=_unavailable, + route="audio transcription", + errors=PROPAGATE, ) - if rust_response is None: - raise RuntimeError("Rust audio transcription bridge is unavailable") return TranscriptionResponse(**rust_response) async def async_audio_transcriptions( @@ -69,16 +81,19 @@ class BedrockAudioTranscriptionRustDispatch: optional_params: dict[str, object], timeout: float | httpx.Timeout | None, ) -> TranscriptionResponse: - rust_response: Final = await rust_transcription_bridge.atranscription( - model=model, - audio=self._audio_payload(audio_file), - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout=timeout, + rust_response: Final = await adispatch( + native=lambda: rust_transcription_bridge.atranscription( + model=model, + audio=self._audio_payload(audio_file), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ), + python=_aunavailable, + route="audio transcription", + errors=PROPAGATE, ) - if rust_response is None: - raise RuntimeError("Rust audio transcription bridge is unavailable") return TranscriptionResponse(**rust_response) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index a75124325ae..e3ee89a2455 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -18,6 +18,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts +from litellm.rust_bridge.dispatch import adispatch, dispatch from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper @@ -423,18 +424,20 @@ class BedrockConverseLLM(BaseAWSLLM): additional_args=rust_logging_args, ) if acompletion: - return rust_chat_completions_bridge.achat_completions_or_fallback( - model=model, - messages=messages, - optional_params=rust_optional_params, - model_response=model_response, - api_key=api_key, - api_base=proxy_endpoint_url, - custom_llm_provider="bedrock", - extra_headers=headers, - timeout=timeout, - on_response=log_rust_post_call, - python_fallback=lambda: self.async_completion( + return adispatch( + native=lambda: rust_chat_completions_bridge.achat_completions( + model=model, + messages=messages, + optional_params=rust_optional_params, + model_response=model_response, + api_key=api_key, + api_base=proxy_endpoint_url, + custom_llm_provider="bedrock", + extra_headers=headers, + timeout=timeout, + on_response=log_rust_post_call, + ), + python=lambda: self.async_completion( model=model, messages=messages, api_base=proxy_endpoint_url, @@ -452,18 +455,25 @@ class BedrockConverseLLM(BaseAWSLLM): api_key=api_key, skip_pre_call_logging=True, ), + route="chat_completions", + errors=rust_chat_completions_bridge.error_handling("bedrock", model), ) - rust_response: Final = rust_chat_completions_bridge.chat_completions( - model=model, - messages=messages, - optional_params=rust_optional_params, - model_response=model_response, - api_key=api_key, - api_base=proxy_endpoint_url, - custom_llm_provider="bedrock", - extra_headers=headers, - timeout=timeout, - on_response=log_rust_post_call, + rust_response: Final = dispatch( + native=lambda: rust_chat_completions_bridge.chat_completions( + model=model, + messages=messages, + optional_params=rust_optional_params, + model_response=model_response, + api_key=api_key, + api_base=proxy_endpoint_url, + custom_llm_provider="bedrock", + extra_headers=headers, + timeout=timeout, + on_response=log_rust_post_call, + ), + python=lambda: None, + route="chat_completions", + errors=rust_chat_completions_bridge.error_handling("bedrock", model), ) if rust_response is not None: return rust_response diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 552b1549db8..e361581ec2a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2408,8 +2408,10 @@ class BaseLLMHTTPHandler: from litellm.rust_bridge import messages as rust_messages_bridge upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"} - try: - rust_response: Final = await rust_messages_bridge.amessages( + from litellm.rust_bridge.dispatch import PYTHON_ON_ERROR, adispatch, async_none + + rust_response: Final = await adispatch( + native=lambda: rust_messages_bridge.amessages( model=model, body=upstream_body, api_key=api_key, @@ -2417,13 +2419,11 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, extra_headers=headers, timeout=timeout, - ) - except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path - verbose_logger.debug( - "Rust Anthropic messages bridge raised %s; falling back to Python path", - type(rust_error).__name__, - ) - return None + ), + python=async_none, + route="messages", + errors=PYTHON_ON_ERROR, + ) if rust_response is None: return None @@ -6511,11 +6511,17 @@ class BaseLLMHTTPHandler: async def _backend_connection(): if _rust_responses_websocket_enabled(custom_llm_provider): from litellm.rust_bridge import responses_websocket as rust_responses_websocket + from litellm.rust_bridge.dispatch import PYTHON_ON_ERROR, adispatch, async_none - rust_backend: Final = await rust_responses_websocket.connect( - url=ws_url, - headers={str(key): str(value) for key, value in headers.items()}, - timeout=timeout, + rust_backend: Final = await adispatch( + native=lambda: rust_responses_websocket.connect( + url=ws_url, + headers={str(key): str(value) for key, value in headers.items()}, + timeout=timeout, + ), + python=async_none, + route="responses_websocket", + errors=PYTHON_ON_ERROR, ) if rust_backend is not None: yield rust_backend diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 36e3b8838e2..03ed110ba34 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -7,8 +7,7 @@ import base64 import mimetypes import os import re -from collections.abc import Callable, Coroutine, Mapping -from dataclasses import dataclass +from collections.abc import Coroutine, Mapping from io import IOBase from typing import Any, Final, cast @@ -18,18 +17,15 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.azure_ai.ocr.common_utils import ( - is_azure_document_intelligence_model, -) +from litellm.llms.azure_ai.ocr.common_utils import is_azure_document_intelligence_model from litellm.llms.base_llm.ocr.transformation import ( OCR_REQUEST_FORMAT_PARAM, - BaseOCRConfig, OCRResponse, parse_ocr_request_format, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge -from litellm.rust_bridge.timeouts import timeout_to_seconds +from litellm.rust_bridge.dispatch import PROPAGATE, adispatch, dispatch from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -38,36 +34,6 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# -@dataclass -class _PreparedOCRRequest: - model: str - document: dict[str, Any] - api_key: str | None - api_base: str | None - custom_llm_provider: str - extra_headers: dict[str, object] | None - provider_config: BaseOCRConfig - optional_params: dict[str, object] - litellm_params: dict[str, object] - effective_timeout: float | httpx.Timeout - litellm_logging_obj: LiteLLMLoggingObj - - -@dataclass -class _PreparedRustOCRCall: - api_key: str | None - api_base: str | None - headers: dict[str, object] - optional_params: dict[str, object] - - -_RUST_OCR_PROVIDERS: Final = { - "mistral", - "azure_ai", - "vertex_ai", -} - - def _prepare_ocr_request( model: str, document: Mapping[str, object], @@ -77,7 +43,7 @@ def _prepare_ocr_request( custom_llm_provider: str | None, extra_headers: dict[str, object] | None, kwargs: dict[str, object], -) -> _PreparedOCRRequest: +) -> rust_ocr_bridge.PreparedOCRRequest: litellm_logging_obj: Final = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) litellm_call_id: Final = cast(str | None, kwargs.get("litellm_call_id", None)) @@ -174,7 +140,7 @@ def _prepare_ocr_request( custom_llm_provider=custom_llm_provider, ) - return _PreparedOCRRequest( + return rust_ocr_bridge.PreparedOCRRequest( model=model, document=document, api_key=api_key, @@ -189,154 +155,6 @@ def _prepare_ocr_request( ) -def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: - if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native": - return False - if not prepared_request.provider_config.supports_rust_bridge(): - return False - return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS - - -def _rust_bridge_optional_params( - prepared_request: _PreparedOCRRequest, - resolve_secret: Callable[[str], str | None], -) -> dict[str, object]: - optional_params: Final = dict(prepared_request.optional_params) - if prepared_request.custom_llm_provider == "vertex_ai": - vertex_project: Final = ( - prepared_request.litellm_params.get("vertex_project") - or prepared_request.litellm_params.get("vertex_ai_project") - or litellm.vertex_project - or resolve_secret("VERTEXAI_PROJECT") - ) - vertex_location: Final = ( - prepared_request.litellm_params.get("vertex_location") - or prepared_request.litellm_params.get("vertex_ai_location") - or litellm.vertex_location - or resolve_secret("VERTEXAI_LOCATION") - or resolve_secret("VERTEX_LOCATION") - ) - if vertex_project is not None: - optional_params["vertex_project"] = vertex_project - if vertex_location is not None: - optional_params["vertex_location"] = vertex_location - return optional_params - - -def _rust_bridge_api_base( - prepared_request: _PreparedOCRRequest, - resolve_secret: Callable[[str], str | None], -) -> str | None: - if prepared_request.api_base is not None: - return prepared_request.api_base - if prepared_request.custom_llm_provider == "azure_ai": - if is_azure_document_intelligence_model(prepared_request.model): - return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") - return resolve_secret("AZURE_AI_API_BASE") - return None - - -def _prepare_rust_ocr_call( - prepared_request: _PreparedOCRRequest, - resolve_api_key: Callable[[str], str | None], -) -> _PreparedRustOCRCall: - provider_config: Final = prepared_request.provider_config - api_key_env_var: Final = provider_config.get_api_key_env_var() - resolved_api_key: Final = prepared_request.api_key or ( - resolve_api_key(api_key_env_var) if api_key_env_var is not None else None - ) - resolved_headers: Final = provider_config.validate_environment( - headers=prepared_request.extra_headers or {}, - model=prepared_request.model, - api_key=resolved_api_key, - api_base=prepared_request.api_base, - litellm_params=prepared_request.litellm_params, - ) - resolved_complete_url: Final = provider_config.get_complete_url( - api_base=prepared_request.api_base, - model=prepared_request.model, - optional_params=prepared_request.optional_params, - litellm_params=prepared_request.litellm_params, - ) - rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key) - rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key) - prepared_request.litellm_logging_obj.pre_call( - input="OCR document processing", - api_key=resolved_api_key, - additional_args={ - "complete_input_dict": { - "model": prepared_request.model, - "document": prepared_request.document, - **rust_optional_params, - }, - "api_base": resolved_complete_url, - "headers": resolved_headers, - }, - ) - return _PreparedRustOCRCall( - api_key=resolved_api_key, - api_base=rust_api_base, - headers=cast(dict[str, object], resolved_headers), - optional_params=rust_optional_params, - ) - - -def _run_rust_ocr( - prepared_request: _PreparedOCRRequest, - resolve_api_key: Callable[[str], str | None], - fallback: Callable[[], OCRResponse | Coroutine[object, object, OCRResponse]], -) -> OCRResponse | Coroutine[object, object, OCRResponse]: - return rust_ocr_bridge.dispatch_ocr( - prepare=lambda: _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ), - call=lambda native, prepared: native( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), - ), - fallback=fallback, - adapt=OCRResponse.model_validate, - model=prepared_request.model, - provider=prepared_request.custom_llm_provider, - eligible=_rust_ocr_supported(prepared_request), - ) - - -async def _run_rust_aocr( - prepared_request: _PreparedOCRRequest, - resolve_api_key: Callable[[str], str | None], - fallback: Callable[[], Coroutine[object, object, OCRResponse]], -) -> OCRResponse: - return await rust_ocr_bridge.adispatch_ocr( - prepare=lambda: _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ), - call=lambda native, prepared: native( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), - ), - fallback=fallback, - adapt=OCRResponse.model_validate, - model=prepared_request.model, - provider=prepared_request.custom_llm_provider, - eligible=_rust_ocr_supported(prepared_request), - ) - - @client async def aocr( model: str, @@ -453,10 +271,11 @@ async def aocr( raise ValueError(f"Got an unexpected None response from the OCR API: {response}") return response - return await _run_rust_aocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - fallback=python_fallback, + return await adispatch( + native=lambda: rust_ocr_bridge.aattempt_ocr(prepared_request=prepared, resolve_api_key=get_secret_str), + python=python_fallback, + route="ocr", + errors=PROPAGATE, ) except Exception as e: raise litellm.exception_type( @@ -714,10 +533,11 @@ def ocr( litellm_params=prepared.litellm_params, ) - return _run_rust_ocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - fallback=python_fallback, + return dispatch( + native=lambda: rust_ocr_bridge.attempt_ocr(prepared_request=prepared, resolve_api_key=get_secret_str), + python=python_fallback, + route="ocr", + errors=PROPAGATE, ) except Exception as e: raise litellm.exception_type( diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py index a170ec4edeb..ab7b7c296b9 100644 --- a/litellm/rust_bridge/bindings.py +++ b/litellm/rust_bridge/bindings.py @@ -67,9 +67,11 @@ def _exception_class(value: object) -> type[BaseException] | None: return None -def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None: - declined: Final = _exception_class(_DECLINED.load()) +def native_upstream_types() -> tuple[type[BaseException], ...]: upstream: Final = _exception_class(_UPSTREAM.load()) - if declined is None or upstream is None: - return None - return declined, upstream + return () if upstream is None else (upstream,) + + +def native_declined_types() -> tuple[type[BaseException], ...]: + declined: Final = _exception_class(_DECLINED.load()) + return () if declined is None else (declined,) diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index ded533c5ad0..24e97993a1f 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -4,16 +4,12 @@ The Rust core owns the conversation translation, the provider call, and the response normalization for the subset of `/chat/completions` requests it accepts. This module only marshals inputs and hands the normalized result to LiteLLM's existing `ModelResponse` builder. - -``None`` means the provider was never called, so the caller is free to serve the -request on the Python path. A failure after the call was issued raises instead: -retrying it there would bill the customer for the same work twice. """ from __future__ import annotations import json -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol import httpx @@ -24,19 +20,15 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo convert_to_model_response_object, ) from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned -from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged from litellm.rust_bridge.configuration import rust_enabled +from litellm.rust_bridge.dispatch import APIErrorMapping, ErrorAction, ErrorHandling from litellm.rust_bridge.protocols import ( RustAchatCompletions, RustChatCompletions, RustChatCompletionsDecline, ) -from litellm.rust_bridge.runtime import ( - BridgeErrorContext, - EndpointBinding, - EndpointDispatch, - async_none, -) +from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.utils import ModelResponse @@ -94,16 +86,10 @@ def response_logger( return log -_CHAT: Final[EndpointDispatch[RustChatCompletions, RustAchatCompletions]] = EndpointDispatch.native( - route="chat_completions", - sync=lambda native: native.chat_completions, - asynchronous=lambda native: native.achat_completions, - enabled=rust_enabled, -) -_CHAT_PREFLIGHT: Final[EndpointBinding[RustChatCompletionsDecline]] = EndpointBinding.native( - route="chat_completions", - select=lambda native: native.chat_completions_decline, - enabled=rust_enabled, +_CHAT: Final[NativeBinding[RustChatCompletions]] = NativeBinding(lambda native: native.chat_completions) +_ACHAT: Final[NativeBinding[RustAchatCompletions]] = NativeBinding(lambda native: native.achat_completions) +_CHAT_PREFLIGHT: Final[NativeBinding[RustChatCompletionsDecline]] = NativeBinding( + lambda native: native.chat_completions_decline ) @@ -117,14 +103,14 @@ def set_rust_chat_completions( patching module attributes.""" if not isinstance(chat_completions, Unchanged): if chat_completions is None: - _CHAT.sync.reset() + _CHAT.reset() else: - _CHAT.sync.override(chat_completions) + _CHAT.override(chat_completions) if not isinstance(achat_completions, Unchanged): if achat_completions is None: - _CHAT.asynchronous.reset() + _ACHAT.reset() else: - _CHAT.asynchronous.override(achat_completions) + _ACHAT.override(achat_completions) if not isinstance(decline, Unchanged): if decline is None: _CHAT_PREFLIGHT.reset() @@ -193,14 +179,24 @@ def rust_chat_completions_accepts( if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") return False - return _CHAT_PREFLIGHT.accepts( - check=lambda decline: decline( + if not rust_enabled(): + return False + decline: Final = _CHAT_PREFLIGHT.load() + if decline is None: + return False + try: + reason: Final = decline( model=model, messages=messages, optional_params=optional_params, custom_llm_provider=custom_llm_provider, - ), - ) + ) + except Exception as error: # noqa: BLE001 # capability checks perform no provider I/O + verbose_logger.debug("Native chat acceptance check failed: %s", error) + return False + if reason is not None: + verbose_logger.debug("Native chat request is ineligible: %s", reason) + return reason is None def _build_model_response( @@ -229,14 +225,13 @@ def chat_completions( extra_headers: Mapping[str, object] | None, timeout: float | httpx.Timeout | None, on_response: ResponseObserver, -) -> ModelResponse | None: +) -> DispatchResult[ModelResponse]: def adapt(rust_response: Mapping[str, object]) -> ModelResponse: on_response(rust_response) return _build_model_response(rust_response, model_response) - return _CHAT.invoke( - prepare=lambda: timeout_to_seconds(timeout), - call=lambda rust_chat_completions, timeout_seconds: rust_chat_completions( + def call(native: RustChatCompletions, timeout_seconds: float | None) -> Mapping[str, object]: + return native( model=model, messages=messages, optional_params=optional_params, @@ -245,10 +240,15 @@ def chat_completions( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, timeout_seconds=timeout_seconds, - ), - fallback=lambda: None, + ) + + return attempt( + load=_CHAT.load, + enabled=rust_enabled(), + eligible=True, + prepare=lambda: timeout_to_seconds(timeout), + call=call, adapt=adapt, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -264,14 +264,13 @@ async def achat_completions( extra_headers: Mapping[str, object] | None, timeout: float | httpx.Timeout | None, on_response: ResponseObserver, -) -> ModelResponse | None: +) -> DispatchResult[ModelResponse]: def adapt(rust_response: Mapping[str, object]) -> ModelResponse: on_response(rust_response) return _build_model_response(rust_response, model_response) - return await _CHAT.ainvoke( - prepare=lambda: timeout_to_seconds(timeout), - call=lambda rust_achat_completions, timeout_seconds: rust_achat_completions( + async def call(native: RustAchatCompletions, timeout_seconds: float | None) -> Mapping[str, object]: + return await native( model=model, messages=messages, optional_params=optional_params, @@ -280,53 +279,21 @@ async def achat_completions( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, timeout_seconds=timeout_seconds, - ), - fallback=async_none, - adapt=adapt, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), - ) + ) - -async def achat_completions_or_fallback( - *, - model: str, - messages: Sequence[object], - optional_params: Mapping[str, object], - model_response: ModelResponse, - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: Mapping[str, object] | None, - timeout: float | httpx.Timeout | None, - on_response: ResponseObserver, - python_fallback: Callable[[], Awaitable[object]], -) -> object: - """Await the Rust path, falling back to the caller's own Python path when - the bridge is unavailable or the call fails. - - The caller supplies the fallback, so the bridge stays free of provider - dispatch. This exists because a caller that dispatches asynchronously has - already returned a coroutine by the time a Rust failure surfaces, and so - cannot fall back on its own. - """ - - def adapt(rust_response: Mapping[str, object]) -> object: - on_response(rust_response) - return _build_model_response(rust_response, model_response) - - return await _CHAT.ainvoke( + return await aattempt( + load=_ACHAT.load, + enabled=rust_enabled(), + eligible=True, prepare=lambda: timeout_to_seconds(timeout), - call=lambda rust_achat_completions, timeout_seconds: rust_achat_completions( - model=model, - messages=messages, - optional_params=optional_params, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_seconds, - ), - fallback=python_fallback, + call=call, adapt=adapt, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), + ) + + +def error_handling(provider: str, model: str) -> ErrorHandling: + return ErrorHandling( + declined=ErrorAction.SKIP, + upstream=APIErrorMapping(provider=provider, model=model), + missing_metadata=ErrorAction.SKIP, ) diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py new file mode 100644 index 00000000000..7572e31a8d2 --- /dev/null +++ b/litellm/rust_bridge/dispatch.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from enum import Enum +from typing import Final, TypeAlias, TypeVar + +from litellm._logging import verbose_logger +from litellm.exceptions import APIError +from litellm.rust_bridge.bindings import native_declined_types, native_upstream_types +from litellm.rust_bridge.runtime import DispatchResult, Handled, NativeFailed, NativeSkipped, NativeSkipReason + +NativeT = TypeVar("NativeT") +PythonT = TypeVar("PythonT") + + +class ErrorAction(Enum): + RAISE = "raise" + SKIP = "skip" + + +@dataclass(frozen=True, slots=True) +class APIErrorMapping: + provider: str + model: str + + +FailureAction: TypeAlias = ErrorAction | APIErrorMapping + + +@dataclass(frozen=True, slots=True) +class ErrorHandling: + declined: FailureAction = ErrorAction.RAISE + upstream: FailureAction = ErrorAction.RAISE + unknown: FailureAction = ErrorAction.RAISE + missing_metadata: FailureAction = ErrorAction.RAISE + unexpected: FailureAction = ErrorAction.RAISE + + +PROPAGATE: Final = ErrorHandling() +PYTHON_ON_ERROR: Final = ErrorHandling( + declined=ErrorAction.SKIP, + upstream=ErrorAction.SKIP, + unknown=ErrorAction.SKIP, + missing_metadata=ErrorAction.SKIP, + unexpected=ErrorAction.SKIP, +) + + +def _handle_error(error: Exception, action: FailureAction, route: str, reason: NativeSkipReason) -> NativeSkipped: + match action: + case ErrorAction.SKIP: + return NativeSkipped(reason, str(error)) + case ErrorAction.RAISE: + raise error + case APIErrorMapping(provider, model): + args: Final[tuple[object, ...]] = error.args + attribute_status: Final = getattr(error, "status_code", None) + attribute_message: Final = getattr(error, "message", None) + status_value: Final = attribute_status if isinstance(attribute_status, int) else (args[0] if args else 0) + message_value: Final = ( + attribute_message if isinstance(attribute_message, str) else (args[1] if len(args) > 1 else str(error)) + ) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + raise APIError( + status_code=status or 500, + message=f"litellm rust {route}: {message}", + llm_provider=provider, + model=model, + ) from error + + +def _resolve(result: DispatchResult[NativeT], errors: ErrorHandling, route: str) -> Handled[NativeT] | NativeSkipped: + if not isinstance(result, NativeFailed): + return result + declined: Final = native_declined_types() + upstream: Final = native_upstream_types() + if not declined or not upstream: + return _handle_error(result.error, errors.missing_metadata, route, NativeSkipReason.FAILED) + if isinstance(result.error, declined): + return _handle_error(result.error, errors.declined, route, NativeSkipReason.DECLINED) + if isinstance(result.error, upstream): + return _handle_error(result.error, errors.upstream, route, NativeSkipReason.FAILED) + return _handle_error(result.error, errors.unknown, route, NativeSkipReason.FAILED) + + +def _log_skip(route: str, skipped: NativeSkipped) -> None: + verbose_logger.debug("Native %s skipped (%s): %s", route, skipped.reason.value, skipped.detail or "") + + +def dispatch( + *, + native: Callable[[], DispatchResult[NativeT]], + python: Callable[[], PythonT], + route: str, + errors: ErrorHandling, +) -> NativeT | PythonT: + try: + attempted: Final = native() + except Exception as error: # noqa: BLE001 # preserve declared handling of loading and adaptation failures + unexpected: Final = _handle_error(error, errors.unexpected, route, NativeSkipReason.FAILED) + _log_skip(route, unexpected) + return python() + result: Final = _resolve(attempted, errors, route) + match result: + case Handled(value): + return value + case NativeSkipped(): + _log_skip(route, result) + return python() + + +async def adispatch( + *, + native: Callable[[], Awaitable[DispatchResult[NativeT]]], + python: Callable[[], Awaitable[PythonT]], + route: str, + errors: ErrorHandling, +) -> NativeT | PythonT: + try: + attempted: Final = await native() + except Exception as error: # noqa: BLE001 # preserve declared handling of loading and adaptation failures + unexpected: Final = _handle_error(error, errors.unexpected, route, NativeSkipReason.FAILED) + _log_skip(route, unexpected) + return await python() + result: Final = _resolve(attempted, errors, route) + match result: + case Handled(value): + return value + case NativeSkipped(): + _log_skip(route, result) + return await python() + + +async def async_none() -> None: + return None diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 580b46d7c38..160e6f0f743 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -6,25 +6,13 @@ from typing import Final import httpx -from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged from litellm.rust_bridge.protocols import RustAmessages, RustMessages -from litellm.rust_bridge.runtime import ( - BridgeErrorContext, - EndpointDispatch, - NativeErrorPolicy, - always_enabled, - async_none, - identity, -) +from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt, identity from litellm.rust_bridge.timeouts import timeout_to_seconds -_MESSAGES: Final[EndpointDispatch[RustMessages, RustAmessages]] = EndpointDispatch.native( - route="messages", - sync=lambda native: native.messages, - asynchronous=lambda native: native.amessages, - enabled=always_enabled, - error_policy=NativeErrorPolicy.PROPAGATE, -) +_MESSAGES: Final[NativeBinding[RustMessages]] = NativeBinding(lambda native: native.messages) +_AMESSAGES: Final[NativeBinding[RustAmessages]] = NativeBinding(lambda native: native.amessages) def set_rust_messages( @@ -34,22 +22,22 @@ def set_rust_messages( ) -> None: if not isinstance(messages, Unchanged): if messages is None: - _MESSAGES.sync.reset() + _MESSAGES.reset() else: - _MESSAGES.sync.override(messages) + _MESSAGES.override(messages) if not isinstance(amessages, Unchanged): if amessages is None: - _MESSAGES.asynchronous.reset() + _AMESSAGES.reset() else: - _MESSAGES.asynchronous.override(amessages) + _AMESSAGES.override(amessages) def load_rust_messages() -> RustMessages | None: - return _MESSAGES.sync.load() + return _MESSAGES.load() def load_rust_amessages() -> RustAmessages | None: - return _MESSAGES.asynchronous.load() + return _AMESSAGES.load() def messages( @@ -61,8 +49,11 @@ def messages( custom_llm_provider: str | None, extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: - return _MESSAGES.invoke( +) -> DispatchResult[dict[str, object]]: + return attempt( + load=_MESSAGES.load, + enabled=True, + eligible=True, prepare=lambda: timeout_to_seconds(timeout), call=lambda rust_messages, timeout_seconds: rust_messages( model=model, @@ -73,9 +64,7 @@ def messages( extra_headers=extra_headers, timeout_seconds=timeout_seconds, ), - fallback=lambda: None, adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -88,8 +77,11 @@ async def amessages( custom_llm_provider: str | None, extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: - return await _MESSAGES.ainvoke( +) -> DispatchResult[dict[str, object]]: + return await aattempt( + load=_AMESSAGES.load, + enabled=True, + eligible=True, prepare=lambda: timeout_to_seconds(timeout), call=lambda rust_amessages, timeout_seconds: rust_amessages( model=model, @@ -100,7 +92,5 @@ async def amessages( extra_headers=extra_headers, timeout_seconds=timeout_seconds, ), - fallback=async_none, adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 169966da8be..4773e106355 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -1,31 +1,59 @@ -"""Thin Python wrapper for the native Rust OCR bridge.""" - from __future__ import annotations -from collections.abc import Awaitable, Callable, Mapping -from typing import Final, TypeVar +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final -from . import configuration as _configuration -from .bindings import UNCHANGED, Unchanged -from .protocols import RustAocr, RustOcr -from .runtime import ( - BridgeErrorContext, - EndpointDispatch, - NativeErrorPolicy, -) +import httpx +from pydantic import TypeAdapter -rust_ocr_enabled = _configuration.rust_ocr_enabled -rust = _configuration.rust -ResultT = TypeVar("ResultT") -RequestT = TypeVar("RequestT") +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.azure_ai.ocr.common_utils import is_azure_document_intelligence_model +from litellm.llms.base_llm.ocr.transformation import OCR_REQUEST_FORMAT_PARAM, BaseOCRConfig, OCRResponse +from litellm.rust_bridge import configuration as _configuration +from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged +from litellm.rust_bridge.protocols import RustAocr, RustOcr +from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt +from litellm.rust_bridge.timeouts import timeout_to_seconds + +rust: Final = _configuration.rust +rust_ocr_enabled: Final = _configuration.rust_ocr_enabled + +_OCR: Final[NativeBinding[RustOcr]] = NativeBinding(lambda native: native.ocr) +_AOCR: Final[NativeBinding[RustAocr]] = NativeBinding(lambda native: native.aocr) +_HEADERS: Final = TypeAdapter(dict[str, object]) -_OCR: Final[EndpointDispatch[RustOcr, RustAocr]] = EndpointDispatch.native( - route="ocr", - sync=lambda native: native.ocr, - asynchronous=lambda native: native.aocr, - enabled=_configuration.rust_ocr_enabled, - error_policy=NativeErrorPolicy.PROPAGATE, +@dataclass(frozen=True, slots=True) +class PreparedOCRRequest: + model: str + document: dict[str, object] + api_key: str | None + api_base: str | None + custom_llm_provider: str + extra_headers: dict[str, object] | None + provider_config: BaseOCRConfig + optional_params: dict[str, object] + litellm_params: dict[str, object] + effective_timeout: float | httpx.Timeout + litellm_logging_obj: LiteLLMLoggingObj + + +@dataclass(frozen=True, slots=True) +class _PreparedRustOCRCall: + api_key: str | None + api_base: str | None + headers: dict[str, object] + optional_params: dict[str, object] + + +_RUST_OCR_PROVIDERS: Final = frozenset( + { + "mistral", + "azure_ai", + "vertex_ai", + } ) @@ -36,59 +64,168 @@ def set_rust_ocr( ) -> None: if not isinstance(ocr, Unchanged): if ocr is None: - _OCR.sync.reset() + _OCR.reset() else: - _OCR.sync.override(ocr) + _OCR.override(ocr) if not isinstance(aocr, Unchanged): if aocr is None: - _OCR.asynchronous.reset() + _AOCR.reset() else: - _OCR.asynchronous.override(aocr) + _AOCR.override(aocr) def load_rust_ocr() -> RustOcr | None: - return _OCR.sync.load() + return _OCR.load() def load_rust_aocr() -> RustAocr | None: - return _OCR.asynchronous.load() + return _AOCR.load() -def dispatch_ocr( - *, - prepare: Callable[[], RequestT], - call: Callable[[RustOcr, RequestT], Mapping[str, object]], - fallback: Callable[[], ResultT], - adapt: Callable[[Mapping[str, object]], ResultT], - model: str, - provider: str, - eligible: bool, -) -> ResultT: - return _OCR.invoke( - prepare=prepare, - call=call, - fallback=fallback, - adapt=adapt, - error_context=BridgeErrorContext(provider=provider, model=model), - eligible=eligible, +def _rust_ocr_supported(prepared_request: PreparedOCRRequest) -> bool: + if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native": + return False + if not prepared_request.provider_config.supports_rust_bridge(): + return False + return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS + + +def _rust_bridge_optional_params( + prepared_request: PreparedOCRRequest, + resolve_secret: Callable[[str], str | None], +) -> dict[str, object]: + if prepared_request.custom_llm_provider != "vertex_ai": + return prepared_request.optional_params + vertex_project: Final = ( + prepared_request.litellm_params.get("vertex_project") + or prepared_request.litellm_params.get("vertex_ai_project") + or litellm.vertex_project + or resolve_secret("VERTEXAI_PROJECT") + ) + vertex_location: Final = ( + prepared_request.litellm_params.get("vertex_location") + or prepared_request.litellm_params.get("vertex_ai_location") + or litellm.vertex_location + or resolve_secret("VERTEXAI_LOCATION") + or resolve_secret("VERTEX_LOCATION") + ) + return { + **prepared_request.optional_params, + **{ + name: value + for name, value in (("vertex_project", vertex_project), ("vertex_location", vertex_location)) + if value is not None + }, + } + + +def _rust_bridge_api_base( + prepared_request: PreparedOCRRequest, + resolve_secret: Callable[[str], str | None], +) -> str | None: + if prepared_request.api_base is not None: + return prepared_request.api_base + if prepared_request.custom_llm_provider == "azure_ai": + if is_azure_document_intelligence_model(prepared_request.model): + return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") + return resolve_secret("AZURE_AI_API_BASE") + return None + + +def _prepare_rust_ocr_call( + prepared_request: PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], +) -> _PreparedRustOCRCall: + provider_config: Final = prepared_request.provider_config + api_key_env_var: Final = provider_config.get_api_key_env_var() + resolved_api_key: Final = prepared_request.api_key or ( + resolve_api_key(api_key_env_var) if api_key_env_var is not None else None + ) + resolved_headers: Final = _HEADERS.validate_python( + provider_config.validate_environment( + headers=prepared_request.extra_headers or {}, + model=prepared_request.model, + api_key=resolved_api_key, + api_base=prepared_request.api_base, + litellm_params=prepared_request.litellm_params, + ) + ) + resolved_complete_url: Final = provider_config.get_complete_url( + api_base=prepared_request.api_base, + model=prepared_request.model, + optional_params=prepared_request.optional_params, + litellm_params=prepared_request.litellm_params, + ) + rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key) + rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key) + prepared_request.litellm_logging_obj.pre_call( + input="OCR document processing", + api_key=resolved_api_key, + additional_args={ + "complete_input_dict": { + "model": prepared_request.model, + "document": prepared_request.document, + **rust_optional_params, + }, + "api_base": resolved_complete_url, + "headers": resolved_headers, + }, + ) + return _PreparedRustOCRCall( + api_key=resolved_api_key, + api_base=rust_api_base, + headers=resolved_headers, + optional_params=rust_optional_params, ) -async def adispatch_ocr( - *, - prepare: Callable[[], RequestT], - call: Callable[[RustAocr, RequestT], Awaitable[Mapping[str, object]]], - fallback: Callable[[], Awaitable[ResultT]], - adapt: Callable[[Mapping[str, object]], ResultT], - model: str, - provider: str, - eligible: bool, -) -> ResultT: - return await _OCR.ainvoke( - prepare=prepare, - call=call, - fallback=fallback, - adapt=adapt, - error_context=BridgeErrorContext(provider=provider, model=model), - eligible=eligible, +def attempt_ocr( + prepared_request: PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], +) -> DispatchResult[OCRResponse]: + return attempt( + load=_OCR.load, + enabled=rust_ocr_enabled(), + prepare=lambda: _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ), + call=lambda native, prepared: native( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), + ), + adapt=OCRResponse.model_validate, + eligible=_rust_ocr_supported(prepared_request), + ) + + +async def aattempt_ocr( + prepared_request: PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], +) -> DispatchResult[OCRResponse]: + return await aattempt( + load=_AOCR.load, + enabled=rust_ocr_enabled(), + prepare=lambda: _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ), + call=lambda native, prepared: native( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), + ), + adapt=OCRResponse.model_validate, + eligible=_rust_ocr_supported(prepared_request), ) diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index d82b3b94669..8b95729bfa9 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -7,24 +7,17 @@ from typing import Final import httpx from websockets.exceptions import ConnectionClosedOK -from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.protocols import ( RustResponsesWebSocket, RustResponsesWebSocketConnection, ) -from litellm.rust_bridge.runtime import ( - BridgeErrorContext, - EndpointBinding, - async_none, - identity, -) +from litellm.rust_bridge.runtime import DispatchResult, aattempt from litellm.rust_bridge.timeouts import timeout_to_seconds -_RESPONSES_WEBSOCKET: Final[EndpointBinding[RustResponsesWebSocketConnection]] = EndpointBinding.native( - route="responses_websocket", - select=lambda native: native.ResponsesWebSocketConnection, - enabled=rust_enabled, +_RESPONSES_WEBSOCKET: Final[NativeBinding[RustResponsesWebSocketConnection]] = NativeBinding( + lambda native: native.ResponsesWebSocketConnection, ) @@ -61,19 +54,16 @@ async def connect( url: str, headers: dict[str, str], timeout: float | httpx.Timeout | None, -) -> _ConnectionAdapter | None: - try: - connection: Final = await _RESPONSES_WEBSOCKET.ainvoke( - prepare=lambda: timeout_to_seconds(timeout), - call=lambda connection_type, timeout_seconds: connection_type.connect( - url=url, - headers=headers, - timeout_seconds=timeout_seconds, - ), - fallback=async_none, - adapt=identity, - error_context=BridgeErrorContext(provider="openai", model="responses websocket"), - ) - except Exception: # noqa: BLE001 # preserve the existing WebSocket connection fallback - return None - return None if connection is None else _ConnectionAdapter(connection) +) -> DispatchResult[_ConnectionAdapter]: + return await aattempt( + load=_RESPONSES_WEBSOCKET.load, + enabled=rust_enabled(), + eligible=True, + prepare=lambda: timeout_to_seconds(timeout), + call=lambda connection_type, timeout_seconds: connection_type.connect( + url=url, + headers=headers, + timeout_seconds=timeout_seconds, + ), + adapt=_ConnectionAdapter, + ) diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index 4cd74ddf24f..48f1ecd23b4 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -1,39 +1,22 @@ from __future__ import annotations from collections.abc import Awaitable, Callable -from dataclasses import dataclass, field +from dataclasses import dataclass from enum import Enum -from typing import Final, Generic, NoReturn, Protocol, TypeAlias, TypeVar - -from litellm.exceptions import APIError -from litellm.rust_bridge.bindings import ( - UNCHANGED, - NativeBinding, - Unchanged, - native_exception_types, -) -from litellm.rust_bridge.protocols import NativeModule +from typing import Final, Generic, TypeAlias, TypeVar BindingT = TypeVar("BindingT") -SelectedT = TypeVar("SelectedT") -SelectedSyncT = TypeVar("SelectedSyncT") -SelectedAsyncT = TypeVar("SelectedAsyncT") NativeT = TypeVar("NativeT") RequestT = TypeVar("RequestT") ResultT = TypeVar("ResultT") -SyncBindingT = TypeVar("SyncBindingT") -AsyncBindingT = TypeVar("AsyncBindingT") -class PythonFallbackReason(Enum): - NATIVE_DISABLED = "native_disabled" - NATIVE_UNAVAILABLE = "native_unavailable" - NATIVE_DECLINED = "native_declined" - - -class NativeErrorPolicy(Enum): - TRANSLATE = "translate" - PROPAGATE = "propagate" +class NativeSkipReason(Enum): + DISABLED = "disabled" + INELIGIBLE = "ineligible" + UNAVAILABLE = "unavailable" + DECLINED = "declined" + FAILED = "failed" @dataclass(frozen=True, slots=True) @@ -42,470 +25,65 @@ class Handled(Generic[ResultT]): @dataclass(frozen=True, slots=True) -class PythonFallback: - reason: PythonFallbackReason +class NativeSkipped: + reason: NativeSkipReason detail: str | None = None -DispatchResult: TypeAlias = Handled[ResultT] | PythonFallback - - @dataclass(frozen=True, slots=True) -class BridgeErrorContext: - provider: str - model: str +class NativeFailed: + error: Exception -class RustEnablement(Protocol): - def __call__(self) -> bool: ... +DispatchResult: TypeAlias = Handled[ResultT] | NativeSkipped | NativeFailed -@dataclass(frozen=True, slots=True) -class EndpointBinding(Generic[BindingT]): - route: str - load: Callable[[], BindingT | None] - enabled: RustEnablement - error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE - _native_binding: NativeBinding[BindingT] | None = field(default=None, repr=False) +def _select(load: Callable[[], BindingT | None], enabled: bool, eligible: bool) -> BindingT | NativeSkipped: + if not enabled: + return NativeSkipped(NativeSkipReason.DISABLED) + if not eligible: + return NativeSkipped(NativeSkipReason.INELIGIBLE) + binding: Final = load() + return NativeSkipped(NativeSkipReason.UNAVAILABLE) if binding is None else binding - @staticmethod - def native( - *, - route: str, - select: Callable[[NativeModule], SelectedT], - enabled: RustEnablement, - error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE, - ) -> EndpointBinding[SelectedT]: - binding: Final = NativeBinding(select) - return EndpointBinding( - route=route, - load=binding.load, - enabled=enabled, - error_policy=error_policy, - _native_binding=binding, - ) - def override(self, value: BindingT | None) -> None: - if self._native_binding is None: - raise RuntimeError("only native Rust bridges support binding overrides") - self._native_binding.override(value) - - def reset(self) -> None: - if self._native_binding is None: - raise RuntimeError("only native Rust bridges support binding resets") - self._native_binding.reset() - - def _attempt( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[BindingT, RequestT], NativeT], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> DispatchResult[ResultT]: - binding_or_fallback: Final = self._binding_or_python_fallback( - eligible=eligible, - ) - if isinstance(binding_or_fallback, PythonFallback): - return binding_or_fallback - preflight_result: Final = preflight() if preflight is not None else None - if preflight_result is not None: - return preflight_result - return self._attempt_call( - call=lambda: call(binding_or_fallback, prepare()), - adapt=adapt, - error_context=error_context, - ) - - async def _aattempt( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[BindingT, RequestT], Awaitable[NativeT]], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> DispatchResult[ResultT]: - binding_or_fallback: Final = self._binding_or_python_fallback( - eligible=eligible, - ) - if isinstance(binding_or_fallback, PythonFallback): - return binding_or_fallback - preflight_result: Final = preflight() if preflight is not None else None - if preflight_result is not None: - return preflight_result - return await self._attempt_acall( - call=lambda: call(binding_or_fallback, prepare()), - adapt=adapt, - error_context=error_context, - ) - - def invoke( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[BindingT, RequestT], NativeT], - fallback: Callable[[], ResultT], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> ResultT: - result: Final = self._attempt( - prepare=prepare, - call=call, - adapt=adapt, - error_context=error_context, - eligible=eligible, - preflight=preflight, - ) - match result: - case Handled(value=value): - return value - case PythonFallback(): - return fallback() - - async def ainvoke( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[BindingT, RequestT], Awaitable[NativeT]], - fallback: Callable[[], Awaitable[ResultT]], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> ResultT: - result: Final = await self._aattempt( - prepare=prepare, - call=call, - adapt=adapt, - error_context=error_context, - eligible=eligible, - preflight=preflight, - ) - match result: - case Handled(value=value): - return value - case PythonFallback(): - return await fallback() - - def assess( - self, - *, - check: Callable[[BindingT], str | None], - ) -> PythonFallback | None: - binding: Final = self._binding_or_python_fallback(eligible=True) - if isinstance(binding, PythonFallback): - return binding - reason: Final = check(binding) - return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, reason) if reason is not None else None - - def accepts( - self, - *, - check: Callable[[BindingT], str | None], - eligible: bool = True, - ) -> bool: - binding_or_fallback: Final = self._binding_or_python_fallback( - eligible=eligible, - ) - if isinstance(binding_or_fallback, PythonFallback): - return False - try: - reason: Final = check(binding_or_fallback) - except Exception: # noqa: BLE001 # preflight performs no provider I/O, so Python handoff is safe - return False - return reason is None - - def require( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[BindingT, RequestT], NativeT], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> ResultT: - result: Final = self._attempt( - prepare=prepare, - call=call, - adapt=adapt, - error_context=error_context, - eligible=eligible, - preflight=preflight, - ) - match result: - case Handled(value=value): - return value - case PythonFallback(): - self._raise_required(result) - - async def arequire( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[BindingT, RequestT], Awaitable[NativeT]], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> ResultT: - result: Final = await self._aattempt( - prepare=prepare, - call=call, - adapt=adapt, - error_context=error_context, - eligible=eligible, - preflight=preflight, - ) - match result: - case Handled(value=value): - return value - case PythonFallback(): - self._raise_required(result) - - def can_attempt( - self, - *, - eligible: bool = True, - ) -> bool: - return not isinstance( - self._binding_or_python_fallback(eligible=eligible), - PythonFallback, - ) - - def _raise_required(self, fallback: PythonFallback) -> NoReturn: - detail: Final = f": {fallback.detail}" if fallback.detail else "" - reason: Final = _required_reason(fallback.reason) - raise RuntimeError(f"native {self.route} endpoint {reason}{detail}") - - def _binding_or_python_fallback( - self, - *, - eligible: bool, - ) -> BindingT | PythonFallback: - if not eligible or not self.enabled(): - return PythonFallback(PythonFallbackReason.NATIVE_DISABLED) - binding: Final = self.load() - if binding is None: - return PythonFallback(PythonFallbackReason.NATIVE_UNAVAILABLE) +def attempt( + *, + load: Callable[[], BindingT | None], + enabled: bool, + eligible: bool, + prepare: Callable[[], RequestT], + call: Callable[[BindingT, RequestT], NativeT], + adapt: Callable[[NativeT], ResultT], +) -> DispatchResult[ResultT]: + binding: Final = _select(load, enabled, eligible) + if isinstance(binding, NativeSkipped): return binding - - def _attempt_call( - self, - *, - call: Callable[[], NativeT], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - ) -> DispatchResult[ResultT]: - if self.error_policy is NativeErrorPolicy.PROPAGATE: - return Handled(adapt(call())) - exceptions: Final = native_exception_types() - if exceptions is None: - try: - value_without_exceptions: Final = call() - except Exception as error: # noqa: BLE001 # preserve chat fallback when native exception classes are absent - return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error)) - return Handled(adapt(value_without_exceptions)) - declined, upstream = exceptions - try: - value: Final = call() - except declined as error: - return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error)) - except upstream as error: - self._raise_upstream(error, error_context) - return Handled(adapt(value)) - - async def _attempt_acall( - self, - *, - call: Callable[[], Awaitable[NativeT]], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - ) -> DispatchResult[ResultT]: - if self.error_policy is NativeErrorPolicy.PROPAGATE: - return Handled(adapt(await call())) - exceptions: Final = native_exception_types() - if exceptions is None: - try: - value_without_exceptions: Final = await call() - except Exception as error: # noqa: BLE001 # preserve chat fallback when native exception classes are absent - return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error)) - return Handled(adapt(value_without_exceptions)) - declined, upstream = exceptions - try: - value: Final = await call() - except declined as error: - return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error)) - except upstream as error: - self._raise_upstream(error, error_context) - return Handled(adapt(value)) - - def _raise_upstream(self, error: BaseException, error_context: BridgeErrorContext) -> NoReturn: - args: Final[tuple[object, ...]] = error.args - attribute_status: Final = getattr(error, "status_code", None) - attribute_message: Final = getattr(error, "message", None) - status_value: Final = attribute_status if isinstance(attribute_status, int) else (args[0] if args else 0) - message_value: Final = ( - attribute_message if isinstance(attribute_message, str) else (args[1] if len(args) > 1 else str(error)) - ) - status: Final = status_value if isinstance(status_value, int) else 0 - message: Final = message_value if isinstance(message_value, str) else str(message_value) - raise APIError( - status_code=status or 500, - message=f"litellm rust {self.route}: {message}", - llm_provider=error_context.provider, - model=error_context.model, - ) from error + try: + value: Final = call(binding, prepare()) + except Exception as error: # noqa: BLE001 # orchestration applies the endpoint's declared error policy + return NativeFailed(error) + return Handled(adapt(value)) -@dataclass(frozen=True, slots=True) -class EndpointDispatch(Generic[SyncBindingT, AsyncBindingT]): - sync: EndpointBinding[SyncBindingT] - asynchronous: EndpointBinding[AsyncBindingT] - - @staticmethod - def native( - *, - route: str, - sync: Callable[[NativeModule], SelectedSyncT], - asynchronous: Callable[[NativeModule], SelectedAsyncT], - enabled: RustEnablement, - error_policy: NativeErrorPolicy = NativeErrorPolicy.TRANSLATE, - ) -> EndpointDispatch[SelectedSyncT, SelectedAsyncT]: - return EndpointDispatch( - sync=EndpointBinding.native(route=route, select=sync, enabled=enabled, error_policy=error_policy), - asynchronous=EndpointBinding.native( - route=route, - select=asynchronous, - enabled=enabled, - error_policy=error_policy, - ), - ) - - def override( - self, - *, - sync: SyncBindingT | None | Unchanged = UNCHANGED, - asynchronous: AsyncBindingT | None | Unchanged = UNCHANGED, - ) -> None: - if not isinstance(sync, Unchanged): - self.sync.override(sync) - if not isinstance(asynchronous, Unchanged): - self.asynchronous.override(asynchronous) - - def reset(self) -> None: - self.sync.reset() - self.asynchronous.reset() - - def invoke( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[SyncBindingT, RequestT], NativeT], - fallback: Callable[[], ResultT], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> ResultT: - return self.sync.invoke( - prepare=prepare, - call=call, - fallback=fallback, - adapt=adapt, - error_context=error_context, - eligible=eligible, - preflight=preflight, - ) - - async def ainvoke( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]], - fallback: Callable[[], Awaitable[ResultT]], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> ResultT: - return await self.asynchronous.ainvoke( - prepare=prepare, - call=call, - fallback=fallback, - adapt=adapt, - error_context=error_context, - eligible=eligible, - preflight=preflight, - ) - - def require( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[SyncBindingT, RequestT], NativeT], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> ResultT: - return self.sync.require( - prepare=prepare, - call=call, - adapt=adapt, - error_context=error_context, - eligible=eligible, - preflight=preflight, - ) - - async def arequire( - self, - *, - prepare: Callable[[], RequestT], - call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]], - adapt: Callable[[NativeT], ResultT], - error_context: BridgeErrorContext, - eligible: bool = True, - preflight: Callable[[], PythonFallback | None] | None = None, - ) -> ResultT: - return await self.asynchronous.arequire( - prepare=prepare, - call=call, - adapt=adapt, - error_context=error_context, - eligible=eligible, - preflight=preflight, - ) - - -def _error_message(error: BaseException) -> str: - reason: Final[object] = error.args[0] if error.args else str(error) - return reason if isinstance(reason, str) else str(reason) - - -def _required_reason(reason: PythonFallbackReason) -> str: - match reason: - case PythonFallbackReason.NATIVE_DISABLED: - return "is disabled" - case PythonFallbackReason.NATIVE_UNAVAILABLE: - return "is unavailable" - case PythonFallbackReason.NATIVE_DECLINED: - return "declined the request" - - -def always_enabled() -> bool: - return True +async def aattempt( + *, + load: Callable[[], BindingT | None], + enabled: bool, + eligible: bool, + prepare: Callable[[], RequestT], + call: Callable[[BindingT, RequestT], Awaitable[NativeT]], + adapt: Callable[[NativeT], ResultT], +) -> DispatchResult[ResultT]: + binding: Final = _select(load, enabled, eligible) + if isinstance(binding, NativeSkipped): + return binding + try: + value: Final = await call(binding, prepare()) + except Exception as error: # noqa: BLE001 # orchestration applies the endpoint's declared error policy + return NativeFailed(error) + return Handled(adapt(value)) def identity(value: ResultT) -> ResultT: return value - - -async def async_none() -> None: - return None diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index b93ab5e625a..4970b86b18d 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -4,25 +4,13 @@ from typing import Final import httpx -from litellm.rust_bridge.bindings import UNCHANGED, Unchanged +from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged from litellm.rust_bridge.protocols import RustAtranscription, RustTranscription -from litellm.rust_bridge.runtime import ( - BridgeErrorContext, - EndpointDispatch, - NativeErrorPolicy, - always_enabled, - async_none, - identity, -) +from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt, identity from litellm.rust_bridge.timeouts import timeout_to_seconds -_TRANSCRIPTION: Final[EndpointDispatch[RustTranscription, RustAtranscription]] = EndpointDispatch.native( - route="audio transcription", - sync=lambda native: native.transcription, - asynchronous=lambda native: native.atranscription, - enabled=always_enabled, - error_policy=NativeErrorPolicy.PROPAGATE, -) +_TRANSCRIPTION: Final[NativeBinding[RustTranscription]] = NativeBinding(lambda native: native.transcription) +_ATRANSCRIPTION: Final[NativeBinding[RustAtranscription]] = NativeBinding(lambda native: native.atranscription) def configure_rust_transcription( @@ -32,22 +20,22 @@ def configure_rust_transcription( ) -> None: if not isinstance(transcription, Unchanged): if transcription is None: - _TRANSCRIPTION.sync.reset() + _TRANSCRIPTION.reset() else: - _TRANSCRIPTION.sync.override(transcription) + _TRANSCRIPTION.override(transcription) if not isinstance(atranscription, Unchanged): if atranscription is None: - _TRANSCRIPTION.asynchronous.reset() + _ATRANSCRIPTION.reset() else: - _TRANSCRIPTION.asynchronous.override(atranscription) + _ATRANSCRIPTION.override(atranscription) def load_rust_transcription() -> RustTranscription | None: - return _TRANSCRIPTION.sync.load() + return _TRANSCRIPTION.load() def load_rust_atranscription() -> RustAtranscription | None: - return _TRANSCRIPTION.asynchronous.load() + return _ATRANSCRIPTION.load() def transcription( @@ -60,8 +48,11 @@ def transcription( extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: - return _TRANSCRIPTION.invoke( +) -> DispatchResult[dict[str, object]]: + return attempt( + load=_TRANSCRIPTION.load, + enabled=True, + eligible=True, prepare=lambda: timeout_to_seconds(timeout), call=lambda rust_transcription, timeout_seconds: rust_transcription( model=model, @@ -73,9 +64,7 @@ def transcription( optional_params=optional_params, timeout_seconds=timeout_seconds, ), - fallback=lambda: None, adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -89,8 +78,11 @@ async def atranscription( extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: - return await _TRANSCRIPTION.ainvoke( +) -> DispatchResult[dict[str, object]]: + return await aattempt( + load=_ATRANSCRIPTION.load, + enabled=True, + eligible=True, prepare=lambda: timeout_to_seconds(timeout), call=lambda rust_atranscription, timeout_seconds: rust_atranscription( model=model, @@ -102,7 +94,5 @@ async def atranscription( optional_params=optional_params, timeout_seconds=timeout_seconds, ), - fallback=async_none, adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 7a9abf1db97..fd3bc01321c 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -156,7 +156,7 @@ "limit": 215 }, "PLW0603": { - "limit": 186 + "limit": 184 }, "PLW1508": { "limit": 190 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1031 + "limit": 1029 }, "TRY002": { "limit": 524 @@ -246,7 +246,7 @@ "limit": 109 }, "TRY300": { - "limit": 848 + "limit": 846 }, "UP028": { "limit": 2 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 e8126600467..1502a47a7d9 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import configuration +from litellm.rust_bridge.runtime import Handled, NativeSkipped, NativeSkipReason from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -133,7 +134,7 @@ def test_load_rust_amessages_returns_injected_impl(): assert rust_messages.load_rust_amessages() is bridge -def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): +def test_messages_wrapper_reports_unavailable(monkeypatch): monkeypatch.setattr( importlib.import_module("litellm.rust_bridge.bindings"), "get_native_bridge", @@ -150,7 +151,7 @@ def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): extra_headers={}, timeout=30.0, ) - assert result is None + assert result == NativeSkipped(NativeSkipReason.UNAVAILABLE) def test_messages_wrapper_forwards_args_and_converts_timeout(): @@ -168,7 +169,7 @@ def test_messages_wrapper_forwards_args_and_converts_timeout(): timeout=httpx.Timeout(600.0, read=42.0), ) - assert response == FAKE_MESSAGES_RESPONSE + assert response == Handled(FAKE_MESSAGES_RESPONSE) assert bridge.calls[0] == { "model": "claude-sonnet-4-5", "body": REQUEST_BODY, @@ -196,7 +197,7 @@ async def test_amessages_wrapper_forwards_args(): timeout=12.5, ) - assert response == FAKE_MESSAGES_RESPONSE + assert response == Handled(FAKE_MESSAGES_RESPONSE) assert bridge.calls[0]["model"] == "claude-sonnet-4-5" assert bridge.calls[0]["timeout_seconds"] == 12.5 @@ -244,7 +245,6 @@ async def test_gate_falls_back_to_python_when_bridge_raises(): rust_messages.set_rust_messages(amessages=bridge) response = await _gate() - assert response is None assert bridge.calls == 1 diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py index 249fbda713e..1ab740c4267 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -12,13 +12,13 @@ import pytest import litellm from litellm.llms.azure_ai.ocr.cohere_parse_transformation import AzureAICohereParseConfig from litellm.llms.cohere.ocr.transformation import CohereParseConfig -from litellm.ocr.main import _PreparedOCRRequest, _rust_ocr_supported +from litellm.rust_bridge.ocr import PreparedOCRRequest, _rust_ocr_supported DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -def _prepared(optional_params: dict[str, object]) -> _PreparedOCRRequest: - return _PreparedOCRRequest( +def _prepared(optional_params: dict[str, object]) -> PreparedOCRRequest: + return PreparedOCRRequest( model="doc-intelligence/prebuilt-layout", document=dict(DOCUMENT), api_key="fake-key", diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index b936b0e75c2..b1bbc73e9e5 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import configuration +from litellm.rust_bridge.runtime import Handled from litellm.rust_bridge.timeouts import timeout_to_seconds # `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` @@ -202,7 +203,7 @@ def build_prepared_request( litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = 12.5, ) -> Any: - return ocr_main._PreparedOCRRequest( + return rust_bridge.PreparedOCRRequest( model=model, document=document, api_key=api_key, @@ -399,99 +400,13 @@ def test_timeout_to_seconds_handles_float_timeout_and_none(): assert timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 -def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): - bridge = RecordingBridge() - - litellm.rust(True) - - rust_bridge.set_rust_ocr(ocr=bridge) - response = rust_bridge.dispatch_ocr( - prepare=lambda: 12.5, - call=lambda native, timeout: native( - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://proxy.internal", - custom_llm_provider="mistral", - extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True, "pages": [0]}, - timeout_seconds=timeout, - ), - fallback=lambda: pytest.fail("unexpected Python fallback"), - adapt=dict, - model="mistral-ocr-latest", - provider="mistral", - eligible=True, - ) - - assert response == FAKE_OCR_RESPONSE - call = bridge.calls[0] - assert call == { - "model": "mistral-ocr-latest", - "document": DOCUMENT, - "api_key": "sk-test", - "api_base": "https://proxy.internal", - "custom_llm_provider": "mistral", - "extra_headers": { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - }, - "optional_params": {"include_image_base64": True, "pages": [0]}, - "timeout_seconds": 12.5, - } - - -@pytest.mark.asyncio -async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): - bridge = RecordingAsyncBridge() - - litellm.rust(True) - - rust_bridge.set_rust_ocr(aocr=bridge) - - async def unexpected_fallback(): - pytest.fail("unexpected Python fallback") - - response = await rust_bridge.adispatch_ocr( - prepare=lambda: 42.0, - call=lambda native, timeout: native( - model="mistral-ocr-maas", - document=DOCUMENT, - api_key=None, - api_base=None, - custom_llm_provider="vertex_ai", - extra_headers=None, - optional_params={"vertex_project": "project-1"}, - timeout_seconds=timeout, - ), - fallback=unexpected_fallback, - adapt=dict, - model="mistral-ocr-maas", - provider="vertex_ai", - eligible=True, - ) - - assert response == FAKE_OCR_RESPONSE - assert bridge.calls[0] == { - "model": "mistral-ocr-maas", - "document": DOCUMENT, - "api_key": None, - "api_base": None, - "custom_llm_provider": "vertex_ai", - "extra_headers": None, - "optional_params": {"vertex_project": "project-1"}, - "timeout_seconds": 42.0, - } - - def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - response = ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + response = rust_bridge.attempt_ocr( prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://proxy.internal", @@ -502,6 +417,8 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): resolve_api_key=lambda _name: None, ) + assert isinstance(response, Handled) + response = response.value assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "hello world" assert bridge.calls[0] == { @@ -524,8 +441,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + rust_bridge.attempt_ocr( prepared_request=build_prepared_request(api_key=None, timeout=None), resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, ) @@ -541,8 +457,7 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver(): def _resolver(name: str) -> str | None: raise AssertionError(f"resolver should not be called for {name}") - ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + rust_bridge.attempt_ocr( prepared_request=build_prepared_request( api_key="sk-explicit", timeout=None, @@ -563,8 +478,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): resolver_calls.append(name) return "sk-provider-env" - ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + rust_bridge.attempt_ocr( prepared_request=build_prepared_request( provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), model="provider-ocr-model", @@ -583,8 +497,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + rust_bridge.attempt_ocr( prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -617,8 +530,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana "VERTEXAI_LOCATION": "us-east5", }.get(name) - ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + rust_bridge.attempt_ocr( prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -636,8 +548,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + rust_bridge.attempt_ocr( prepared_request=build_prepared_request( custom_llm_provider="azure_ai", model="pixtral-12b-2409", @@ -655,8 +566,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + rust_bridge.attempt_ocr( prepared_request=build_prepared_request( custom_llm_provider="azure_ai", model="doc-intelligence/prebuilt-layout", @@ -677,8 +587,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - ocr_main._run_rust_ocr( - fallback=lambda: pytest.fail("unexpected Python fallback"), + rust_bridge.attempt_ocr( prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://api.mistral.ai/v1", @@ -851,7 +760,7 @@ async def test_ocr_fallback_skips_native_preparation( def unexpected_preparation(*_args: object, **_kwargs: object) -> None: pytest.fail("Python fallback must not resolve native credentials or emit native pre_call") - monkeypatch.setattr(ocr_main, "_prepare_rust_ocr_call", unexpected_preparation) + monkeypatch.setattr(rust_bridge, "_prepare_rust_ocr_call", unexpected_preparation) monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fallback) response: Final = ( diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 6e24d6def7a..d00ef5c8127 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -4,6 +4,7 @@ import pytest from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled from litellm.rust_bridge import configuration, responses_websocket +from litellm.rust_bridge.runtime import Handled, NativeSkipped, NativeSkipReason, NativeFailed class _FakeNativeConnection: @@ -64,18 +65,15 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: @pytest.mark.asyncio -async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_bridge_reports_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: configuration.rust(True) responses_websocket._RESPONSES_WEBSOCKET.override(None) - assert ( - await responses_websocket.connect( - url="wss://example.test/responses", - headers={}, - timeout=None, - ) - is None - ) + assert await responses_websocket.connect( + url="wss://example.test/responses", + headers={}, + timeout=None, + ) == NativeSkipped(NativeSkipReason.UNAVAILABLE) @pytest.mark.asyncio @@ -91,7 +89,8 @@ async def test_enabled_bridge_connects_and_adapts_socket( timeout=1.0, ) - assert connection is not None + assert isinstance(connection, Handled) + connection = connection.value await connection.send("response.create") assert await connection.recv() == "response.completed" await connection.close() @@ -110,15 +109,9 @@ class _FailingNativeBridge: @pytest.mark.asyncio -async def test_connection_failure_preserves_python_fallback() -> None: +async def test_connection_failure_is_reported_to_orchestration() -> None: configuration.rust(True) responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge) - - assert ( - await responses_websocket.connect( - url="wss://example.test/responses", - headers={}, - timeout=None, - ) - is None - ) + result = await responses_websocket.connect(url="wss://example.test/responses", headers={}, timeout=None) + assert isinstance(result, NativeFailed) + assert str(result.error) == "connection failed" diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py index f9d7f6dbd4e..cd562ee91c1 100644 --- a/tests/test_litellm/rust_bridge/test_bindings.py +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -48,7 +48,8 @@ def test_native_exception_types_reject_non_exception_classes(monkeypatch: pytest native: Final = SimpleNamespace(RustBridgeDeclined=invalid, RustUpstreamError=RuntimeError) monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) - assert bindings.native_exception_types() is None + assert bindings.native_declined_types() == () + assert bindings.native_upstream_types() == (RuntimeError,) @pytest.mark.parametrize( @@ -59,10 +60,6 @@ def test_native_exception_types_reject_non_exception_classes(monkeypatch: pytest "wrong: NativeBinding[RustAchatCompletions] = NativeBinding(lambda native: native.chat_completions)", "reportAssignmentType", ), - ( - 'EndpointBinding.native(route="chat", select=lambda native: native.chat_completion, enabled=always_enabled)', - "reportAttributeAccessIssue", - ), ("NativeBinding(lambda native: native.ocrr)", "reportAttributeAccessIssue"), ( "wrong: NativeBinding[RustAmessages] = NativeBinding(lambda native: native.messages)", @@ -85,14 +82,8 @@ def test_selectors_are_checked_by_type_checker(tmp_path: Path, expression: str, "from litellm.rust_bridge.bindings import NativeBinding\n" "from litellm.rust_bridge.protocols import RustChatCompletions, RustAchatCompletions, " "RustMessages, RustAmessages, RustOcr, RustAocr, RustTranscription, RustAtranscription\n" - "from litellm.rust_bridge.runtime import EndpointBinding, EndpointDispatch, always_enabled\n" "binding = NativeBinding(lambda native: native.chat_completions)\n" "assert_type(binding, NativeBinding[RustChatCompletions])\n" - 'bridge = EndpointBinding.native(route="chat", select=lambda native: native.chat_completions, enabled=always_enabled)\n' - "assert_type(bridge, EndpointBinding[RustChatCompletions])\n" - 'endpoint = EndpointDispatch.native(route="chat", sync=lambda native: native.chat_completions, ' - "asynchronous=lambda native: native.achat_completions, enabled=always_enabled)\n" - "assert_type(endpoint, EndpointDispatch[RustChatCompletions, RustAchatCompletions])\n" "assert_type(NativeBinding(lambda native: native.messages), NativeBinding[RustMessages])\n" "assert_type(NativeBinding(lambda native: native.ocr), NativeBinding[RustOcr])\n" "assert_type(NativeBinding(lambda native: native.transcription), NativeBinding[RustTranscription])\n" @@ -117,4 +108,6 @@ def test_selectors_are_checked_by_type_checker(tmp_path: Path, expression: str, ) diagnostics: Final = json.loads(result.stdout)["generalDiagnostics"] assert result.returncode == 1, result.stdout + result.stderr - assert [(item["rule"], item["range"]["start"]["line"]) for item in diagnostics] == [(expected_rule, 13)] + assert [(item["rule"], item["range"]["start"]["line"]) for item in diagnostics] == [ + (expected_rule, len(source.read_text().splitlines()) - 1) + ] diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 49cb1822ccc..f3f2ac651bf 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm.rust_bridge import bindings, configuration from litellm.rust_bridge import chat_completions as bridge +from litellm.rust_bridge.runtime import Handled, NativeSkipped, NativeFailed from litellm.types.utils import ModelResponse RUST_RESPONSE = { @@ -250,7 +251,8 @@ class TestSyncCall: result = bridge.chat_completions(**_call_kwargs(model_response)) - assert result is not None + assert isinstance(result, Handled) + result = result.value assert result.choices[0].message.content == "hello from rust" assert result.choices[0].finish_reason == "stop" assert result.model == "claude-sonnet-4-5-20260101" @@ -266,14 +268,14 @@ class TestSyncCall: bridge.chat_completions(**_call_kwargs(ModelResponse())) assert native.calls[0]["timeout_seconds"] == 30.0 - def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch): + def test_reports_unavailable_bridge(self, monkeypatch): _hide_native_bridge(monkeypatch) - assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None + assert isinstance(bridge.chat_completions(**_call_kwargs(ModelResponse())), NativeSkipped) - def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch): + def test_reports_native_decline_to_orchestration(self, monkeypatch): _fake_native_bridge(monkeypatch) bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming"))) - assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None + assert isinstance(bridge.chat_completions(**_call_kwargs(ModelResponse())), NativeFailed) class TestAsyncCall: @@ -281,134 +283,18 @@ class TestAsyncCall: async def test_builds_a_model_response(self): bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall()) result = await bridge.achat_completions(**_call_kwargs(ModelResponse())) - assert result is not None + assert isinstance(result, Handled) + result = result.value assert result.choices[0].message.content == "hello from rust" assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} @pytest.mark.asyncio - async def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch): + async def test_reports_unavailable_bridge(self, monkeypatch): _hide_native_bridge(monkeypatch) - assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None + assert isinstance(await bridge.achat_completions(**_call_kwargs(ModelResponse())), NativeSkipped) @pytest.mark.asyncio - async def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch): + async def test_reports_native_decline_to_orchestration(self, monkeypatch): _fake_native_bridge(monkeypatch) bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming"))) - assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None - - -class TestAsyncFallbackWrapper: - @pytest.mark.asyncio - async def test_returns_the_rust_response_without_running_the_fallback(self): - bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall()) - ran = [] - - async def fallback(): - ran.append(True) - return "python" - - result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) - assert result.choices[0].message.content == "hello from rust" - assert ran == [] - - @pytest.mark.asyncio - async def test_runs_the_fallback_when_the_core_declines(self, monkeypatch): - _fake_native_bridge(monkeypatch) - bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming"))) - - async def fallback(): - return "python" - - result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) - assert result == "python" - - @pytest.mark.asyncio - async def test_runs_the_fallback_when_the_bridge_is_unavailable(self, monkeypatch): - _hide_native_bridge(monkeypatch) - - async def fallback(): - return "python" - - result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) - assert result == "python" - - -class TestFailureClassification: - """A failure the provider already saw must not be retried on the Python - path: it would bill the customer for the same work twice.""" - - @pytest.fixture(autouse=True) - def _native_exceptions(self, monkeypatch): - _fake_native_bridge(monkeypatch) - - def test_a_decline_falls_back_because_nothing_was_sent(self): - bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming"))) - assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None - - def test_an_upstream_failure_is_surfaced_with_its_status(self): - from litellm.exceptions import APIError - - bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited"))) - with pytest.raises(APIError) as raised: - bridge.chat_completions(**_call_kwargs(ModelResponse())) - assert raised.value.status_code == 429 - assert "rate limited" in str(raised.value) - - def test_a_transport_failure_with_no_response_surfaces_as_a_500(self): - from litellm.exceptions import APIError - - bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset"))) - with pytest.raises(APIError) as raised: - bridge.chat_completions(**_call_kwargs(ModelResponse())) - assert raised.value.status_code == 500 - - def test_an_unrecognized_error_is_not_swallowed(self): - bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=RuntimeError("something else"))) - with pytest.raises(RuntimeError): - bridge.chat_completions(**_call_kwargs(ModelResponse())) - - @pytest.mark.asyncio - async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self): - from litellm.exceptions import APIError - - bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom"))) - ran = [] - - async def fallback(): - ran.append(True) - return "python" - - with pytest.raises(APIError): - await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) - assert ran == [], "a request the provider already served must not be re-issued" - - @pytest.mark.asyncio - async def test_the_async_wrapper_falls_back_on_a_decline(self): - bridge.set_rust_chat_completions( - achat_completions=_RecordingAsyncCall(error=_FakeDeclined("blank message text")) - ) - - async def fallback(): - return "python" - - result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) - assert result == "python" - - -@pytest.mark.asyncio -async def test_missing_native_exception_types_preserves_python_fallback(monkeypatch): - _hide_native_bridge(monkeypatch) - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=RuntimeError("connection failed")), - achat_completions=_RecordingAsyncCall(error=RuntimeError("connection failed")), - ) - - assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None - - async def fallback(): - return "python" - - assert ( - await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) - == "python" - ) + assert isinstance(await bridge.achat_completions(**_call_kwargs(ModelResponse())), NativeFailed) diff --git a/tests/test_litellm/rust_bridge/test_dispatch.py b/tests/test_litellm/rust_bridge/test_dispatch.py new file mode 100644 index 00000000000..3641729e521 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_dispatch.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import asyncio +import logging +from types import SimpleNamespace +from typing import Final + +import pytest + +from litellm.exceptions import APIError +from litellm.rust_bridge import bindings +from litellm.rust_bridge.chat_completions import error_handling +from litellm.rust_bridge.dispatch import PROPAGATE, PYTHON_ON_ERROR, adispatch, dispatch +from litellm.rust_bridge.runtime import DispatchResult, Handled, NativeFailed, NativeSkipped, NativeSkipReason + + +class Declined(Exception): + pass + + +class Upstream(Exception): + pass + + +@pytest.fixture(autouse=True) +def native_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + bindings, "get_native_bridge", lambda: SimpleNamespace(RustBridgeDeclined=Declined, RustUpstreamError=Upstream) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("reason", tuple(NativeSkipReason)) +async def test_shared_dispatch_calls_python_once_and_logs_skip( + asynchronous: bool, reason: NativeSkipReason, caplog: pytest.LogCaptureFixture +) -> None: + caplog.set_level(logging.DEBUG, logger="LiteLLM") + calls: Final[list[str]] = [] + + def native() -> DispatchResult[str]: + calls.append("native") + return NativeSkipped(reason, "diagnostic detail") + + async def anative() -> DispatchResult[str]: + return native() + + def python() -> str: + calls.append("python") + return "python response" + + async def apython() -> str: + return python() + + result: Final = ( + await adispatch(native=anative, python=apython, route="test", errors=PROPAGATE) + if asynchronous + else dispatch(native=native, python=python, route="test", errors=PROPAGATE) + ) + assert result == "python response" + assert calls == ["native", "python"] + assert f"Native test skipped ({reason.value}): diagnostic detail" in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_native_success_does_not_run_python_even_when_value_is_none(asynchronous: bool) -> None: + async def native() -> DispatchResult[None]: + return Handled(None) + + def python() -> str: + pytest.fail("handled results must not run Python") + + async def apython() -> str: + return python() + + result: Final = ( + await adispatch(native=native, python=apython, route="test", errors=PYTHON_ON_ERROR) + if asynchronous + else dispatch(native=lambda: Handled(None), python=python, route="test", errors=PYTHON_ON_ERROR) + ) + assert result is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("policy", ("chat", "propagate", "python")) +@pytest.mark.parametrize("kind", ("declined", "upstream", "unknown", "unexpected", "missing")) +async def test_declarations_preserve_endpoint_error_behavior( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, policy: str, kind: str +) -> None: + if kind == "missing": + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + error: Final = ( + Declined("unsupported") + if kind == "declined" + else Upstream(429, "rate limited") + if kind == "upstream" + else RuntimeError("failed") + ) + rules: Final = ( + error_handling("anthropic", "model") + if policy == "chat" + else PYTHON_ON_ERROR + if policy == "python" + else PROPAGATE + ) + calls: Final[list[str]] = [] + + def native() -> DispatchResult[str]: + if kind == "unexpected": + raise error + return NativeFailed(error) + + async def anative() -> DispatchResult[str]: + return native() + + def python() -> str: + calls.append("python") + return "python response" + + async def apython() -> str: + return python() + + async def run() -> str: + if asynchronous: + return await adispatch(native=anative, python=apython, route="chat_completions", errors=rules) + return dispatch(native=native, python=python, route="chat_completions", errors=rules) + + if policy == "python" or (policy == "chat" and kind in ("declined", "missing")): + assert await run() == "python response" + assert calls == ["python"] + elif policy == "chat" and kind == "upstream": + with pytest.raises(APIError) as caught: + await run() + assert caught.value.status_code == 429 + assert caught.value.model == "model" + assert caught.value.llm_provider == "anthropic" + assert caught.value.__cause__ is error + assert calls == [] + else: + with pytest.raises(type(error)) as caught_original: + await run() + assert caught_original.value is error + assert calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_python_failure_is_never_reclassified_as_native_failure(asynchronous: bool) -> None: + error: Final = RuntimeError("Python failed") + calls: Final[list[str]] = [] + + async def native() -> DispatchResult[str]: + return NativeSkipped(NativeSkipReason.UNAVAILABLE) + + def python() -> str: + calls.append("python") + raise error + + async def apython() -> str: + return python() + + async def run() -> str: + if asynchronous: + return await adispatch(native=native, python=apython, route="test", errors=PYTHON_ON_ERROR) + return dispatch( + native=lambda: NativeSkipped(NativeSkipReason.UNAVAILABLE), + python=python, + route="test", + errors=PYTHON_ON_ERROR, + ) + + with pytest.raises(RuntimeError) as caught: + await run() + assert caught.value is error + assert calls == ["python"] + + +@pytest.mark.asyncio +async def test_cancellation_does_not_run_python() -> None: + async def native() -> DispatchResult[str]: + raise asyncio.CancelledError + + async def python() -> str: + pytest.fail("cancellation must not dispatch Python") + + with pytest.raises(asyncio.CancelledError): + await adispatch(native=native, python=python, route="test", errors=PYTHON_ON_ERROR) + + +@pytest.mark.parametrize("status", (0, 401, 403, 429, 500, 503)) +def test_chat_upstream_mapping_preserves_status_message_and_context(status: int) -> None: + error: Final = Upstream(status, "upstream failed") + with pytest.raises(APIError, match="upstream failed") as caught: + dispatch( + native=lambda: NativeFailed(error), + python=lambda: pytest.fail("upstream errors must not run Python"), + route="chat_completions", + errors=error_handling("anthropic", "model"), + ) + assert caught.value.status_code == (status or 500) + assert caught.value.model == "model" + assert caught.value.llm_provider == "anthropic" + assert caught.value.__cause__ is error diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index 1ba13fa99fd..a882e23fb58 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -1,577 +1,117 @@ from __future__ import annotations -from dataclasses import dataclass -from types import SimpleNamespace from typing import Final import pytest -from litellm.exceptions import APIError -from litellm.rust_bridge import bindings, runtime +from litellm.rust_bridge import runtime -class RustBridgeDeclined(Exception): - pass - - -class RustUpstreamError(Exception): - pass - - -@pytest.fixture(autouse=True) -def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - bindings, - "get_native_bridge", - lambda: SimpleNamespace( - RustBridgeDeclined=RustBridgeDeclined, - RustUpstreamError=RustUpstreamError, - ), - ) - - -def context() -> runtime.BridgeErrorContext: - return runtime.BridgeErrorContext(provider="anthropic", model="model") - - -def enabled() -> bool: - return True - - -@dataclass(frozen=True, slots=True) -class FallbackCase: - process_enabled: bool | None = None - eligible: bool = True - binding_available: bool = True - declined: bool = False - expected_events: tuple[str, ...] = () - - -FALLBACK_CASES: Final = ( - pytest.param( - FallbackCase(process_enabled=False, expected_events=("python",)), - id="process-disabled", - ), - pytest.param( - FallbackCase(eligible=False, expected_events=("python",)), - id="request-ineligible", - ), - pytest.param( - FallbackCase(binding_available=False, expected_events=("load", "python")), - id="bridge-unavailable", - ), - pytest.param( - FallbackCase(declined=True, expected_events=("load", "prepare", "rust", "python")), - id="bridge-declined", - ), -) - - -@pytest.mark.parametrize("case", FALLBACK_CASES) -def test_invoke_falls_back_only_before_provider_success(case: FallbackCase) -> None: - events: list[str] = [] - - def load() -> object | None: - events.append("load") - return object() if case.binding_available else None - - def call(_binding: object, _request: object) -> int: - events.append("rust") - if case.declined: - raise RustBridgeDeclined("unsupported") - return 3 - - bridge: Final = runtime.EndpointBinding( - route="messages", load=load, enabled=lambda: case.process_enabled is not False - ) - result: Final = bridge.invoke( - prepare=lambda: events.append("prepare"), - call=call, - fallback=lambda: events.append("python") or "fallback", - adapt=str, - error_context=context(), - eligible=case.eligible, - ) - - assert result == "fallback" - assert tuple(events) == case.expected_events - - -@pytest.mark.asyncio -@pytest.mark.parametrize("case", FALLBACK_CASES) -async def test_ainvoke_matches_sync_fallback_contract(case: FallbackCase) -> None: - events: list[str] = [] - - def load() -> object | None: - events.append("load") - return object() if case.binding_available else None - - async def call(_binding: object, _request: object) -> int: - events.append("rust") - if case.declined: - raise RustBridgeDeclined("unsupported") - return 3 - - async def fallback() -> str: - events.append("python") - return "fallback" - - bridge: Final = runtime.EndpointBinding( - route="messages", load=load, enabled=lambda: case.process_enabled is not False - ) - result: Final = await bridge.ainvoke( - prepare=lambda: events.append("prepare"), - call=call, - fallback=fallback, - adapt=str, - error_context=context(), - eligible=case.eligible, - ) - - assert result == "fallback" - assert tuple(events) == case.expected_events - - -def test_invoke_adapts_native_success_without_fallback() -> None: - bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) - - result: Final = bridge.invoke( - prepare=lambda: 3, - call=lambda _binding, request: request * 2, - fallback=lambda: pytest.fail("fallback must not run"), - adapt=lambda value: f"adapted-{value}", - error_context=context(), - ) - - assert result == "adapted-6" - - -@pytest.mark.asyncio -async def test_ainvoke_adapts_native_success_without_fallback() -> None: - async def call(_binding: object, request: int) -> int: - return request * 2 - - async def fallback() -> str: - pytest.fail("fallback must not run") - - bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) - result: Final = await bridge.ainvoke( - prepare=lambda: 3, - call=call, - fallback=fallback, - adapt=lambda value: f"adapted-{value}", - error_context=context(), - ) - - assert result == "adapted-6" - - -@pytest.mark.parametrize( - ("error", "expected_type", "expected_status", "expected_message"), - ( - pytest.param(RustUpstreamError(401, "unauthorized"), APIError, 401, "unauthorized", id="auth"), - pytest.param(RustUpstreamError(429, "rate limited"), APIError, 429, "rate limited", id="rate-limit"), - pytest.param(RustUpstreamError(500, "failed"), APIError, 500, "failed", id="server-error"), - pytest.param(RustUpstreamError(0, "connection reset"), APIError, 500, "connection reset", id="transport"), - pytest.param(RustUpstreamError(403, "forbidden"), APIError, 403, "forbidden", id="other-status"), - ), -) @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", (False, True)) -async def test_upstream_failure_maps_to_api_error_without_fallback( - asynchronous: bool, - error: RustUpstreamError, - expected_type: type[BaseException], - expected_status: int, - expected_message: str, -) -> None: - def fail(_binding: object, _request: object) -> object: - raise error - - async def afail(binding: object, request: object) -> object: - return fail(binding, request) - - async def fallback() -> str: - pytest.fail("fallback must not run") - - bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) - - async def invoke() -> None: - if asynchronous: - await bridge.ainvoke( - prepare=lambda: None, call=afail, fallback=fallback, adapt=str, error_context=context() - ) - else: - bridge.invoke( - prepare=lambda: None, - call=fail, - fallback=lambda: pytest.fail("fallback must not run"), - adapt=str, - error_context=context(), - ) - - with pytest.raises(expected_type, match=expected_message) as caught: - await invoke() - - assert type(caught.value) is expected_type - assert caught.value.status_code == expected_status - assert caught.value.llm_provider == "anthropic" - assert caught.value.model == "model" - assert caught.value.__cause__ is error - - -@pytest.mark.asyncio -async def test_async_upstream_failure_maps_to_api_error_without_fallback() -> None: - async def fail(_binding: object, _request: object) -> object: - raise RustUpstreamError(503, "overloaded") - - async def fallback() -> object: - pytest.fail("fallback must not run") - - bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) - - with pytest.raises(APIError, match="overloaded") as caught: - await bridge.ainvoke(prepare=lambda: None, call=fail, fallback=fallback, adapt=str, error_context=context()) - - assert caught.value.status_code == 503 - - -def test_unknown_failure_is_preserved_without_fallback() -> None: - error: Final = RuntimeError("unknown") - bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) - - with pytest.raises(RuntimeError, match="unknown") as caught: - bridge.invoke( - prepare=lambda: None, - call=lambda _binding, _request: (_ for _ in ()).throw(error), - fallback=lambda: pytest.fail("fallback must not run"), - adapt=str, - error_context=context(), - ) - - assert caught.value is error - - -@pytest.mark.parametrize( - ("process_enabled", "binding_available", "declined", "expected_message"), - ( - pytest.param(False, True, False, "native messages endpoint is disabled", id="disabled"), - pytest.param(None, False, False, "native messages endpoint is unavailable", id="unavailable"), - pytest.param( - None, - True, - True, - "native messages endpoint declined the request: unsupported", - id="declined", - ), - ), -) -def test_require_explains_why_rust_did_not_handle_request( - process_enabled: bool | None, - binding_available: bool, - declined: bool, - expected_message: str, -) -> None: - def call(_binding: object, _request: object) -> object: - if declined: - raise RustBridgeDeclined("unsupported") - return object() - - bridge: Final = runtime.EndpointBinding( - route="messages", - load=object if binding_available else lambda: None, - enabled=lambda: process_enabled is not False, - ) - - with pytest.raises(RuntimeError, match=f"^{expected_message}$"): - bridge.require( - prepare=lambda: None, - call=call, - adapt=str, - error_context=context(), - ) - - -@pytest.mark.asyncio -async def test_arequire_explains_unavailable_native_binding() -> None: - endpoint: Final = runtime.EndpointBinding(route="messages", load=lambda: None, enabled=enabled) - - with pytest.raises(RuntimeError, match=r"^native messages endpoint is unavailable$"): - await endpoint.arequire( - prepare=lambda: pytest.fail("must not prepare"), - call=lambda _binding, _request: pytest.fail("must not invoke"), - adapt=str, - error_context=context(), - ) - - -@pytest.mark.parametrize( - ("state", "expected", "expected_events"), - ( - pytest.param("disabled", False, (), id="disabled"), - pytest.param("ineligible", False, (), id="ineligible"), - pytest.param("unavailable", False, ("load",), id="unavailable"), - pytest.param("available", True, ("load",), id="available"), - ), -) -def test_can_attempt_only_enabled_available_requests( - state: str, - expected: bool, - expected_events: tuple[str, ...], -) -> None: - events: list[str] = [] +@pytest.mark.parametrize("state", ("disabled", "ineligible", "unavailable", "handled")) +async def test_attempt_only_prepares_selected_requests(asynchronous: bool, state: str) -> None: + events: Final[list[str]] = [] def load() -> object | None: events.append("load") return None if state == "unavailable" else object() - bridge: Final = runtime.EndpointBinding(route="messages", load=load, enabled=lambda: state != "disabled") - - assert ( - bridge.can_attempt( - eligible=state != "ineligible", - ) - is expected - ) - assert tuple(events) == expected_events - - -def test_native_endpoint_applies_partial_overrides_and_reset(monkeypatch: pytest.MonkeyPatch) -> None: - def native_sync() -> str: - return "native" - - async def native_async() -> str: - return "native async" - - def replacement_sync() -> str: - return "replacement" - - monkeypatch.setattr( - bindings, - "get_native_bridge", - lambda: SimpleNamespace(chat_completions=native_sync, achat_completions=native_async), - ) - endpoint: Final[runtime.EndpointDispatch[object, object]] = runtime.EndpointDispatch.native( - route="test", - sync=lambda native: native.chat_completions, - asynchronous=lambda native: native.achat_completions, - enabled=enabled, - ) - - assert endpoint.sync.load() is native_sync - assert endpoint.asynchronous.load() is native_async - endpoint.override(sync=replacement_sync) - assert endpoint.sync.load() is replacement_sync - assert endpoint.asynchronous.load() is native_async - endpoint.override(asynchronous=None) - assert endpoint.sync.load() is replacement_sync - assert endpoint.asynchronous.load() is None - endpoint.reset() - assert endpoint.sync.load() is native_sync - assert endpoint.asynchronous.load() is native_async - - -def test_direct_endpoint_binding_rejects_native_state_controls() -> None: - endpoint: Final = runtime.EndpointBinding(route="test", load=object, enabled=enabled) - - with pytest.raises(RuntimeError, match="only native Rust bridges support binding overrides"): - endpoint.override(object()) - with pytest.raises(RuntimeError, match="only native Rust bridges support binding resets"): - endpoint.reset() - - -@pytest.mark.parametrize( - ("enabled_state", "reason", "expected"), - ( - pytest.param(False, None, runtime.PythonFallbackReason.NATIVE_DISABLED, id="disabled"), - pytest.param(True, "unsupported model", runtime.PythonFallbackReason.NATIVE_DECLINED, id="declined"), - pytest.param(True, None, None, id="accepted"), - ), -) -def test_assess_reports_binding_eligibility( - enabled_state: bool, - reason: str | None, - expected: runtime.PythonFallbackReason | None, -) -> None: - binding: Final = object() - checked: list[object] = [] - endpoint: Final = runtime.EndpointBinding(route="test", load=lambda: binding, enabled=lambda: enabled_state) - - result: Final = endpoint.assess(check=lambda value: checked.append(value) or reason) - - assert (result.reason if result is not None else None) is expected - assert (result.detail if result is not None else None) == reason - assert checked == ([binding] if enabled_state else []) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -async def test_dispatch_require_returns_adapted_native_success_without_exception_metadata( - monkeypatch: pytest.MonkeyPatch, - asynchronous: bool, -) -> None: - monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) - endpoint: Final = runtime.EndpointDispatch( - sync=runtime.EndpointBinding(route="test", load=object, enabled=enabled), - asynchronous=runtime.EndpointBinding(route="test", load=object, enabled=enabled), - ) - - async def acall(_binding: object, request: int) -> int: - return request * 2 - - result: Final = ( - await endpoint.arequire(prepare=lambda: 3, call=acall, adapt=str, error_context=context()) - if asynchronous - else endpoint.require( - prepare=lambda: 3, - call=lambda _binding, request: request * 2, - adapt=str, - error_context=context(), - ) - ) - - assert result == "6" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -async def test_response_adaptation_failure_never_authorizes_fallback(asynchronous: bool) -> None: - def adapt(value: str) -> str: - assert value == "provider response" - raise RustBridgeDeclined("adapter failed after provider response") - - async def native(binding: object, request: object) -> str: - return "provider response" - - async def fallback() -> str: - pytest.fail("a received response must not be retried") - - bridge = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) - - async def invoke() -> None: - if asynchronous: - await bridge.ainvoke( - prepare=lambda: None, call=native, fallback=fallback, adapt=adapt, error_context=context() - ) - else: - bridge.invoke( - prepare=lambda: None, - call=lambda binding, request: "provider response", - fallback=lambda: pytest.fail("a received response must not be retried"), - adapt=adapt, - error_context=context(), - ) - - with pytest.raises(RustBridgeDeclined, match="adapter failed"): - await invoke() - - -@pytest.mark.parametrize("error", (RustBridgeDeclined("unsupported"), RustUpstreamError(429, "rate limited"))) -def test_propagate_policy_preserves_native_errors(error: Exception) -> None: - def fail(_binding: object, _request: object) -> object: - raise error - - bridge: Final = runtime.EndpointBinding( - route="ocr", load=object, enabled=enabled, error_policy=runtime.NativeErrorPolicy.PROPAGATE - ) - - with pytest.raises(type(error)) as caught: - bridge.invoke( - prepare=lambda: None, - call=fail, - fallback=lambda: pytest.fail("fallback must not run"), - adapt=str, - error_context=context(), - ) - - assert caught.value is error - - -@pytest.mark.asyncio -@pytest.mark.parametrize("error", (RustBridgeDeclined("unsupported"), RustUpstreamError(429, "rate limited"))) -async def test_async_propagate_policy_preserves_native_errors(error: Exception) -> None: - async def fail(_binding: object, _request: object) -> object: - raise error - - async def fallback() -> str: - pytest.fail("fallback must not run") - - bridge: Final = runtime.EndpointBinding( - route="messages", load=object, enabled=enabled, error_policy=runtime.NativeErrorPolicy.PROPAGATE - ) - - with pytest.raises(type(error)) as caught: - await bridge.ainvoke(prepare=lambda: None, call=fail, fallback=fallback, adapt=str, error_context=context()) - - assert caught.value is error - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -@pytest.mark.parametrize("available, accepted", ((False, False), (True, False), (True, True))) -async def test_preflight_runs_after_binding_selection_before_preparation( - asynchronous: bool, available: bool, accepted: bool -) -> None: - events: list[str] = [] - - def load() -> object | None: - events.append("load") - return object() if available else None - - def preflight() -> runtime.PythonFallback | None: - events.append("preflight") - return None if accepted else runtime.PythonFallback(runtime.PythonFallbackReason.NATIVE_DECLINED) - def prepare() -> int: events.append("prepare") - return 7 + return 3 - def call(binding: object, request: int) -> int: - events.append("native") - return request + def call(_binding: object, request: int) -> int: + events.append("call") + return request * 2 async def acall(binding: object, request: int) -> int: return call(binding, request) - def fallback() -> str: - events.append("python") - return "3" + def adapt(value: int) -> str: + events.append("adapt") + return str(value) - async def afallback() -> str: - return fallback() - - endpoint: Final = runtime.EndpointBinding(route="ocr", load=load, enabled=enabled) result: Final = ( - await endpoint.ainvoke( - prepare=prepare, call=acall, fallback=afallback, adapt=str, error_context=context(), preflight=preflight + await runtime.aattempt( + load=load, + enabled=state != "disabled", + eligible=state != "ineligible", + prepare=prepare, + call=acall, + adapt=adapt, ) if asynchronous - else endpoint.invoke( - prepare=prepare, call=call, fallback=fallback, adapt=str, error_context=context(), preflight=preflight + else runtime.attempt( + load=load, + enabled=state != "disabled", + eligible=state != "ineligible", + prepare=prepare, + call=call, + adapt=adapt, ) ) - assert result == ("7" if available and accepted else "3") - assert events == ( - ["load", "preflight", "prepare", "native"] - if available and accepted - else ["load", "preflight", "python"] - if available - else ["load", "python"] + if state == "handled": + assert result == runtime.Handled("6") + assert events == ["load", "prepare", "call", "adapt"] + else: + assert result == runtime.NativeSkipped(runtime.NativeSkipReason(state)) + assert events == (["load"] if state == "unavailable" else []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("phase", ("prepare", "call")) +async def test_attempt_reports_failure_without_deciding_retry(asynchronous: bool, phase: str) -> None: + error: Final = RuntimeError("native failure") + + def prepare() -> int: + if phase == "prepare": + raise error + return 3 + + def call(_binding: object, request: int) -> int: + raise error + + async def acall(binding: object, request: int) -> int: + return call(binding, request) + + def adapt(value: int) -> str: + pytest.fail("failed attempts cannot be adapted") + + result: Final = ( + await runtime.aattempt(load=object, enabled=True, eligible=True, prepare=prepare, call=acall, adapt=adapt) + if asynchronous + else runtime.attempt(load=object, enabled=True, eligible=True, prepare=prepare, call=call, adapt=adapt) ) + assert isinstance(result, runtime.NativeFailed) + assert result.error is error -def test_preflight_failure_is_not_a_native_decline() -> None: - endpoint: Final = runtime.EndpointBinding(route="ocr", load=object, enabled=enabled) +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_adaptation_failure_remains_distinct_from_native_failure(asynchronous: bool) -> None: + error: Final = ValueError("invalid response") - def preflight() -> runtime.PythonFallback | None: - raise ValueError("invalid acceptance contract") + async def acall(_binding: object, request: int) -> int: + return request - with pytest.raises(ValueError, match="invalid acceptance contract"): - endpoint.invoke( - prepare=lambda: pytest.fail("must not prepare"), - call=lambda binding, request: pytest.fail("must not invoke"), - fallback=lambda: pytest.fail("must not fall back"), - adapt=str, - error_context=context(), - preflight=preflight, - ) + def adapt(value: int) -> str: + raise error + + async def run() -> None: + if asynchronous: + await runtime.aattempt(load=object, enabled=True, eligible=True, prepare=lambda: 3, call=acall, adapt=adapt) + else: + runtime.attempt( + load=object, + enabled=True, + eligible=True, + prepare=lambda: 3, + call=lambda binding, request: request, + adapt=adapt, + ) + + with pytest.raises(ValueError, match="invalid response") as caught: + await run() + assert caught.value is error diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index 1b0fcfacd1d..0cbe0bf5277 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -4,6 +4,7 @@ import pytest import litellm from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch +from litellm.rust_bridge.runtime import Handled rust_bridge = importlib.import_module("litellm.rust_bridge.transcription") @@ -55,7 +56,8 @@ def test_enabled_sync_bridge_receives_audio() -> None: optional_params={"temperature": 0}, timeout=5.0, ) - assert result == {"text": "hello"} + assert isinstance(result, Handled) + assert result.value == {"text": "hello"} assert bridge.calls[0]["audio"] == {"data": "AQI=", "format": "wav", "filename": "audio.wav"} @@ -72,7 +74,7 @@ async def test_enabled_async_bridge() -> None: optional_params={}, timeout=None, ) - assert result == {"text": "async"} + assert result == Handled({"text": "async"}) def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None: @@ -83,7 +85,8 @@ def test_loader_returns_none_without_native_extension(monkeypatch: pytest.Monkey def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None) + rust_bridge.configure_rust_transcription(transcription=None) + monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None) with pytest.raises(RuntimeError, match="bridge is unavailable"): BedrockAudioTranscriptionRustDispatch().audio_transcriptions( @@ -100,10 +103,8 @@ def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> @pytest.mark.asyncio async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: - async def unavailable(**_: object) -> None: - return None - - monkeypatch.setattr(rust_bridge, "atranscription", unavailable) + rust_bridge.configure_rust_transcription(atranscription=None) + monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None) with pytest.raises(RuntimeError, match="bridge is unavailable"): await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions( diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 1c5671c0b57..144732cfd47 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22173 + "limit": 22165 }, "LIT002": { "limit": 26729 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1027 + "limit": 1022 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16426 + "limit": 16419 }, "LIT011": { "limit": 5506