diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index 87847953a8f..b1f8c957ff4 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -1,38 +1,14 @@ import base64 -from io import IOBase -from typing import Final, NoReturn +from typing import Final import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file -from litellm.litellm_core_utils.litellm_logging import Logging from litellm.rust_bridge import transcription as rust_transcription_bridge -from litellm.rust_bridge.dispatch import anative_first, native_first, provider_errors -from litellm.rust_bridge.request import request_context -from litellm.rust_bridge.runtime import DispatchResult, adapt_result 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 _input_source_kind(audio_file: FileTypes) -> str: - content: Final = audio_file[1] if isinstance(audio_file, tuple) else audio_file - if isinstance(content, (bytes, bytearray, memoryview)): - return "bytes" - if isinstance(content, IOBase): - return "file" - if isinstance(content, str): - return "path" - return "opaque" - @staticmethod def _audio_payload(audio_file: FileTypes) -> dict[str, object]: processed_audio: Final = process_audio_file(audio_file) @@ -55,44 +31,6 @@ class BedrockAudioTranscriptionRustDispatch: "filename": processed_audio.filename, } - def _attempt_audio_transcriptions( - self, - *, - model: str, - audio_file: FileTypes, - api_key: str | None, - api_base: str | None, - custom_llm_provider: str, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, - logging_obj: Logging | None = None, - ) -> DispatchResult[TranscriptionResponse]: - result: 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, - input_source_kind=self._input_source_kind(audio_file), - context=request_context( - logging_obj=logging_obj, - request_model=logging_obj.model if logging_obj is not None else model, - litellm_params=logging_obj.litellm_params if logging_obj is not None else None, - ), - ) - return adapt_result(result, lambda response: TranscriptionResponse(**response)) - - @native_first( - native=_attempt_audio_transcriptions, - route="audio transcription", - errors=lambda self, model, audio_file, api_key, api_base, custom_llm_provider, extra_headers, optional_params, timeout, logging_obj=None: ( - provider_errors(custom_llm_provider, model) - ), - ) def audio_transcriptions( self, *, @@ -104,24 +42,8 @@ class BedrockAudioTranscriptionRustDispatch: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, - logging_obj: Logging | None = None, ) -> TranscriptionResponse: - _unavailable() - - async def _attempt_async_audio_transcriptions( - self, - *, - model: str, - audio_file: FileTypes, - api_key: str | None, - api_base: str | None, - custom_llm_provider: str, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, - logging_obj: Logging | None = None, - ) -> DispatchResult[TranscriptionResponse]: - result: Final = await rust_transcription_bridge.atranscription( + rust_response: Final = rust_transcription_bridge.transcription( model=model, audio=self._audio_payload(audio_file), api_key=api_key, @@ -130,22 +52,11 @@ class BedrockAudioTranscriptionRustDispatch: extra_headers=extra_headers, optional_params=optional_params, timeout=timeout, - input_source_kind=self._input_source_kind(audio_file), - context=request_context( - logging_obj=logging_obj, - request_model=logging_obj.model if logging_obj is not None else model, - litellm_params=logging_obj.litellm_params if logging_obj is not None else None, - ), ) - return adapt_result(result, lambda response: TranscriptionResponse(**response)) + if rust_response is None: + raise RuntimeError("Rust audio transcription bridge is unavailable") + return TranscriptionResponse(**rust_response) - @anative_first( - native=_attempt_async_audio_transcriptions, - route="audio transcription", - errors=lambda self, model, audio_file, api_key, api_base, custom_llm_provider, extra_headers, optional_params, timeout, logging_obj=None: ( - provider_errors(custom_llm_provider, model) - ), - ) async def async_audio_transcriptions( self, *, @@ -157,6 +68,17 @@ class BedrockAudioTranscriptionRustDispatch: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, - logging_obj: Logging | None = None, ) -> TranscriptionResponse: - await _aunavailable() + 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, + ) + if rust_response is None: + raise RuntimeError("Rust audio transcription bridge is unavailable") + return TranscriptionResponse(**rust_response) diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index d687fb6ef0a..4dd93f6e998 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -10,7 +10,8 @@ from litellm.rust_bridge.bindings import ( UNCHANGED, NativeBinding, Unchanged, - native_exception_types, + native_declined_types, + native_upstream_types, ) from litellm.rust_bridge.protocols import NativeModule, RustRouteDecline from litellm.rust_bridge.request import NativeRequestCapabilities, NativeRequestContext @@ -32,6 +33,14 @@ class PythonFallbackReason(Enum): NATIVE_DECLINED = "native_declined" +class NativeSkipReason(Enum): + DISABLED = "disabled" + INELIGIBLE = "ineligible" + UNAVAILABLE = "unavailable" + DECLINED = "declined" + FAILED = "failed" + + @dataclass(frozen=True, slots=True) class Handled(Generic[ResultT]): value: ResultT @@ -43,7 +52,18 @@ class PythonFallback: detail: str | None = None -DispatchResult: TypeAlias = Handled[ResultT] | PythonFallback +@dataclass(frozen=True, slots=True) +class NativeSkipped: + reason: NativeSkipReason + detail: str | None = None + + +@dataclass(frozen=True, slots=True) +class NativeFailed: + error: Exception + + +DispatchResult: TypeAlias = Handled[ResultT] | PythonFallback | NativeSkipped | NativeFailed @dataclass(frozen=True, slots=True) @@ -296,10 +316,10 @@ class EndpointBinding(Generic[BindingT]): adapt: Callable[[NativeT], ResultT], error_context: BridgeErrorContext, ) -> DispatchResult[ResultT]: - exceptions: Final = native_exception_types() - if exceptions is None: + declined: Final = native_declined_types() + upstream: Final = native_upstream_types() + if not declined or not upstream: return Handled(adapt(call())) - declined, upstream = exceptions try: value: Final = call() except declined as error: @@ -315,10 +335,10 @@ class EndpointBinding(Generic[BindingT]): adapt: Callable[[NativeT], ResultT], error_context: BridgeErrorContext, ) -> DispatchResult[ResultT]: - exceptions: Final = native_exception_types() - if exceptions is None: + declined: Final = native_declined_types() + upstream: Final = native_upstream_types() + if not declined or not upstream: return Handled(adapt(await call())) - declined, upstream = exceptions try: value: Final = await call() except declined as error: @@ -505,6 +525,12 @@ def identity(value: ResultT) -> ResultT: return value +def adapt_result(result: DispatchResult[NativeT], adapt: Callable[[NativeT], ResultT]) -> DispatchResult[ResultT]: + if isinstance(result, Handled): + return Handled(adapt(result.value)) + return result + + async def async_none() -> None: return None diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 3a10f60624d..ffdcdeec8b3 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -57,7 +57,6 @@ _PREFLIGHT: Final[EndpointBinding[RustRouteDecline]] = EndpointBinding.native( def configure_rust_transcription( - enabled: bool = True, *, transcription: RustTranscription | None | Unchanged = UNCHANGED, atranscription: RustAtranscription | None | Unchanged = UNCHANGED, diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 9e0de8e6157..236b15a55e6 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -12,8 +12,13 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import configuration -from litellm.rust_bridge.request import NativeOCRRequest, NativeRequestContext, NativeRequestOptions -from litellm.rust_bridge.runtime import Handled +from litellm.rust_bridge.request import ( + NativeOCRRequest, + NativeRequestContext, + NativeRequestOptions, + NativeVertexOptions, + PreparedNativeCall, +) from litellm.rust_bridge.timeouts import timeout_to_seconds # `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` @@ -47,7 +52,6 @@ class RecordingBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] - self.contexts: list[NativeRequestContext] = [] def __call__( self, @@ -65,10 +69,10 @@ class RecordingBridge: "custom_llm_provider": options.custom_llm_provider, "extra_headers": options.extra_headers, "optional_params": request.optional_params, + "vertex": options.vertex, "timeout_seconds": options.timeout_seconds, } ) - self.contexts.append(context) return dict(FAKE_OCR_RESPONSE) @@ -77,7 +81,6 @@ class RecordingAsyncBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] - self.contexts: list[NativeRequestContext] = [] async def __call__( self, @@ -95,10 +98,10 @@ class RecordingAsyncBridge: "custom_llm_provider": options.custom_llm_provider, "extra_headers": options.extra_headers, "optional_params": request.optional_params, + "vertex": options.vertex, "timeout_seconds": options.timeout_seconds, } ) - self.contexts.append(context) return dict(FAKE_OCR_RESPONSE) @@ -192,7 +195,7 @@ def build_prepared_request( litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = 12.5, ) -> Any: - return rust_bridge.PreparedOCRRequest( + return ocr_main._PreparedOCRRequest( model=model, document=document, api_key=api_key, @@ -397,13 +400,105 @@ 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: PreparedNativeCall( + request=NativeOCRRequest( + model="mistral-ocr-latest", + document=DOCUMENT, + optional_params={"include_image_base64": True, "pages": [0]}, + ), + options=NativeRequestOptions( + api_key="sk-test", + api_base="https://proxy.internal", + custom_llm_provider="mistral", + extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, + timeout_seconds=12.5, + ), + ), + 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]}, + "vertex": None, + "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: PreparedNativeCall( + request=NativeOCRRequest( + model="mistral-ocr-maas", + document=DOCUMENT, + optional_params={}, + ), + options=NativeRequestOptions( + custom_llm_provider="vertex_ai", + vertex=NativeVertexOptions(project="project-1"), + timeout_seconds=42.0, + ), + ), + 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": NativeVertexOptions(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 = rust_bridge.attempt_ocr( + response = ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://proxy.internal", @@ -414,13 +509,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.contexts[0].capabilities.execution_mode == "sync" - assert bridge.contexts[0].capabilities.input_source_kind == "document_url" - assert bridge.contexts[0].capabilities.native_response_format is False assert bridge.calls[0] == { "model": "mistral-ocr-latest", "document": DOCUMENT, @@ -432,6 +522,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): "x-trace-id": "trace-1", }, "optional_params": {"include_image_base64": True}, + "vertex": NativeVertexOptions(), "timeout_seconds": 12.5, } @@ -441,7 +532,8 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - rust_bridge.attempt_ocr( + ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), 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, ) @@ -457,7 +549,8 @@ 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}") - rust_bridge.attempt_ocr( + ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( api_key="sk-explicit", timeout=None, @@ -478,7 +571,8 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): resolver_calls.append(name) return "sk-provider-env" - rust_bridge.attempt_ocr( + ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), model="provider-ocr-model", @@ -497,7 +591,8 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - rust_bridge.attempt_ocr( + ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -512,11 +607,8 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): resolve_api_key=lambda _name: None, ) - assert bridge.calls[0]["optional_params"] == { - "include_image_base64": True, - "vertex_project": "project-1", - "vertex_location": "us-central1", - } + assert bridge.calls[0]["optional_params"] == {"include_image_base64": True} + assert bridge.calls[0]["vertex"] == NativeVertexOptions(project="project-1", location="us-central1") def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): @@ -530,7 +622,8 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana "VERTEXAI_LOCATION": "us-east5", }.get(name) - rust_bridge.attempt_ocr( + ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -539,8 +632,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana resolve_api_key=_resolver, ) - assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" - assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" + assert bridge.calls[0]["vertex"] == NativeVertexOptions(project="project-from-secret", location="us-east5") def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): @@ -548,7 +640,8 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - rust_bridge.attempt_ocr( + ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( custom_llm_provider="azure_ai", model="pixtral-12b-2409", @@ -566,7 +659,8 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - rust_bridge.attempt_ocr( + ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( custom_llm_provider="azure_ai", model="doc-intelligence/prebuilt-layout", @@ -587,7 +681,8 @@ def test_run_rust_ocr_runs_pre_call_logging(): litellm.rust(True) rust_bridge.set_rust_ocr(ocr=bridge) - rust_bridge.attempt_ocr( + ocr_main._run_rust_ocr( + fallback=lambda: pytest.fail("unexpected Python fallback"), prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://api.mistral.ai/v1", @@ -734,7 +829,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(rust_bridge, "_prepare_rust_ocr_call", unexpected_preparation) + monkeypatch.setattr(ocr_main, "_prepare_rust_ocr_call", unexpected_preparation) monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fallback) response: Final = ( @@ -747,25 +842,6 @@ async def test_ocr_fallback_skips_native_preparation( fallback.assert_called_once() -@pytest.mark.asyncio -async def test_aocr_rejects_empty_python_fallback_response(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict[str, object] = {} - - def fake_exception_type(**kwargs: object) -> CapturedException: - captured.update(kwargs) - return CapturedException("wrapped") - - monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) - monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", AsyncMock(return_value=None)) - - with pytest.raises(CapturedException, match="wrapped"): - await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - original: Final = captured["original_exception"] - assert isinstance(original, ValueError) - assert str(original) == "Got an unexpected None response from the OCR API: None" - - def test_ocr_provider_configs_expose_api_key_env_vars(): from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( AzureDocumentIntelligenceOCRConfig, diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index a882e23fb58..4b6e5c1290b 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -1,117 +1,460 @@ from __future__ import annotations +from dataclasses import dataclass +from types import SimpleNamespace from typing import Final import pytest -from litellm.rust_bridge import runtime +from litellm.exceptions import APIError, AuthenticationError, InternalServerError, RateLimitError +from litellm.rust_bridge import bindings, 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"), AuthenticationError, 401, "unauthorized", id="auth"), + pytest.param(RustUpstreamError(429, "rate limited"), RateLimitError, 429, "rate limited", id="rate-limit"), + pytest.param(RustUpstreamError(500, "failed"), InternalServerError, 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)) -@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]] = [] +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.parametrize( + ("state", "expected", "expected_events"), + ( + pytest.param("disabled", False, (), id="disabled"), + pytest.param("ineligible", False, (), id="ineligible"), + pytest.param("unavailable", False, ("load",), id="unavailable"), + pytest.param("available", True, ("load",), id="available"), + ), +) +def test_can_attempt_only_enabled_available_requests( + state: str, + expected: bool, + expected_events: tuple[str, ...], +) -> None: + events: list[str] = [] def load() -> object | None: events.append("load") return None if state == "unavailable" else object() - def prepare() -> int: - events.append("prepare") - return 3 + bridge: Final = runtime.EndpointBinding(route="messages", load=load, enabled=lambda: state != "disabled") - 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 adapt(value: int) -> str: - events.append("adapt") - return str(value) - - result: Final = ( - await runtime.aattempt( - load=load, - enabled=state != "disabled", + assert ( + bridge.can_attempt( eligible=state != "ineligible", - prepare=prepare, - call=acall, - adapt=adapt, - ) - if asynchronous - else runtime.attempt( - load=load, - enabled=state != "disabled", - eligible=state != "ineligible", - prepare=prepare, - call=call, - adapt=adapt, ) + is expected ) - 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 []) + assert tuple(events) == expected_events + + +def test_native_endpoint_applies_partial_overrides_and_reset(monkeypatch: pytest.MonkeyPatch) -> None: + def native_sync() -> str: + return "native" + + async def native_async() -> str: + return "native async" + + def replacement_sync() -> str: + return "replacement" + + monkeypatch.setattr( + bindings, + "get_native_bridge", + lambda: SimpleNamespace(chat_completions=native_sync, achat_completions=native_async), + ) + endpoint: Final[runtime.EndpointDispatch[object, object]] = runtime.EndpointDispatch.native( + route="test", + sync=lambda native: native.chat_completions, + asynchronous=lambda native: native.achat_completions, + enabled=enabled, + ) + + assert endpoint.sync.load() is native_sync + assert endpoint.asynchronous.load() is native_async + endpoint.override(sync=replacement_sync) + assert endpoint.sync.load() is replacement_sync + assert endpoint.asynchronous.load() is native_async + endpoint.override(asynchronous=None) + assert endpoint.sync.load() is replacement_sync + assert endpoint.asynchronous.load() is None + endpoint.reset() + assert endpoint.sync.load() is native_sync + assert endpoint.asynchronous.load() is native_async @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", (False, True)) -@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") +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") - def prepare() -> int: - if phase == "prepare": - raise error - return 3 + async def native(binding: object, request: object) -> str: + return "provider response" - def call(_binding: object, request: int) -> int: - raise error + async def fallback() -> str: + pytest.fail("a received response must not be retried") - async def acall(binding: object, request: int) -> int: - return call(binding, request) + bridge = runtime.EndpointBinding(route="messages", load=object, enabled=enabled) - 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 - - -@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") - - async def acall(_binding: object, request: int) -> int: - return request - - def adapt(value: int) -> str: - raise error - - async def run() -> None: + async def invoke() -> None: if asynchronous: - await runtime.aattempt(load=object, enabled=True, eligible=True, prepare=lambda: 3, call=acall, adapt=adapt) + await bridge.ainvoke( + prepare=lambda: None, call=native, fallback=fallback, adapt=adapt, error_context=context() + ) else: - runtime.attempt( - load=object, - enabled=True, - eligible=True, - prepare=lambda: 3, - call=lambda binding, request: request, + 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(ValueError, match="invalid response") as caught: - await run() - assert caught.value is error + with pytest.raises(RustBridgeDeclined, match="adapter failed"): + await invoke() + + +@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 + + def call(binding: object, request: int) -> int: + events.append("native") + return request + + async def acall(binding: object, request: int) -> int: + return call(binding, request) + + def fallback() -> str: + events.append("python") + return "3" + + 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 + ) + if asynchronous + else endpoint.invoke( + prepare=prepare, call=call, fallback=fallback, adapt=str, error_context=context(), preflight=preflight + ) + ) + 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"] + ) + + +def test_preflight_failure_is_not_a_native_decline() -> None: + endpoint: Final = runtime.EndpointBinding(route="ocr", load=object, enabled=enabled) + + def preflight() -> runtime.PythonFallback | None: + raise ValueError("invalid acceptance contract") + + 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, + ) diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index c7864b1319a..85303fed5ce 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -4,8 +4,11 @@ import pytest import litellm from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch -from litellm.rust_bridge.request import NativeRequestContext, NativeRequestOptions, NativeTranscriptionRequest -from litellm.rust_bridge.runtime import Handled +from litellm.rust_bridge.request import ( + NativeRequestContext, + NativeRequestOptions, + NativeTranscriptionRequest, +) rust_bridge = importlib.import_module("litellm.rust_bridge.transcription") @@ -24,7 +27,6 @@ def reset_rust_transcription() -> None: class SyncBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] - self.contexts: list[NativeRequestContext] = [] def __call__( self, @@ -34,9 +36,13 @@ class SyncBridge: context: NativeRequestContext, ) -> dict[str, object]: self.calls.append( - {"model": request.model, "audio": request.audio, "optional_params": request.optional_params} + { + "model": request.model, + "audio": request.audio, + "optional_params": request.optional_params, + "bedrock": options.bedrock, + } ) - self.contexts.append(context) return {"text": "hello"} @@ -63,17 +69,9 @@ def test_enabled_sync_bridge_receives_audio() -> None: extra_headers=None, optional_params={"temperature": 0}, timeout=5.0, - stream=True, - has_custom_client=True, - input_source_kind="file", ) - assert isinstance(result, Handled) - assert result.value == {"text": "hello"} + assert result == {"text": "hello"} assert bridge.calls[0]["audio"] == {"data": "AQI=", "format": "wav", "filename": "audio.wav"} - assert bridge.contexts[0].capabilities.execution_mode == "sync" - assert bridge.contexts[0].capabilities.stream is True - assert bridge.contexts[0].capabilities.has_custom_client is True - assert bridge.contexts[0].capabilities.input_source_kind == "file" @pytest.mark.asyncio @@ -89,7 +87,7 @@ async def test_enabled_async_bridge() -> None: optional_params={}, timeout=None, ) - assert result == Handled({"text": "async"}) + assert result == {"text": "async"} def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None: @@ -100,8 +98,7 @@ def test_loader_returns_none_without_native_extension(monkeypatch: pytest.Monkey def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: - rust_bridge.configure_rust_transcription(transcription=None) - monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None) + monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None) with pytest.raises(RuntimeError, match="bridge is unavailable"): BedrockAudioTranscriptionRustDispatch().audio_transcriptions( @@ -118,8 +115,10 @@ 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: - rust_bridge.configure_rust_transcription(atranscription=None) - monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: None) + async def unavailable(**_: object) -> None: + return None + + monkeypatch.setattr(rust_bridge, "atranscription", unavailable) with pytest.raises(RuntimeError, match="bridge is unavailable"): await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions( @@ -136,7 +135,7 @@ async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPat def test_bedrock_transcription_uses_rust_only_path() -> None: rust_bridge.configure_rust_transcription( - transcription=lambda *_args, **_: {"text": "rust"}, + transcription=lambda request, *, options, context: {"text": "rust"}, atranscription=None, ) try: @@ -152,7 +151,9 @@ def test_bedrock_transcription_uses_rust_only_path() -> None: @pytest.mark.asyncio async def test_bedrock_atranscription_uses_rust_only_path() -> None: - async def rust_response(*_args: object, **_: object) -> dict[str, object]: + async def rust_response( + request: NativeTranscriptionRequest, *, options: object, context: NativeRequestContext + ) -> dict[str, object]: return {"text": "rust"} rust_bridge.configure_rust_transcription(transcription=None, atranscription=rust_response)