diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index 35007003150..4820c393bd7 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -35,8 +35,10 @@ pub async fn ocr( .await } -pub fn ocr_provider_supported(model: &str, provider: &str) -> bool { - common_utils::ocr_provider_config(provider, model).is_some() +pub fn ocr_provider_supported(model: &str, provider: &str, request_format: Option<&str>) -> bool { + common_utils::ocr_provider_config(provider, model).is_some_and(|config| { + request_format != Some("native") || config.supported_ocr_params().contains(&"req_format") + }) } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/eligibility.rs b/litellm-rust/crates/core/src/eligibility.rs index febd25d7e57..f28e5c73b07 100644 --- a/litellm-rust/crates/core/src/eligibility.rs +++ b/litellm-rust/crates/core/src/eligibility.rs @@ -6,7 +6,6 @@ pub enum NativeRouteDecline { Streaming, AgenticHook, CustomClient, - NativeResponseFormat, } impl NativeRouteDecline { @@ -16,7 +15,6 @@ impl NativeRouteDecline { Self::Streaming => "native streaming is unavailable", Self::AgenticHook => "native agentic hooks are unavailable", Self::CustomClient => "native custom clients are unavailable", - Self::NativeResponseFormat => "native OCR response format is unavailable", } } } @@ -37,8 +35,7 @@ pub fn native_route_decline( if capabilities.has_custom_client { return Some(NativeRouteDecline::CustomClient); } - (capabilities.request_format.as_deref() == Some("native")) - .then_some(NativeRouteDecline::NativeResponseFormat) + None } #[cfg(test)] @@ -52,6 +49,7 @@ mod tests { has_agentic_hook: true, has_custom_client: true, request_format: Some("native".into()), + ..Default::default() }; assert_eq!( native_route_decline(false, &all_unsupported), @@ -84,13 +82,6 @@ mod tests { }, NativeRouteDecline::CustomClient, ), - ( - RequestCapabilities { - request_format: Some("native".into()), - ..Default::default() - }, - NativeRouteDecline::NativeResponseFormat, - ), ]; for (capabilities, expected) in cases { assert_eq!(native_route_decline(true, &capabilities), Some(expected)); diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 899969b5787..c2350845cbd 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -332,6 +332,11 @@ litellm_context = replace( ) assert routes.ocr_decline('model', 'mistral', context=native_context) is not None assert routes.ocr_decline('model', 'mistral', context=litellm_context) is None +assert routes.ocr_decline( + 'doc-intelligence/prebuilt-layout', + 'azure_ai', + context=native_context, +) is None ", Some(&locals), Some(&locals), diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 139f0ba48b4..1ed77211c00 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -23,11 +23,12 @@ fn prepare_ocr( options: NativeRequestOptions, context: NativeRequestContext, ) -> PyResult> + Send + 'static> { + let context: LiteLlmRequestContext = context.into(); let provider_supported = litellm_ai_gateway::io::ocr::ocr_provider_supported( &input.model, options.provider("mistral"), + context.capabilities.request_format.as_deref(), ); - let context: LiteLlmRequestContext = context.into(); if let Some(reason) = super::definition::request_decline(provider_supported, &context) { return Err(crate::errors::RustBridgeDeclined::new_err(reason)); } @@ -58,10 +59,12 @@ fn ocr_decline( context: NativeRequestContext, ) -> Option { let context: LiteLlmRequestContext = context.into(); - super::definition::request_decline( - litellm_ai_gateway::io::ocr::ocr_provider_supported(model, custom_llm_provider), - &context, - ) + let provider_supported = litellm_ai_gateway::io::ocr::ocr_provider_supported( + model, + custom_llm_provider, + context.capabilities.request_format.as_deref(), + ); + super::definition::request_decline(provider_supported, &context) } bridge_route! { diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index d68267d04b1..5797f77bd7f 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -467,7 +467,7 @@ class AnthropicChatCompletion(BaseLLM): speed=optional_params.get("speed") if optional_params else None, tool_name_reverse_map=( litellm_params.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY) - if isinstance(litellm_params, dict) + if isinstance(litellm_params, dict) # pyright: ignore[reportUnnecessaryIsInstance] # runtime callers can still supply non-dict values else None ), ) @@ -482,8 +482,6 @@ class AnthropicChatCompletion(BaseLLM): else: if client is None or not isinstance(client, HTTPHandler): client = _get_httpx_client(params={"timeout": timeout}) - else: - client = client try: response: Final = client.post( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 60ede52b84e..d8258721683 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -316,7 +316,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return self.custom_llm_provider or "anthropic" @classmethod - def get_config_for_model(cls, model: str) -> dict[str, object]: + def get_config_for_model( + cls, model: str + ) -> dict[str, object]: # mutable-ok: callers merge provider defaults into a request copy from pydantic import TypeAdapter return TypeAdapter(dict[str, object]).validate_python(cls.get_config(model=model)) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 8afc1bdee01..9360f4fbf14 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -693,7 +693,9 @@ def anthropic_messages_handler( ) -def _native_messages_body(params: dict[str, object], model: str, drop_params: bool, provider: str) -> dict[str, object]: +def _native_messages_body( # mutable-ok: provider mapping consumes and returns an owned request body + params: dict[str, object], model: str, drop_params: bool, provider: str +) -> dict[str, object]: from pydantic import TypeAdapter requested: Final = TypeAdapter(dict[str, object]).validate_python( diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 507285d17d5..453a50c7e19 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -71,7 +71,9 @@ class BaseConfig(ABC): pass @classmethod - def get_config_for_model(cls, model: str) -> dict[str, object]: + def get_config_for_model( + cls, model: str + ) -> dict[str, object]: # mutable-ok: callers merge provider defaults into a request copy return TypeAdapter(dict[str, object]).validate_python(cls.get_config()) @classmethod diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 27058ad35bc..18472092eeb 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -442,12 +442,10 @@ class BedrockConverseLLM(BaseAWSLLM): if client is None or isinstance(client, AsyncHTTPHandler): _params: Final = {} if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, int): + if isinstance(timeout, (float, int)): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout client = _get_httpx_client(_params) - else: - client = client if stream is not None and stream is True: completion_stream, response_headers = make_sync_call( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f2319e3acbd..5f7e101b9e8 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -230,7 +230,9 @@ def _responses_api_optional_request_param_names() -> frozenset[str]: return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys()) -def _custom_logger_callbacks(logging_obj: LiteLLMLoggingObj | None) -> list["CustomLogger"]: +def _custom_logger_callbacks( + logging_obj: LiteLLMLoggingObj | None, +) -> list["CustomLogger"]: # mutable-ok: callback runtime requires an owned ordered list from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import ( get_custom_logger_compatible_class, @@ -6397,6 +6399,7 @@ class BaseLLMHTTPHandler: }, ) + from litellm.rust_bridge.request import request_context from litellm.rust_bridge.responses_websocket import open_connection async with open_connection( @@ -6405,6 +6408,11 @@ class BaseLLMHTTPHandler: timeout=timeout, model=model, provider=custom_llm_provider, + context=request_context( + logging_obj=logging_obj, + request_model=logging_obj.model, + litellm_params=litellm_params.model_dump(), + ), fallback=lambda: websockets.connect( ws_url, additional_headers=headers, diff --git a/litellm/main.py b/litellm/main.py index 52ae8cdbd64..e6d678a6920 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7851,11 +7851,7 @@ def transcription( litellm_params=litellm_params_dict, model_response=model_response, atranscription=atranscription, - client=( - client - if client is not None and (isinstance(client, HTTPHandler) or isinstance(client, AsyncHTTPHandler)) - else None - ), + client=(client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None), timeout=timeout, max_retries=max_retries, logging_obj=litellm_logging_obj, @@ -7872,11 +7868,7 @@ def transcription( litellm_params=litellm_params_dict, model_response=model_response, atranscription=atranscription, - client=( - client - if client is not None and (isinstance(client, HTTPHandler) or isinstance(client, AsyncHTTPHandler)) - else None - ), + client=(client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None), timeout=timeout, max_retries=max_retries, logging_obj=litellm_logging_obj, diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index fe5f7e2536f..c35743a65fc 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -31,8 +31,10 @@ from litellm.llms.base_llm.ocr.transformation import ( from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge from litellm.rust_bridge.request import ( + NativeRequestCapabilities, NativeRequestOptions, PreparedNativeCall, + request_context, vertex_options, ) from litellm.rust_bridge.timeouts import timeout_to_seconds @@ -47,16 +49,17 @@ base_llm_http_handler = BaseLLMHTTPHandler() @dataclass class _PreparedOCRRequest: model: str - document: dict[str, Any] + document: dict[str, Any] # mutable-ok: public OCR document shape is a mutable SDK dictionary api_key: str | None api_base: str | None custom_llm_provider: str - extra_headers: dict[str, object] | None + extra_headers: dict[str, object] | None # mutable-ok: preserves the caller-owned SDK header contract provider_config: BaseOCRConfig - optional_params: dict[str, object] - litellm_params: dict[str, object] + optional_params: dict[str, object] # mutable-ok: provider mapping consumes an owned parameter copy + litellm_params: dict[str, object] # mutable-ok: preserves the existing SDK parameter contract effective_timeout: float | httpx.Timeout litellm_logging_obj: LiteLLMLoggingObj + execution_mode: str = "sync" def _prepare_ocr_request( @@ -68,6 +71,7 @@ def _prepare_ocr_request( custom_llm_provider: str | None, extra_headers: dict[str, object] | None, kwargs: dict[str, object], + execution_mode: str = "sync", ) -> _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)) @@ -177,13 +181,14 @@ def _prepare_ocr_request( litellm_params=dict(litellm_params), effective_timeout=effective_timeout, litellm_logging_obj=litellm_logging_obj, + execution_mode=execution_mode, ) def _rust_bridge_optional_params( prepared_request: _PreparedOCRRequest, resolve_secret: Callable[[str], str | None], -) -> dict[str, object]: +) -> dict[str, object]: # mutable-ok: returns an owned provider-parameter copy optional_params: Final = dict(prepared_request.optional_params) if prepared_request.custom_llm_provider == "vertex_ai": vertex_project: Final = ( @@ -272,6 +277,21 @@ def _prepare_rust_ocr_call( ), timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), ), + context=request_context( + logging_obj=prepared_request.litellm_logging_obj, + request_model=getattr(prepared_request.litellm_logging_obj, "model", prepared_request.model), + litellm_params=prepared_request.litellm_params, + capabilities=NativeRequestCapabilities( + execution_mode=prepared_request.execution_mode, + input_source_kind=str(prepared_request.document.get("type") or "unknown"), + request_format=( + value + if isinstance((value := prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM)), str) + else None + ), + native_response_format=(prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native"), + ), + ), ) @@ -422,6 +442,7 @@ async def aocr( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, kwargs=kwargs, + execution_mode="async", ) model = prepared.model custom_llm_provider = prepared.custom_llm_provider @@ -445,7 +466,7 @@ async def aocr( litellm_params=prepared.litellm_params, ) response: Final = await pending if asyncio.iscoroutine(pending) else pending - if response is None: + if response is None: # pyright: ignore[reportUnnecessaryComparison] # provider adapters can violate their declared return type raise ValueError(f"Got an unexpected None response from the OCR API: {response}") return response diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index 68a8100c0f7..1ee325a1e3c 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -47,6 +47,8 @@ from litellm.rust_bridge.request import ( anthropic_options, bedrock_options, call_native, + request_context, + with_capabilities, ) from litellm.rust_bridge.runtime import ( BridgeErrorContext, @@ -166,12 +168,14 @@ def _provider_eligibility_options( def _eligibility_context( *, + execution_mode: str | None = None, stream: bool, has_custom_client: bool = False, has_agentic_hook: bool = False, ) -> NativeRequestContext: return NativeRequestContext( capabilities=NativeRequestCapabilities( + execution_mode=execution_mode, stream=stream, has_custom_client=has_custom_client, has_agentic_hook=has_agentic_hook, @@ -179,6 +183,11 @@ def _eligibility_context( ) +def _execution_context(context: NativeRequestContext | None, mode: str) -> NativeRequestContext: + current = context or NativeRequestContext() + return with_capabilities(current, replace(current.capabilities, execution_mode=mode)) + + def rust_chat_completions_accepts( *, model: str, @@ -235,6 +244,7 @@ def chat_completions( on_response: ResponseObserver, bedrock: NativeBedrockOptions | None = None, anthropic: NativeAnthropicOptions | None = None, + context: NativeRequestContext | None = None, ) -> ModelResponse | None: def adapt(rust_response: Mapping[str, object]) -> ModelResponse: on_response(rust_response) @@ -256,7 +266,7 @@ def chat_completions( bedrock=bedrock, anthropic=anthropic, ), - context=NativeRequestContext(), + context=_execution_context(context, "sync"), ), call=call_native, fallback=lambda: None, @@ -279,6 +289,7 @@ async def achat_completions( on_response: ResponseObserver, bedrock: NativeBedrockOptions | None = None, anthropic: NativeAnthropicOptions | None = None, + context: NativeRequestContext | None = None, ) -> ModelResponse | None: def adapt(rust_response: Mapping[str, object]) -> ModelResponse: on_response(rust_response) @@ -300,7 +311,7 @@ async def achat_completions( bedrock=bedrock, anthropic=anthropic, ), - context=NativeRequestContext(), + context=_execution_context(context, "async"), ), call=call_native, fallback=async_none, @@ -324,6 +335,7 @@ async def achat_completions_or_fallback( python_fallback: Callable[[], Awaitable[object]], bedrock: NativeBedrockOptions | None = None, anthropic: NativeAnthropicOptions | None = None, + context: NativeRequestContext | None = None, ) -> object: """Await the Rust path, falling back to the caller's own Python path when the bridge is unavailable or the call fails. @@ -354,7 +366,7 @@ async def achat_completions_or_fallback( bedrock=bedrock, anthropic=anthropic, ), - context=NativeRequestContext(), + context=_execution_context(context, "async"), ), call=call_native, fallback=python_fallback, @@ -383,6 +395,7 @@ class _ChatOperation: custom_llm_provider=ctx.custom_llm_provider, options=_provider_eligibility_options(ctx.custom_llm_provider, ctx.litellm_params, ctx.optional_params), context=_eligibility_context( + execution_mode="async" if ctx.acompletion else "sync", stream=bool(ctx.stream), has_custom_client=ctx.client is not None or ctx.shared_session is not None, has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(ctx.logging), @@ -450,10 +463,16 @@ class _ChatOperation: extra_headers=headers, timeout_seconds=timeout_to_seconds(float(ctx.timeout) if isinstance(ctx.timeout, str) else ctx.timeout), ), - context=_eligibility_context( - stream=bool(ctx.stream), - has_custom_client=ctx.client is not None or ctx.shared_session is not None, - has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(ctx.logging), + context=request_context( + logging_obj=ctx.logging, + request_model=ctx.logging.model, + litellm_params=ctx.litellm_params, + capabilities=NativeRequestCapabilities( + execution_mode="async" if ctx.acompletion else "sync", + stream=bool(ctx.stream), + has_custom_client=ctx.client is not None or ctx.shared_session is not None, + has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(ctx.logging), + ), ), ) diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index fc33f9bf6a5..b59e8e0906e 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -25,10 +25,13 @@ from litellm.rust_bridge.protocols import RustAmessages, RustMessages, RustRoute from litellm.rust_bridge.request import ( NativeMessagesRequest, NativePreCallDetails, + NativeRequestCapabilities, NativeRequestContext, NativeRequestOptions, PreparedNativeCall, call_native, + request_context, + with_capabilities, ) from litellm.rust_bridge.runtime import ( BridgeErrorContext, @@ -101,6 +104,10 @@ def messages( custom_llm_provider: str | None, extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, + stream: bool = False, + has_custom_client: bool = False, + has_agentic_hook: bool = False, + context: NativeRequestContext | None = None, ) -> dict[str, object] | None: return _MESSAGES.invoke( prepare=lambda: PreparedNativeCall( @@ -115,7 +122,15 @@ def messages( extra_headers=extra_headers, timeout_seconds=timeout_to_seconds(timeout), ), - context=NativeRequestContext(), + context=with_capabilities( + context or NativeRequestContext(), + NativeRequestCapabilities( + execution_mode="sync", + stream=stream, + has_custom_client=has_custom_client, + has_agentic_hook=has_agentic_hook, + ), + ), ), call=call_native, preflight=lambda: assess_route(_PREFLIGHT, model, custom_llm_provider or ""), @@ -134,6 +149,10 @@ async def amessages( custom_llm_provider: str | None, extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, + stream: bool = False, + has_custom_client: bool = False, + has_agentic_hook: bool = False, + context: NativeRequestContext | None = None, ) -> dict[str, object] | None: return await _MESSAGES.ainvoke( prepare=lambda: PreparedNativeCall( @@ -148,7 +167,15 @@ async def amessages( extra_headers=extra_headers, timeout_seconds=timeout_to_seconds(timeout), ), - context=NativeRequestContext(), + context=with_capabilities( + context or NativeRequestContext(), + NativeRequestCapabilities( + execution_mode="async", + stream=stream, + has_custom_client=has_custom_client, + has_agentic_hook=has_agentic_hook, + ), + ), ), call=call_native, preflight=lambda: assess_route(_PREFLIGHT, model, custom_llm_provider or ""), @@ -190,6 +217,9 @@ class _MessagesOperation: api_key: str | None api_base: str | None python: Callable[[], MessagesResult] + stream: bool = False + asynchronous: bool = False + has_custom_client: bool = False logged: bool = False def prepare(self) -> PreparedNativeCall[NativeMessagesRequest]: @@ -275,7 +305,17 @@ class _MessagesOperation: BaseLLMHTTPHandler.resolve_anthropic_messages_timeout(self.params, False, self.provider) ), ), - context=NativeRequestContext(), + context=request_context( + logging_obj=self.logging, + request_model=self.logging.model if self.logging is not None else self.model, + litellm_params=self.params.model_dump(), + capabilities=NativeRequestCapabilities( + execution_mode="async" if self.asynchronous else "sync", + stream=self.stream, + has_custom_client=self.has_custom_client, + has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(self.logging), + ), + ), ) def fallback(self) -> MessagesResult: @@ -313,7 +353,20 @@ def dispatch_messages( has_custom_client: bool, fallback: Callable[[], MessagesResult], ) -> MessagesResult: - operation: Final = _MessagesOperation(model, provider, messages, body, params, logging, api_key, api_base, fallback) + operation: Final = _MessagesOperation( + model, + provider, + messages, + body, + params, + logging, + api_key, + api_base, + fallback, + stream, + asynchronous, + has_custom_client, + ) def preflight() -> PythonFallback | None: return assess_route( diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index c7617c5de1c..9512f5ed00a 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -6,7 +6,6 @@ from collections.abc import Awaitable, Callable, Mapping from typing import Final, TypeVar from . import configuration as _configuration -from .bindings import UNCHANGED, Unchanged from .protocols import RustAocr, RustOcr, RustRouteDecline from .request import NativeOCRRequest, PreparedNativeCall, call_native from .runtime import ( @@ -16,8 +15,6 @@ from .runtime import ( assess_route, ) -rust_ocr_enabled = _configuration.rust_ocr_enabled -rust = _configuration.rust ResultT = TypeVar("ResultT") @@ -25,40 +22,17 @@ _OCR: Final[EndpointDispatch[RustOcr, RustAocr]] = EndpointDispatch.native( route="ocr", sync=lambda native: native.ocr, asynchronous=lambda native: native.aocr, - enabled=_configuration.rust_ocr_enabled, + enabled=_configuration.rust_enabled, ) _PREFLIGHT: Final[EndpointBinding[RustRouteDecline]] = EndpointBinding.native( route="ocr", select=lambda native: native.ocr_decline, - enabled=_configuration.rust_ocr_enabled, + enabled=_configuration.rust_enabled, ) -def set_rust_ocr( - *, - ocr: RustOcr | None | Unchanged = UNCHANGED, - aocr: RustAocr | None | Unchanged = UNCHANGED, - decline: RustRouteDecline | None | Unchanged = UNCHANGED, -) -> None: - if not isinstance(decline, Unchanged): - if decline is None: - _PREFLIGHT.reset() - else: - _PREFLIGHT.override(decline) - if not isinstance(ocr, Unchanged): - if ocr is None: - _OCR.sync.reset() - else: - _OCR.sync.override(ocr) - if not isinstance(aocr, Unchanged): - if aocr is None: - _OCR.asynchronous.reset() - else: - _OCR.asynchronous.override(aocr) - - def load_rust_ocr() -> RustOcr | None: return _OCR.sync.load() diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 570a1c6634e..8b7d4fe6a23 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -17,11 +17,13 @@ from litellm.rust_bridge.protocols import ( RustRouteDecline, ) from litellm.rust_bridge.request import ( + NativeRequestCapabilities, NativeRequestContext, NativeRequestOptions, NativeResponsesWebSocketRequest, PreparedNativeCall, call_native, + with_capabilities, ) from litellm.rust_bridge.runtime import ( BridgeErrorContext, @@ -93,6 +95,7 @@ async def connect( model: str = "responses websocket", provider: str = "openai", fallback: Callable[[], Awaitable[Connection | None]] = async_none, + context: NativeRequestContext | None = None, ) -> Connection | None: return await _RESPONSES_WEBSOCKET.ainvoke( prepare=lambda: PreparedNativeCall( @@ -104,7 +107,14 @@ async def connect( timeout_seconds=timeout_to_seconds(timeout), custom_llm_provider=provider, ), - context=NativeRequestContext(), + context=with_capabilities( + context or NativeRequestContext(), + NativeRequestCapabilities( + execution_mode="async", + websocket_mode="native", + requires_connection=True, + ), + ), ), call=lambda connection_type, request: call_native(connection_type.connect, request), preflight=lambda: assess_route(_PREFLIGHT, model, provider), @@ -123,6 +133,7 @@ async def open_connection( model: str, provider: str, fallback: Callable[[], AbstractAsyncContextManager[Connection]], + context: NativeRequestContext | None = None, ) -> AsyncGenerator[Connection]: async with AsyncExitStack() as stack: @@ -136,6 +147,7 @@ async def open_connection( model=model, provider=provider, fallback=python_connection, + context=context, ) if backend is None: raise RuntimeError("WebSocket connection returned no connection") diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index 4dd93f6e998..acc13ae0df4 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -5,6 +5,8 @@ from dataclasses import dataclass, field from enum import Enum from typing import Final, Generic, NoReturn, Protocol, TypeAlias, TypeVar +from typing_extensions import assert_never + from litellm.exceptions import APIError, AuthenticationError, InternalServerError, RateLimitError from litellm.rust_bridge.bindings import ( UNCHANGED, @@ -180,6 +182,8 @@ class EndpointBinding(Generic[BindingT]): return value case PythonFallback(): return fallback() + case _ as unreachable: + assert_never(unreachable) async def ainvoke( self, @@ -205,6 +209,8 @@ class EndpointBinding(Generic[BindingT]): return value case PythonFallback(): return await fallback() + case _ as unreachable: + assert_never(unreachable) def assess( self, @@ -257,6 +263,8 @@ class EndpointBinding(Generic[BindingT]): return value case PythonFallback(): self._raise_required(result) + case _ as unreachable: + assert_never(unreachable) async def arequire( self, @@ -281,6 +289,8 @@ class EndpointBinding(Generic[BindingT]): return value case PythonFallback(): self._raise_required(result) + case _ as unreachable: + assert_never(unreachable) def can_attempt( self, diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index ffdcdeec8b3..3aa5fddea1a 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -20,12 +20,15 @@ from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.protocols import RustAtranscription, RustRouteDecline, RustTranscription from litellm.rust_bridge.request import ( NativePreCallDetails, + NativeRequestCapabilities, NativeRequestContext, NativeRequestOptions, NativeTranscriptionRequest, PreparedNativeCall, bedrock_options, call_native, + request_context, + with_capabilities, ) from litellm.rust_bridge.runtime import ( BridgeErrorContext, @@ -90,13 +93,17 @@ def load_rust_atranscription() -> RustAtranscription | None: def transcription( *, model: str, - audio: dict[str, object], + audio: object, api_key: str | None, api_base: str | None, custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, + stream: bool = False, + has_custom_client: bool = False, + input_source_kind: str | None = None, + context: NativeRequestContext | None = None, ) -> dict[str, object] | None: return _TRANSCRIPTION.invoke( prepare=lambda: PreparedNativeCall( @@ -113,7 +120,15 @@ def transcription( timeout_seconds=timeout_to_seconds(timeout), bedrock=bedrock_options(optional_params), ), - context=NativeRequestContext(), + context=with_capabilities( + context or NativeRequestContext(), + NativeRequestCapabilities( + execution_mode="sync", + stream=stream, + has_custom_client=has_custom_client, + input_source_kind=input_source_kind, + ), + ), ), call=call_native, preflight=lambda: assess_route(_PREFLIGHT, model, custom_llm_provider or ""), @@ -126,13 +141,17 @@ def transcription( async def atranscription( *, model: str, - audio: dict[str, object], + audio: object, api_key: str | None, api_base: str | None, custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, + stream: bool = False, + has_custom_client: bool = False, + input_source_kind: str | None = None, + context: NativeRequestContext | None = None, ) -> dict[str, object] | None: return await _TRANSCRIPTION.ainvoke( prepare=lambda: PreparedNativeCall( @@ -149,7 +168,15 @@ async def atranscription( timeout_seconds=timeout_to_seconds(timeout), bedrock=bedrock_options(optional_params), ), - context=NativeRequestContext(), + context=with_capabilities( + context or NativeRequestContext(), + NativeRequestCapabilities( + execution_mode="async", + stream=stream, + has_custom_client=has_custom_client, + input_source_kind=input_source_kind, + ), + ), ), call=call_native, preflight=lambda: assess_route(_PREFLIGHT, model, custom_llm_provider or ""), @@ -162,6 +189,17 @@ async def atranscription( TranscriptionResult = TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] +def _input_source_kind(file: FileTypes) -> str: + content: Final = file[1] if isinstance(file, tuple) else file + if isinstance(content, (bytes, bytearray, memoryview)): + return "bytes" + if isinstance(content, IOBase): + return "file" + if isinstance(content, str): + return "path" + return "opaque" + + @dataclass class _TranscriptionOperation: model: str @@ -174,6 +212,8 @@ class _TranscriptionOperation: timeout: float | httpx.Timeout | None logging: Logging python: Callable[[FileTypes], TranscriptionResult] + asynchronous: bool = False + has_custom_client: bool = False fallback_file: FileTypes | None = None logged: bool = False @@ -228,7 +268,17 @@ class _TranscriptionOperation: timeout_seconds=timeout_to_seconds(self.timeout), bedrock=bedrock_options(self.optional_params), ), - context=NativeRequestContext(), + context=request_context( + logging_obj=self.logging, + request_model=self.logging.model, + litellm_params=self.logging.litellm_params, + capabilities=NativeRequestCapabilities( + execution_mode="async" if self.asynchronous else "sync", + stream=self.optional_params.get("stream") is True, + has_custom_client=self.has_custom_client, + input_source_kind=_input_source_kind(self.file), + ), + ), ) def fallback(self) -> TranscriptionResult: @@ -265,7 +315,18 @@ def dispatch_transcription( fallback: Callable[[FileTypes], TranscriptionResult], ) -> TranscriptionResult: operation: Final = _TranscriptionOperation( - model, provider, file, api_key, api_base, headers, optional_params, timeout, logging, fallback + model, + provider, + file, + api_key, + api_base, + headers, + optional_params, + timeout, + logging, + fallback, + asynchronous, + has_custom_client, ) def preflight() -> PythonFallback | None: 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 45ef8aa59aa..f2abfc98a23 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -128,16 +128,6 @@ def test_load_rust_messages_returns_injected_impl(): assert rust_messages.load_rust_messages() is bridge -def test_bare_rust_still_toggles_ocr(): - from litellm.rust_bridge.ocr import rust_ocr_enabled - - litellm.rust(True) - assert rust_ocr_enabled() is True - - litellm.rust(False) - assert rust_ocr_enabled() is False - - def test_load_rust_amessages_returns_injected_impl(): bridge = RecordingAsyncMessages() litellm.rust(True) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 0babbace05c..1f13c59bd8e 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -10,9 +10,10 @@ import pytest import litellm from litellm._logging import verbose_logger from litellm.integrations.code_interpreter_interception.handler import ( - CodeInterpreterInterceptionLogger, LITELLM_CODE_EXECUTION_TOOL_NAME, + CodeInterpreterInterceptionLogger, ) +from litellm.llms.azure.videos.transformation import AzureVideoConfig from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -24,9 +25,7 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _collect_ws_project_quota_callbacks, _google_genai_streaming_hidden_params, _has_pre_call_deployment_hook, - _rust_responses_websocket_enabled, ) -from litellm.llms.azure.videos.transformation import AzureVideoConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams @@ -1467,6 +1466,7 @@ def _make_responses_handler_call(signed_body): signing provider (e.g. Bedrock Mantle). """ from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams @@ -1522,6 +1522,7 @@ def test_responses_handler_signs_after_fake_stream_prep_strips_stream(): We snapshot request_data at sign time and assert "stream" is already gone. """ from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.openai import ResponsesAPIResponse @@ -1585,6 +1586,7 @@ def _make_compact_handler_call(signed_body, is_async): signing provider (e.g. Bedrock Mantle SigV4 / bearer). """ from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams @@ -2688,21 +2690,6 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h assert "".join([chunk.choices[0].delta.content or "" for chunk in collected]) == "hi" -@pytest.mark.parametrize( - "custom_llm_provider, enabled, expected", - [("openai", True, True), ("openai", False, False), ("azure", True, False), - ("hosted_vllm", True, False), (None, True, False)], -) -def test_the_rust_responses_websocket_needs_openai_and_process_enablement( - custom_llm_provider, enabled, expected, monkeypatch -): - from litellm.rust_bridge import configuration - - configuration.reset_rust_configuration() - monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0") - assert _rust_responses_websocket_enabled(custom_llm_provider) is expected - - def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch): from litellm.integrations.custom_logger import CustomLogger diff --git a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py index c25e52e4421..433845c652a 100644 --- a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py +++ b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py @@ -12,10 +12,8 @@ from litellm.llms.azure_ai.ocr.common_utils import ( is_azure_document_intelligence_model, ) from litellm.ocr.main import _prepare_ocr_request -from litellm.rust_bridge.ocr import _rust_bridge_api_base _DOC = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -_DOC_INTELLIGENCE_ENDPOINT = "https://di.cognitiveservices.azure.com" _AZURE_AI_API_BASE = "https://generic-azure-ai.example.com" @@ -24,13 +22,6 @@ class _FakeLogging: return None -def _resolve_secret(name: str) -> str | None: - return { - "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT": _DOC_INTELLIGENCE_ENDPOINT, - "AZURE_AI_API_BASE": _AZURE_AI_API_BASE, - }.get(name) - - def _prepare(model: str, api_base: str | None): return _prepare_ocr_request( model=model, @@ -57,15 +48,13 @@ class TestIsAzureDocumentIntelligenceModel: class TestDocIntelligenceApiBaseResolution: def test_generic_azure_ai_base_does_not_hijack_doc_intelligence(self, monkeypatch): - """Without an explicit api_base, the AZURE_AI_API_BASE fallback must not - overwrite the endpoint, so it resolves to the Document Intelligence one.""" + """The generic Azure AI fallback must stay out of native admission.""" monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", raising=False) prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", None) assert prepared.api_base is None - assert _rust_bridge_api_base(prepared, _resolve_secret) == _DOC_INTELLIGENCE_ENDPOINT def test_explicit_api_base_is_honoured_for_doc_intelligence(self, monkeypatch): """A caller-supplied api_base must always win, even for doc-intelligence.""" @@ -75,7 +64,6 @@ class TestDocIntelligenceApiBaseResolution: prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", custom) assert prepared.api_base == custom - assert _rust_bridge_api_base(prepared, _resolve_secret) == custom def test_generic_azure_ai_base_still_applies_to_mistral_ocr(self, monkeypatch): """Non doc-intelligence azure_ai models keep using AZURE_AI_API_BASE.""" diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py index 1ab740c4267..6fbdbf061c5 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -1,54 +1,14 @@ """ -Tests for the OCR `req_format` option in the SDK request path: -providers that don't support a native response must reject it, and the Rust -bridge (which only returns the normalized shape) must not serve native requests. +Tests for the OCR `req_format` option in the SDK request path. """ -import dataclasses -from unittest.mock import MagicMock - 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.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( - model="doc-intelligence/prebuilt-layout", - document=dict(DOCUMENT), - api_key="fake-key", - api_base="https://example.cognitiveservices.azure.com", - custom_llm_provider="azure_ai", - extra_headers=None, - provider_config=MagicMock(), - optional_params=optional_params, - litellm_params={}, - effective_timeout=60.0, - litellm_logging_obj=MagicMock(), - ) - - -@pytest.mark.parametrize("optional_params", [{}, {"req_format": "litellm"}]) -def test_rust_ocr_serves_default_format(optional_params): - assert _rust_ocr_supported(_prepared(optional_params)) is True - - -def test_rust_ocr_skipped_for_native_format(): - assert _rust_ocr_supported(_prepared({"req_format": "native"})) is False - - -@pytest.mark.parametrize("provider_config", [CohereParseConfig(), AzureAICohereParseConfig()]) -def test_rust_ocr_skipped_for_configs_without_bridge_support(provider_config): - prepared = dataclasses.replace(_prepared({}), provider_config=provider_config) - - assert _rust_ocr_supported(prepared) is False - - @pytest.mark.asyncio async def test_native_format_rejected_for_provider_without_support_as_bad_request(): with pytest.raises(litellm.BadRequestError, match="not supported for provider") as exc_info: diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 236b15a55e6..cbafa9a5a01 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -213,19 +213,26 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.set_rust_ocr(ocr=None, aocr=None, decline=None) + rust_bridge._OCR.sync.reset() + rust_bridge._OCR.asynchronous.reset() + rust_bridge._PREFLIGHT.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL - rust_bridge.set_rust_ocr( - decline=lambda model, custom_llm_provider, *, context: ( + rust_bridge._PREFLIGHT.override( + lambda model, custom_llm_provider, *, context: ( "unsupported feature" if any(getattr(context.capabilities, key) for key in ("stream", "has_agentic_hook", "has_custom_client")) - or context.capabilities.request_format == "native" + or ( + context.capabilities.request_format == "native" + and not (custom_llm_provider == "azure_ai" and "doc-intelligence" in model) + ) else None ) ) yield - rust_bridge.set_rust_ocr(ocr=None, aocr=None, decline=None) + rust_bridge._OCR.sync.reset() + rust_bridge._OCR.asynchronous.reset() + rust_bridge._PREFLIGHT.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -235,7 +242,7 @@ def fake_bridge(): """Enable the Rust path with an injected recording bridge (no native wheel).""" bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) return bridge @@ -244,27 +251,27 @@ def fake_async_bridge(): """Enable the async Rust path with an injected recording bridge.""" bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._OCR.asynchronous.override(bridge) return bridge def test_rust_toggles_flag(): - assert rust_bridge.rust_ocr_enabled() is False + assert configuration.rust_enabled() is False litellm.rust(True) - assert rust_bridge.rust_ocr_enabled() is True + assert configuration.rust_enabled() is True litellm.rust(False) - assert rust_bridge.rust_ocr_enabled() is False + assert configuration.rust_enabled() is False def test_env_var_enables_rust_ocr(monkeypatch): monkeypatch.setenv("LITELLM_RUST", "1") - assert rust_bridge.rust_ocr_enabled() is True + assert configuration.rust_enabled() is True def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) assert rust_bridge.load_rust_ocr() is bridge @@ -328,7 +335,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch): def test_load_rust_aocr_returns_injected_impl(): bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._OCR.asynchronous.override(bridge) assert rust_bridge.load_rust_aocr() is bridge @@ -337,7 +344,8 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) + rust_bridge._OCR.sync.override(bridge) + rust_bridge._OCR.asynchronous.override(async_bridge) litellm.rust(False) assert rust_bridge.load_rust_ocr() is bridge @@ -356,9 +364,12 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch): bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) + rust_bridge._OCR.sync.override(bridge) + rust_bridge._OCR.asynchronous.override(async_bridge) - rust_bridge.set_rust_ocr(ocr=None, aocr=None, decline=None) + rust_bridge._OCR.sync.reset() + rust_bridge._OCR.asynchronous.reset() + rust_bridge._PREFLIGHT.reset() assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None @@ -405,7 +416,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) response = rust_bridge.dispatch_ocr( prepare=lambda: PreparedNativeCall( request=NativeOCRRequest( @@ -452,7 +463,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._OCR.asynchronous.override(bridge) async def unexpected_fallback(): pytest.fail("unexpected Python fallback") @@ -495,7 +506,7 @@ 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) + rust_bridge._OCR.sync.override(bridge) response = ocr_main._run_rust_ocr( fallback=lambda: pytest.fail("unexpected Python fallback"), @@ -530,7 +541,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( fallback=lambda: pytest.fail("unexpected Python fallback"), @@ -544,7 +555,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) def _resolver(name: str) -> str | None: raise AssertionError(f"resolver should not be called for {name}") @@ -565,7 +576,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): bridge = RecordingBridge() resolver_calls = [] litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) def _resolver(name): resolver_calls.append(name) @@ -589,7 +600,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( fallback=lambda: pytest.fail("unexpected Python fallback"), @@ -614,7 +625,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) def _resolver(name: str) -> str | None: return { @@ -638,7 +649,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( fallback=lambda: pytest.fail("unexpected Python fallback"), @@ -657,7 +668,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( fallback=lambda: pytest.fail("unexpected Python fallback"), @@ -679,7 +690,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): logging_obj = RecordingLogging() bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.sync.override(bridge) ocr_main._run_rust_ocr( fallback=lambda: pytest.fail("unexpected Python fallback"), @@ -744,7 +755,7 @@ def test_ocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=RaisingBridge()) + rust_bridge._OCR.sync.override(RaisingBridge()) with pytest.raises(CapturedException): litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -790,7 +801,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge()) + rust_bridge._OCR.asynchronous.override(RaisingAsyncBridge()) with pytest.raises(CapturedException): await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")