From fba6af26ad0f3e89330c8e5c5b0401c344de289e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 21:35:34 -0700 Subject: [PATCH] fix(native): preserve typed capability context --- .../crates/core/src/request_context.rs | 5 + .../crates/python-bridge/src/marshal.rs | 15 + litellm/llms/anthropic/chat/handler.py | 200 ++++--- litellm/llms/bedrock/chat/converse_handler.py | 318 ++++++------ litellm/llms/custom_httpx/llm_http_handler.py | 7 + litellm/ocr/main.py | 282 +++------- litellm/rust_bridge/chat_completions.py | 173 +++---- litellm/rust_bridge/messages.py | 83 +-- litellm/rust_bridge/ocr.py | 308 +++++++++-- litellm/rust_bridge/request.py | 5 + litellm/rust_bridge/responses_websocket.py | 73 ++- litellm/rust_bridge/transcription.py | 86 ++-- .../test_rust_bridge_messages.py | 191 ++++++- .../chat/test_anthropic_chat_handler.py | 486 +++++++++++++----- .../chat/test_bedrock_converse_handler.py | 124 +++-- tests/test_litellm/ocr/test_rust_bridge.py | 175 ++----- .../responses/test_rust_bridge_websocket.py | 96 +++- .../rust_bridge/test_chat_completions.py | 244 ++------- .../test_audio_transcription_rust_bridge.py | 47 +- 19 files changed, 1595 insertions(+), 1323 deletions(-) diff --git a/litellm-rust/crates/core/src/request_context.rs b/litellm-rust/crates/core/src/request_context.rs index 2796edf9993..8f02af57f59 100644 --- a/litellm-rust/crates/core/src/request_context.rs +++ b/litellm-rust/crates/core/src/request_context.rs @@ -7,10 +7,15 @@ pub struct RequestAttribution { #[derive(Clone, Debug, Default, PartialEq)] pub struct RequestCapabilities { + pub execution_mode: Option, pub stream: bool, pub has_agentic_hook: bool, pub has_custom_client: bool, pub request_format: Option, + pub input_source_kind: Option, + pub native_response_format: bool, + pub websocket_mode: Option, + pub requires_connection: bool, } #[derive(Clone, Debug, Default, PartialEq)] diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 1b6b37f04de..122a4b6f9e1 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -118,10 +118,15 @@ pub(crate) struct NativeRequestContext { #[derive(FromPyObject)] struct NativeRequestCapabilities { + execution_mode: Option, stream: bool, has_agentic_hook: bool, has_custom_client: bool, request_format: Option, + input_source_kind: Option, + native_response_format: bool, + websocket_mode: Option, + requires_connection: bool, } impl From for litellm_core::request_context::LiteLlmRequestContext { @@ -136,10 +141,15 @@ impl From for litellm_core::request_context::LiteLlmReques user_api_key_team_id: input.attribution.user_api_key_team_id, }, capabilities: litellm_core::request_context::RequestCapabilities { + execution_mode: input.capabilities.execution_mode, stream: input.capabilities.stream, has_agentic_hook: input.capabilities.has_agentic_hook, has_custom_client: input.capabilities.has_custom_client, request_format: input.capabilities.request_format, + input_source_kind: input.capabilities.input_source_kind, + native_response_format: input.capabilities.native_response_format, + websocket_mode: input.capabilities.websocket_mode, + requires_connection: input.capabilities.requires_connection, }, } } @@ -206,10 +216,15 @@ class BedrockOptions: @dataclass(frozen=True) class Capabilities: + execution_mode: object = None stream: object = False has_agentic_hook: object = False has_custom_client: object = False request_format: object = None + input_source_kind: object = None + native_response_format: object = False + websocket_mode: object = None + requires_connection: object = False @dataclass(frozen=True) class VertexOptions: diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index cc9430f6bf4..a0c84615224 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -27,7 +27,9 @@ 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 anative_first, native_first, provider_errors from litellm.rust_bridge.request import anthropic_options +from litellm.rust_bridge.runtime import DispatchResult from litellm.types.llms.anthropic import ( ContentBlockDelta, ContentBlockStart, @@ -369,15 +371,7 @@ class AnthropicChatCompletion(BaseLLM): if config is None: raise ValueError(f"Provider config not found for model: {model} and provider: {custom_llm_provider}") - def build_request() -> tuple[dict, dict]: # mutable-ok: rewritten in place downstream - """Translate the request the Python way, returning `(headers, data)`. - - The pair stays mutable because the streaming path rewrites it in - place (`data["stream"] = True`) before sending. - - Shared by the normal path and by the Rust path's fallback, which - builds it only when the Rust call did not serve the request. - """ + def prepare_python() -> tuple[dict[str, str], dict[str, object]]: # mutable-ok: stream mutates data request_data: Final = config.transform_request( model=model, messages=messages, @@ -385,12 +379,29 @@ class AnthropicChatCompletion(BaseLLM): litellm_params=litellm_params, headers=headers, ) - return update_request_with_filtered_beta( + python_headers, data = update_request_with_filtered_beta( headers=headers, request_data=request_data, provider=custom_llm_provider, ) + ## LOGGING + # Reaching here with `serves_via_rust` set means the Rust attempt + # declined at call time, before the provider was called, and already + # logged this request. That is the same attempt continuing. + if not serves_via_rust: + logging_obj.pre_call( + input=messages, + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": python_headers, + }, + ) + print_verbose(f"_is_function_call: {_is_function_call}") + return python_headers, data + # The Rust core owns the whole call for the subset it accepts, so ask # before transforming: whichever path runs emits pre_call exactly once. # `get_config` merges the class-level defaults (Anthropic's required @@ -407,68 +418,26 @@ class AnthropicChatCompletion(BaseLLM): litellm_params=litellm_params, stream=stream, ) + rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict + "complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent + "model": model, + "messages": messages, + **rust_optional_params, + }, + "api_base": api_base, + "headers": headers, + } if serves_via_rust: - rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict - "complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent - "model": model, - "messages": messages, - **rust_optional_params, - }, - "api_base": api_base, - "headers": headers, - } logging_obj.pre_call(input=messages, api_key=api_key, additional_args=rust_logging_args) - log_rust_post_call: Final = rust_chat_completions_bridge.response_logger( - logging_obj=logging_obj, - messages=messages, - api_key=api_key, - additional_args=rust_logging_args, - ) - if acompletion is True: + log_rust_post_call: Final = rust_chat_completions_bridge.response_logger( + logging_obj=logging_obj, + messages=messages, + api_key=api_key, + additional_args=rust_logging_args, + ) - async def python_fallback() -> "ModelResponse | CustomStreamWrapper": - # pre_call already fired for this request above. The Rust - # path only declines before the provider is called, so this - # is the same attempt continuing, not a second one. - fallback_headers, fallback_data = build_request() - return await self.acompletion_function( - model=model, - messages=messages, - data=fallback_data, - api_base=api_base, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - encoding=encoding, - api_key=api_key, - provider_config=config, - logging_obj=logging_obj, - optional_params=optional_params, - stream=stream, - _is_function_call=_is_function_call, - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=fallback_headers, - client=client, - json_mode=json_mode, - timeout=timeout, - ) - - 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=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=headers, - timeout=timeout, - on_response=log_rust_post_call, - python_fallback=python_fallback, - anthropic=anthropic_options(litellm_params), - ) - rust_response: Final = rust_chat_completions_bridge.chat_completions( + def native_completion() -> DispatchResult[ModelResponse]: + return rust_chat_completions_bridge.chat_completions( model=model, messages=messages, optional_params=rust_optional_params, @@ -480,34 +449,42 @@ class AnthropicChatCompletion(BaseLLM): timeout=timeout, on_response=log_rust_post_call, anthropic=anthropic_options(litellm_params), + stream=bool(stream), + has_custom_client=client is not None, + eligible=serves_via_rust, ) - if rust_response is not None: - return rust_response - headers, data = build_request() - - ## LOGGING - # Reaching here with `serves_via_rust` set means the Rust attempt - # declined at call time, before the provider was called, and already - # logged this request. That is the same attempt continuing. - if not serves_via_rust: - logging_obj.pre_call( - input=messages, + async def native_acompletion() -> DispatchResult[ModelResponse]: + return await rust_chat_completions_bridge.achat_completions( + model=model, + messages=messages, + optional_params=rust_optional_params, + model_response=model_response, api_key=api_key, - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": headers, - }, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + on_response=log_rust_post_call, + anthropic=anthropic_options(litellm_params), + stream=bool(stream), + has_custom_client=client is not None, + eligible=serves_via_rust, ) - print_verbose(f"_is_function_call: {_is_function_call}") - if acompletion is True: + + @anative_first( + native=native_acompletion, + route="chat_completions", + errors=lambda: provider_errors(custom_llm_provider or "", model), + ) + async def execute_async() -> ModelResponse | CustomStreamWrapper: + headers, data = prepare_python() if ( stream is True ): # if function call - fake the streaming (need complete blocks for output parsing in openai format) print_verbose("makes async anthropic streaming POST request") data["stream"] = stream - return self.acompletion_stream_function( + return await self.acompletion_stream_function( model=model, messages=messages, data=data, @@ -529,7 +506,7 @@ class AnthropicChatCompletion(BaseLLM): client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), ) else: - return self.acompletion_function( + return await self.acompletion_function( model=model, messages=messages, data=data, @@ -551,7 +528,14 @@ class AnthropicChatCompletion(BaseLLM): json_mode=json_mode, timeout=timeout, ) - else: + + @native_first( + native=native_completion, + route="chat_completions", + errors=lambda: provider_errors(custom_llm_provider or "", model), + ) + def execute_sync() -> ModelResponse | CustomStreamWrapper: + headers, data = prepare_python() ## COMPLETION CALL if ( stream is True @@ -583,13 +567,12 @@ class AnthropicChatCompletion(BaseLLM): ) else: - if client is None or not isinstance(client, HTTPHandler): - client = _get_httpx_client(params={"timeout": timeout}) - else: - client = client + python_client: Final = ( + client if isinstance(client, HTTPHandler) else _get_httpx_client(params={"timeout": timeout}) + ) try: - response: Final = client.post( + response: Final = python_client.post( api_base, headers=headers, data=json.dumps(data), @@ -610,20 +593,21 @@ class AnthropicChatCompletion(BaseLLM): status_code=status_code, headers=error_headers, ) + return config.transform_response( + model=model, + raw_response=response, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + request_data=data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + json_mode=json_mode, + ) - return config.transform_response( - model=model, - raw_response=response, - model_response=model_response, - logging_obj=logging_obj, - api_key=api_key, - request_data=data, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - encoding=encoding, - json_mode=json_mode, - ) + return execute_async() if acompletion else execute_sync() def embedding(self): # logic for parsing in - calling - parsing out model embedding calls diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 19ffb695af1..982af97a6c4 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -10,10 +10,6 @@ from litellm.anthropic_beta_headers_manager import ( update_headers_with_filtered_beta, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject -from litellm.llms.bedrock.request_metadata import ( - get_bedrock_request_metadata_fields, - resolve_bedrock_request_metadata, -) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -22,7 +18,9 @@ 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.request import NativeBedrockOptions +from litellm.rust_bridge.dispatch import anative_first, native_first, provider_errors +from litellm.rust_bridge.request import bedrock_options +from litellm.rust_bridge.runtime import DispatchResult from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper @@ -398,15 +396,11 @@ class BedrockConverseLLM(BaseAWSLLM): # resolved so both paths sign as the same principal. Bearer-token auth # resolves no SigV4 principal at all, and each path reads that token # itself. - rust_optional_params: Final = optional_params - rust_bedrock_options: Final = NativeBedrockOptions( - aws_access_key_id=None if credentials is None else credentials.access_key, - aws_secret_access_key=None if credentials is None else credentials.secret_key, - aws_session_token=None if credentials is None else credentials.token, - aws_region_name=aws_region_name, - request_metadata_fields=get_bedrock_request_metadata_fields(), - request_metadata=resolve_bedrock_request_metadata(litellm_params, optional_params.get("requestMetadata")), - ) + rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy + **optional_params, + **_sigv4_principal(credentials), + "aws_region_name": aws_region_name, + } serves_via_rust: Final = rust_chat_completions_accepts( model=model, messages=messages, @@ -415,55 +409,25 @@ class BedrockConverseLLM(BaseAWSLLM): litellm_params=litellm_params, stream=stream, ) + rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict + "complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent + "messages": messages, + **optional_params, + }, + "api_base": proxy_endpoint_url, + "headers": headers, + } if serves_via_rust: - rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict - "complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent - "messages": messages, - **optional_params, - }, - "api_base": proxy_endpoint_url, - "headers": headers, - } logging_obj.pre_call(input=messages, api_key="", additional_args=rust_logging_args) - log_rust_post_call: Final = rust_chat_completions_bridge.response_logger( - logging_obj=logging_obj, - messages=messages, - api_key="", - 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, - bedrock=rust_bedrock_options, - python_fallback=lambda: self.async_completion( - model=model, - messages=messages, - api_base=proxy_endpoint_url, - model_response=model_response, - encoding=encoding, - logging_obj=logging_obj, - optional_params=optional_params, - stream=stream, - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=headers, - timeout=timeout, - client=client, - credentials=credentials, - api_key=api_key, - skip_pre_call_logging=True, - ), - ) - rust_response: Final = rust_chat_completions_bridge.chat_completions( + log_rust_post_call: Final = rust_chat_completions_bridge.response_logger( + logging_obj=logging_obj, + messages=messages, + api_key="", + additional_args=rust_logging_args, + ) + + def native_completion() -> DispatchResult[ModelResponse]: + return rust_chat_completions_bridge.chat_completions( model=model, messages=messages, optional_params=rust_optional_params, @@ -474,17 +438,39 @@ class BedrockConverseLLM(BaseAWSLLM): extra_headers=headers, timeout=timeout, on_response=log_rust_post_call, - bedrock=rust_bedrock_options, + bedrock=bedrock_options(rust_optional_params), + stream=bool(stream), + has_custom_client=client is not None, + eligible=serves_via_rust, ) - if rust_response is not None: - return rust_response - ### ROUTING (ASYNC, STREAMING, SYNC) - if acompletion: - if isinstance(client, HTTPHandler): - client = None + async def native_acompletion() -> DispatchResult[ModelResponse]: + return await 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, + bedrock=bedrock_options(rust_optional_params), + stream=bool(stream), + has_custom_client=client is not None, + eligible=serves_via_rust, + ) + + @anative_first( + native=native_acompletion, + route="chat_completions", + errors=lambda: provider_errors("bedrock", model), + ) + async def execute_async() -> ModelResponse | CustomStreamWrapper: + python_client: Final = None if isinstance(client, HTTPHandler) else client if stream is True: - return self.async_streaming( + return await self.async_streaming( model=model, messages=messages, api_base=proxy_endpoint_url, @@ -497,7 +483,7 @@ class BedrockConverseLLM(BaseAWSLLM): logger_fn=logger_fn, headers=headers, timeout=timeout, - client=client, + client=python_client, json_mode=json_mode, fake_stream=fake_stream, credentials=credentials, @@ -505,7 +491,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream_chunk_size=stream_chunk_size, ) ### ASYNC COMPLETION - return self.async_completion( + return await self.async_completion( model=model, messages=messages, api_base=proxy_endpoint_url, @@ -518,108 +504,112 @@ class BedrockConverseLLM(BaseAWSLLM): logger_fn=logger_fn, headers=headers, timeout=timeout, - client=client, + client=python_client, credentials=credentials, api_key=api_key, + skip_pre_call_logging=serves_via_rust, + ) + + @native_first( + native=native_completion, + route="chat_completions", + errors=lambda: provider_errors("bedrock", model), + ) + def execute_sync() -> ModelResponse | CustomStreamWrapper: + ## TRANSFORMATION ## + + _data: Final = litellm.AmazonConverseConfig()._transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=extra_headers, + ) + data: Final = json.dumps(_data) + + prepped: Final = self.get_request_headers( + credentials=credentials, + aws_region_name=aws_region_name, + extra_headers=extra_headers, + endpoint_url=proxy_endpoint_url, + data=data, + headers=headers, + api_key=api_key, ) - ## TRANSFORMATION ## + ## LOGGING + # Reaching here with `serves_via_rust` set means the synchronous Rust + # attempt declined at call time, before the provider was called, and + # already logged this request. That is the same attempt continuing. + if not serves_via_rust: + logging_obj.pre_call( + input=messages, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": proxy_endpoint_url, + "headers": prepped.headers, + }, + ) + resolved_timeout: Final = httpx.Timeout(timeout) if isinstance(timeout, (float, int)) else timeout + python_client: Final = ( + _get_httpx_client({"timeout": resolved_timeout} if resolved_timeout is not None else None) + if client is None or isinstance(client, AsyncHTTPHandler) + else client + ) - _data: Final = litellm.AmazonConverseConfig()._transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=extra_headers, - ) - data: Final = json.dumps(_data) + if stream is not None and stream is True: + completion_stream, response_headers = make_sync_call( + client=python_client, + api_base=proxy_endpoint_url, + headers=prepped.headers, + data=data, + model=model, + messages=messages, + logging_obj=logging_obj, + json_mode=json_mode, + fake_stream=fake_stream, + stream_chunk_size=stream_chunk_size, + ) + streaming_response: Final = CustomStreamWrapper( + completion_stream=completion_stream, + model=model, + custom_llm_provider="bedrock", + logging_obj=logging_obj, + _response_headers=response_headers, + ) - prepped: Final = self.get_request_headers( - credentials=credentials, - aws_region_name=aws_region_name, - extra_headers=extra_headers, - endpoint_url=proxy_endpoint_url, - data=data, - headers=headers, - api_key=api_key, - ) + return streaming_response - ## LOGGING - # Reaching here with `serves_via_rust` set means the synchronous Rust - # attempt declined at call time, before the provider was called, and - # already logged this request. That is the same attempt continuing. - # The asynchronous branch above returns before this point, and hands - # its own fallback `skip_pre_call_logging=True` for the same reason. - if not serves_via_rust: - logging_obj.pre_call( - input=messages, + ### COMPLETION + + try: + response: Final = python_client.post( + url=proxy_endpoint_url, + headers=prepped.headers, + data=data, + logging_obj=logging_obj, + ) + response.raise_for_status() + except httpx.HTTPStatusError as err: + error_code: Final = err.response.status_code + raise BedrockError(status_code=error_code, message=err.response.text) + except httpx.TimeoutException: + raise BedrockError(status_code=408, message="Timeout error occurred.") + + sync_transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=model_response, + stream=stream if isinstance(stream, bool) else False, + logging_obj=logging_obj, api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": proxy_endpoint_url, - "headers": prepped.headers, - }, - ) - if client is None or isinstance(client, AsyncHTTPHandler): - _params: Final = {} - if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, 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( - client=(client if client is not None and isinstance(client, HTTPHandler) else None), - api_base=proxy_endpoint_url, - headers=prepped.headers, data=data, - model=model, messages=messages, - logging_obj=logging_obj, - json_mode=json_mode, - fake_stream=fake_stream, - stream_chunk_size=stream_chunk_size, - ) - streaming_response: Final = CustomStreamWrapper( - completion_stream=completion_stream, - model=model, - custom_llm_provider="bedrock", - logging_obj=logging_obj, - _response_headers=response_headers, + optional_params=optional_params, + encoding=encoding, ) + sync_transformed_response.set_provider_response_headers(response.headers) + return sync_transformed_response - return streaming_response - - ### COMPLETION - - try: - response: Final = client.post( - url=proxy_endpoint_url, - headers=prepped.headers, - data=data, - logging_obj=logging_obj, - ) - response.raise_for_status() - except httpx.HTTPStatusError as err: - error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) - except httpx.TimeoutException: - raise BedrockError(status_code=408, message="Timeout error occurred.") - - sync_transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=model_response, - stream=stream if isinstance(stream, bool) else False, - logging_obj=logging_obj, - api_key="", - data=data, - messages=messages, - optional_params=optional_params, - encoding=encoding, - ) - sync_transformed_response.set_provider_response_headers(response.headers) - return sync_transformed_response + return execute_async() if acompletion else execute_sync() diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ee5a49374af..5aa580ae47a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2232,6 +2232,8 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, has_agentic_hook=self._has_agentic_completion_hook(logging_obj), + stream=bool(stream), + has_custom_client=client is not None, model=model, api_key=api_key, api_base=api_base, @@ -2388,6 +2390,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, has_agentic_hook: bool, + stream: bool, + has_custom_client: bool, model: str, api_key: str | None, api_base: str | None, @@ -2415,6 +2419,9 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, extra_headers=headers, timeout=timeout, + stream=stream, + has_custom_client=has_custom_client, + has_agentic_hook=has_agentic_hook, ) def adapt(rust_response: dict[str, object]) -> AnthropicMessagesResponse: diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index bebb4b44589..7fea86e10b8 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -8,7 +8,6 @@ import mimetypes import os import re from collections.abc import Callable, Coroutine, Mapping -from dataclasses import dataclass from io import IOBase from typing import Any, Final, cast @@ -18,23 +17,16 @@ 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.request import ( - NativeRequestOptions, - PreparedNativeCall, - vertex_options, -) -from litellm.rust_bridge.timeouts import timeout_to_seconds +from litellm.rust_bridge.dispatch import anative_first, native_first, provider_errors +from litellm.rust_bridge.runtime import DispatchResult from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -43,28 +35,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 - - -_RUST_OCR_PROVIDERS: Final = { - "mistral", - "azure_ai", - "vertex_ai", -} - - def _prepare_ocr_request( model: str, document: Mapping[str, object], @@ -74,7 +44,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)) @@ -171,7 +141,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, @@ -186,142 +156,70 @@ 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, +@anative_first( + native=rust_ocr_bridge.aattempt_ocr, + route="ocr", + errors=lambda prepared_request, resolve_api_key: provider_errors( + prepared_request.custom_llm_provider, prepared_request.model + ), +) +async def _execute_aocr( + prepared_request: rust_ocr_bridge.PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> PreparedNativeCall[rust_ocr_bridge.NativeOCRRequest]: - 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 PreparedNativeCall( - request=rust_ocr_bridge.NativeOCRRequest( - model=prepared_request.model, - document=prepared_request.document, - optional_params=prepared_request.optional_params, - ), - options=NativeRequestOptions( - vertex=vertex_options(rust_optional_params), - api_key=resolved_api_key, - api_base=rust_api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=cast( # cast-ok: provider header normalization returns string-object pairs - dict[str, object], resolved_headers - ), - timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), - ), - ) - - -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, - ), - 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, - ), - fallback=fallback, - adapt=OCRResponse.model_validate, + pending: Final = base_llm_http_handler.ocr( model=prepared_request.model, - provider=prepared_request.custom_llm_provider, - eligible=_rust_ocr_supported(prepared_request), + document=prepared_request.document, + optional_params=prepared_request.optional_params, + timeout=prepared_request.effective_timeout, + logging_obj=prepared_request.litellm_logging_obj, + api_key=prepared_request.api_key, + api_base=prepared_request.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + aocr=True, + headers=prepared_request.extra_headers, + provider_config=prepared_request.provider_config, + litellm_params=prepared_request.litellm_params, + ) + response: Final = await pending if asyncio.iscoroutine(pending) else pending + if response is None: + raise ValueError(f"Got an unexpected None response from the OCR API: {response}") + return response + + +def _attempt_ocr( + prepared_request: rust_ocr_bridge.PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], + is_async: bool, +) -> DispatchResult[OCRResponse]: + return rust_ocr_bridge.attempt_ocr(prepared_request=prepared_request, resolve_api_key=resolve_api_key) + + +@native_first( + native=_attempt_ocr, + route="ocr", + errors=lambda prepared_request, resolve_api_key, is_async: provider_errors( + prepared_request.custom_llm_provider, prepared_request.model + ), +) +def _execute_ocr( + prepared_request: rust_ocr_bridge.PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], + is_async: bool, +) -> OCRResponse | Coroutine[object, object, OCRResponse]: + return base_llm_http_handler.ocr( + model=prepared_request.model, + document=prepared_request.document, + optional_params=prepared_request.optional_params, + timeout=prepared_request.effective_timeout, + logging_obj=prepared_request.litellm_logging_obj, + api_key=prepared_request.api_key, + api_base=prepared_request.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + aocr=is_async, + headers=prepared_request.extra_headers, + provider_config=prepared_request.provider_config, + litellm_params=prepared_request.litellm_params, ) @@ -421,31 +319,7 @@ async def aocr( from litellm.secret_managers.main import get_secret_str - async def python_fallback() -> OCRResponse: - pending: Final = base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=True, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, - ) - response: Final = await pending if asyncio.iscoroutine(pending) else pending - if response is None: - 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 _execute_aocr(prepared_request=prepared, resolve_api_key=get_secret_str) except Exception as e: raise litellm.exception_type( model=model, @@ -686,27 +560,7 @@ def ocr( from litellm.secret_managers.main import get_secret_str - def python_fallback() -> OCRResponse | Coroutine[object, object, OCRResponse]: - return base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=_is_async, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, - ) - - return _run_rust_ocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - fallback=python_fallback, - ) + return _execute_ocr(prepared_request=prepared, resolve_api_key=get_secret_str, is_async=_is_async) except Exception as e: raise litellm.exception_type( model=model, diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index a46a1f541bd..6af8471e2b3 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,7 +20,7 @@ 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.protocols import ( RustAchatCompletions, @@ -35,17 +31,13 @@ from litellm.rust_bridge.request import ( NativeAnthropicOptions, NativeBedrockOptions, NativeChatCompletionsRequest, + NativeRequestCapabilities, NativeRequestContext, NativeRequestOptions, PreparedNativeCall, call_native, ) -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 @@ -103,16 +95,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 ) @@ -126,14 +112,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() @@ -202,14 +188,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( @@ -240,18 +236,23 @@ def chat_completions( on_response: ResponseObserver, bedrock: NativeBedrockOptions | None = None, anthropic: NativeAnthropicOptions | None = None, -) -> ModelResponse | None: + stream: bool = False, + has_custom_client: bool = False, + eligible: bool = True, +) -> 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( + def call(native: RustChatCompletions, prepared: PreparedNativeCall[NativeChatCompletionsRequest]) -> Mapping[str, object]: + return call_native(native, prepared) + + return attempt( + load=_CHAT.load, + enabled=rust_enabled(), + eligible=eligible, prepare=lambda: PreparedNativeCall( - NativeChatCompletionsRequest( - model=model, - messages=messages, - optional_params=optional_params, - ), + request=NativeChatCompletionsRequest(model=model, messages=messages, optional_params=optional_params), options=NativeRequestOptions( api_key=api_key, api_base=api_base, @@ -261,12 +262,16 @@ def chat_completions( bedrock=bedrock, anthropic=anthropic, ), - context=NativeRequestContext(), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + execution_mode="sync", + stream=stream, + has_custom_client=has_custom_client, + ) + ), ), - call=call_native, - fallback=lambda: None, + call=call, adapt=adapt, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -284,18 +289,26 @@ async def achat_completions( on_response: ResponseObserver, bedrock: NativeBedrockOptions | None = None, anthropic: NativeAnthropicOptions | None = None, -) -> ModelResponse | None: + stream: bool = False, + has_custom_client: bool = False, + eligible: bool = True, +) -> 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( + async def call( + native: RustAchatCompletions, + prepared: PreparedNativeCall[NativeChatCompletionsRequest], + ) -> Mapping[str, object]: + return await call_native(native, prepared) + + return await aattempt( + load=_ACHAT.load, + enabled=rust_enabled(), + eligible=eligible, prepare=lambda: PreparedNativeCall( - NativeChatCompletionsRequest( - model=model, - messages=messages, - optional_params=optional_params, - ), + request=NativeChatCompletionsRequest(model=model, messages=messages, optional_params=optional_params), options=NativeRequestOptions( api_key=api_key, api_base=api_base, @@ -305,64 +318,14 @@ async def achat_completions( bedrock=bedrock, anthropic=anthropic, ), - context=NativeRequestContext(), - ), - call=call_native, - 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]], - bedrock: NativeBedrockOptions | None = None, - anthropic: NativeAnthropicOptions | 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. - - 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( - prepare=lambda: PreparedNativeCall( - NativeChatCompletionsRequest( - model=model, - messages=messages, - optional_params=optional_params, + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + execution_mode="async", + stream=stream, + has_custom_client=has_custom_client, + ) ), - options=NativeRequestOptions( - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), - bedrock=bedrock, - anthropic=anthropic, - ), - context=NativeRequestContext(), ), - call=call_native, - fallback=python_fallback, + call=call, adapt=adapt, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index f41cff140f8..445009e4cfb 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -6,30 +6,21 @@ 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.request import ( NativeMessagesRequest, + NativeRequestCapabilities, NativeRequestContext, NativeRequestOptions, PreparedNativeCall, call_native, ) -from litellm.rust_bridge.runtime import ( - BridgeErrorContext, - EndpointDispatch, - 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, -) +_MESSAGES: Final[NativeBinding[RustMessages]] = NativeBinding(lambda native: native.messages) +_AMESSAGES: Final[NativeBinding[RustAmessages]] = NativeBinding(lambda native: native.amessages) def set_rust_messages( @@ -39,22 +30,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( @@ -66,13 +57,16 @@ 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( + stream: bool = False, + has_custom_client: bool = False, + has_agentic_hook: bool = False, +) -> DispatchResult[dict[str, object]]: + return attempt( + load=_MESSAGES.load, + enabled=True, + eligible=True, prepare=lambda: PreparedNativeCall( - NativeMessagesRequest( - model=model, - body=body, - ), + request=NativeMessagesRequest(model=model, body=body), options=NativeRequestOptions( api_key=api_key, api_base=api_base, @@ -80,12 +74,17 @@ def messages( extra_headers=extra_headers, timeout_seconds=timeout_to_seconds(timeout), ), - context=NativeRequestContext(), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + execution_mode="sync", + stream=stream, + has_custom_client=has_custom_client, + has_agentic_hook=has_agentic_hook, + ) + ), ), call=call_native, - fallback=lambda: None, adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -98,13 +97,16 @@ 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( + stream: bool = False, + has_custom_client: bool = False, + has_agentic_hook: bool = False, +) -> DispatchResult[dict[str, object]]: + return await aattempt( + load=_AMESSAGES.load, + enabled=True, + eligible=True, prepare=lambda: PreparedNativeCall( - NativeMessagesRequest( - model=model, - body=body, - ), + request=NativeMessagesRequest(model=model, body=body), options=NativeRequestOptions( api_key=api_key, api_base=api_base, @@ -112,10 +114,15 @@ async def amessages( extra_headers=extra_headers, timeout_seconds=timeout_to_seconds(timeout), ), - context=NativeRequestContext(), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + execution_mode="async", + stream=stream, + has_custom_client=has_custom_client, + has_agentic_hook=has_agentic_hook, + ) + ), ), call=call_native, - 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 320f4b6718f..e2236317853 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -1,29 +1,68 @@ -"""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 .request import NativeOCRRequest, PreparedNativeCall, call_native -from .runtime import ( - BridgeErrorContext, - EndpointDispatch, +import httpx +from pydantic import TypeAdapter + +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.request import ( + NativeOCRRequest, + NativeRequestCapabilities, + NativeRequestContext, + NativeRequestOptions, + PreparedNativeCall, + call_native, + vertex_options, ) +from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt +from litellm.rust_bridge.timeouts import timeout_to_seconds -rust_ocr_enabled = _configuration.rust_ocr_enabled -rust = _configuration.rust -ResultT = TypeVar("ResultT") +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, +@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", + } ) @@ -34,57 +73,212 @@ 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[[], PreparedNativeCall[NativeOCRRequest]], - fallback: Callable[[], ResultT], - adapt: Callable[[Mapping[str, object]], ResultT], - model: str, - provider: str, - eligible: bool, -) -> ResultT: - return _OCR.invoke( - prepare=prepare, - call=call_native, - 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 _ocr_input_source_kind(document: dict[str, object]) -> str: + if "document_url" in document: + return "document_url" + if "image_url" in document: + return "image_url" + if "file" in document: + return "file" + return "inline" + + +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[[], PreparedNativeCall[NativeOCRRequest]], - 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_native, - 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: call_native( + native, + PreparedNativeCall( + request=NativeOCRRequest( + model=prepared_request.model, + document=prepared_request.document, + optional_params=prepared.optional_params, + ), + options=NativeRequestOptions( + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), + vertex=vertex_options(prepared.optional_params), + ), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + execution_mode="sync", + input_source_kind=_ocr_input_source_kind(prepared_request.document), + native_response_format=( + prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native" + ), + ) + ), + ), + ), + 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: call_native( + native, + PreparedNativeCall( + request=NativeOCRRequest( + model=prepared_request.model, + document=prepared_request.document, + optional_params=prepared.optional_params, + ), + options=NativeRequestOptions( + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout), + vertex=vertex_options(prepared.optional_params), + ), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + execution_mode="async", + input_source_kind=_ocr_input_source_kind(prepared_request.document), + native_response_format=( + prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native" + ), + ) + ), + ), + ), + adapt=OCRResponse.model_validate, + eligible=_rust_ocr_supported(prepared_request), ) diff --git a/litellm/rust_bridge/request.py b/litellm/rust_bridge/request.py index 4239b9763a1..e81d10f587f 100644 --- a/litellm/rust_bridge/request.py +++ b/litellm/rust_bridge/request.py @@ -90,10 +90,15 @@ class RequestAttribution: @dataclass(frozen=True, slots=True) class NativeRequestCapabilities: + execution_mode: str | None = None stream: bool = False has_agentic_hook: bool = False has_custom_client: bool = False request_format: str | None = None + input_source_kind: str | None = None + native_response_format: bool = False + websocket_mode: str | None = None + requires_connection: bool = False @dataclass(frozen=True, slots=True) diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index b7098aaddb4..2dc3e18bd2c 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -2,36 +2,32 @@ from __future__ import annotations +from collections.abc import AsyncGenerator +from contextlib import AbstractAsyncContextManager, asynccontextmanager 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.request import ( + NativeRequestCapabilities, NativeRequestContext, NativeRequestOptions, NativeResponsesWebSocketRequest, PreparedNativeCall, call_native, ) -from litellm.rust_bridge.runtime import ( - BridgeErrorContext, - EndpointBinding, - async_none, - identity, -) +from litellm.rust_bridge.runtime import DispatchResult, aattempt, adapt_result 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, ) @@ -46,7 +42,7 @@ def set_rust_responses_websocket( _RESPONSES_WEBSOCKET.override(connection) -class _ConnectionAdapter: +class ConnectionAdapter: def __init__(self, connection: RustResponsesWebSocket): self._connection: Final[RustResponsesWebSocket] = connection @@ -68,18 +64,49 @@ async def connect( url: str, headers: dict[str, str], timeout: float | httpx.Timeout | None, -) -> _ConnectionAdapter | None: - connection: Final = await _RESPONSES_WEBSOCKET.ainvoke( + websocket_mode: str = "native", + requires_connection: bool = True, +) -> DispatchResult[ConnectionAdapter]: + return await aattempt( + load=_RESPONSES_WEBSOCKET.load, + enabled=rust_enabled(), + eligible=True, prepare=lambda: PreparedNativeCall( - NativeResponsesWebSocketRequest( - url=url, - ), + request=NativeResponsesWebSocketRequest(url=url), options=NativeRequestOptions(extra_headers=headers, timeout_seconds=timeout_to_seconds(timeout)), - context=NativeRequestContext(), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + websocket_mode=websocket_mode, + requires_connection=requires_connection, + ) + ), ), - call=lambda connection_type, request: call_native(connection_type.connect, request), - fallback=async_none, - adapt=identity, - error_context=BridgeErrorContext(provider="openai", model="responses websocket"), + call=lambda connection_type, prepared: call_native(connection_type.connect, prepared), + adapt=ConnectionAdapter, ) - return None if connection is None else _ConnectionAdapter(connection) + + +@asynccontextmanager +async def _connection_context(connection: ConnectionAdapter) -> AsyncGenerator[ConnectionAdapter, None]: + try: + yield connection + finally: + await connection.close() + + +async def managed_connect( + *, + url: str, + headers: dict[str, str], + timeout: float | httpx.Timeout | None, + websocket_mode: str = "managed", + requires_connection: bool = True, +) -> DispatchResult[AbstractAsyncContextManager[ConnectionAdapter]]: + result: Final = await connect( + url=url, + headers=headers, + timeout=timeout, + websocket_mode=websocket_mode, + requires_connection=requires_connection, + ) + return adapt_result(result, _connection_context) diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 0300cd84883..97f4eb485d1 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -4,9 +4,10 @@ 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.request import ( + NativeRequestCapabilities, NativeRequestContext, NativeRequestOptions, NativeTranscriptionRequest, @@ -14,47 +15,36 @@ from litellm.rust_bridge.request import ( bedrock_options, call_native, ) -from litellm.rust_bridge.runtime import ( - BridgeErrorContext, - EndpointDispatch, - 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, -) +_TRANSCRIPTION: Final[NativeBinding[RustTranscription]] = NativeBinding(lambda native: native.transcription) +_ATRANSCRIPTION: Final[NativeBinding[RustAtranscription]] = NativeBinding(lambda native: native.atranscription) def configure_rust_transcription( - enabled: bool = True, *, transcription: RustTranscription | None | Unchanged = UNCHANGED, atranscription: RustAtranscription | None | Unchanged = UNCHANGED, ) -> 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( @@ -67,14 +57,16 @@ 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( + stream: bool = False, + has_custom_client: bool = False, + input_source_kind: str | None = None, +) -> DispatchResult[dict[str, object]]: + return attempt( + load=_TRANSCRIPTION.load, + enabled=True, + eligible=True, prepare=lambda: PreparedNativeCall( - NativeTranscriptionRequest( - model=model, - audio=audio, - optional_params=optional_params, - ), + request=NativeTranscriptionRequest(model=model, audio=audio, optional_params=optional_params), options=NativeRequestOptions( api_key=api_key, api_base=api_base, @@ -83,12 +75,17 @@ def transcription( timeout_seconds=timeout_to_seconds(timeout), bedrock=bedrock_options(optional_params), ), - context=NativeRequestContext(), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + execution_mode="sync", + stream=stream, + has_custom_client=has_custom_client, + input_source_kind=input_source_kind, + ) + ), ), call=call_native, - fallback=lambda: None, adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) @@ -102,14 +99,16 @@ 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( + stream: bool = False, + has_custom_client: bool = False, + input_source_kind: str | None = None, +) -> DispatchResult[dict[str, object]]: + return await aattempt( + load=_ATRANSCRIPTION.load, + enabled=True, + eligible=True, prepare=lambda: PreparedNativeCall( - NativeTranscriptionRequest( - model=model, - audio=audio, - optional_params=optional_params, - ), + request=NativeTranscriptionRequest(model=model, audio=audio, optional_params=optional_params), options=NativeRequestOptions( api_key=api_key, api_base=api_base, @@ -118,10 +117,15 @@ async def atranscription( timeout_seconds=timeout_to_seconds(timeout), bedrock=bedrock_options(optional_params), ), - context=NativeRequestContext(), + context=NativeRequestContext( + capabilities=NativeRequestCapabilities( + execution_mode="async", + stream=stream, + has_custom_client=has_custom_client, + input_source_kind=input_source_kind, + ) + ), ), call=call_native, - fallback=async_none, adapt=identity, - error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model), ) diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 59787822ab6..7af9a080da6 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -9,7 +9,8 @@ import pytest import litellm from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import configuration -from litellm.rust_bridge.request import NativeMessagesRequest, NativeRequestContext +from litellm.rust_bridge.request import NativeMessagesRequest, NativeRequestContext, NativeRequestOptions +from litellm.rust_bridge.runtime import Handled, NativeFailed, NativeSkipped, NativeSkipReason from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -38,12 +39,13 @@ REQUEST_BODY: dict[str, object] = { class RecordingMessages: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.contexts: list[NativeRequestContext] = [] def __call__( self, request: NativeMessagesRequest, *, - options: object, + options: NativeRequestOptions, context: NativeRequestContext, ) -> dict[str, object]: self.calls.append( @@ -57,18 +59,20 @@ class RecordingMessages: "timeout_seconds": options.timeout_seconds, } ) + self.contexts.append(context) return dict(FAKE_MESSAGES_RESPONSE) class RecordingAsyncMessages: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.contexts: list[NativeRequestContext] = [] async def __call__( self, request: NativeMessagesRequest, *, - options: object, + options: NativeRequestOptions, context: NativeRequestContext, ) -> dict[str, object]: self.calls.append( @@ -82,6 +86,7 @@ class RecordingAsyncMessages: "timeout_seconds": options.timeout_seconds, } ) + self.contexts.append(context) return dict(FAKE_MESSAGES_RESPONSE) @@ -89,9 +94,7 @@ class ExplodingAsyncMessages: def __init__(self) -> None: self.calls = 0 - async def __call__( - self, request: NativeMessagesRequest, *, options: object, context: NativeRequestContext - ) -> dict[str, object]: + async def __call__(self, *args: object, **kwargs: object) -> dict[str, object]: self.calls += 1 raise AssertionError("bridge must not be called") @@ -100,9 +103,7 @@ class RaisingAsyncMessages: def __init__(self) -> None: self.calls = 0 - async def __call__( - self, request: NativeMessagesRequest, *, options: object, context: NativeRequestContext - ) -> dict[str, object]: + async def __call__(self, *args: object, **kwargs: object) -> dict[str, object]: self.calls += 1 raise RuntimeError("upstream request failed with status 400: bad request") @@ -142,7 +143,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", @@ -159,7 +160,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(): @@ -177,7 +178,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, @@ -205,16 +206,43 @@ 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 +@pytest.mark.asyncio +async def test_amessages_wrapper_preserves_capability_facts(): + bridge = RecordingAsyncMessages() + rust_messages.set_rust_messages(amessages=bridge) + + await rust_messages.amessages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key=None, + api_base=None, + custom_llm_provider="anthropic", + extra_headers=None, + timeout=None, + stream=True, + has_custom_client=True, + has_agentic_hook=True, + ) + + capabilities = bridge.contexts[0].capabilities + assert capabilities.execution_mode == "async" + assert capabilities.stream is True + assert capabilities.has_custom_client is True + assert capabilities.has_agentic_hook is True + + def _gate(**overrides): kwargs = { "custom_llm_provider": "azure_ai", "litellm_params": GenericLiteLLMParams(api_key="sk-azure"), "has_agentic_hook": False, + "stream": False, + "has_custom_client": False, "model": "claude-sonnet-4-5", "api_key": "sk-azure", "api_base": "https://resource.services.ai.azure.com/anthropic", @@ -223,7 +251,7 @@ def _gate(**overrides): "timeout": 30.0, } kwargs.update(overrides) - return BaseLLMHTTPHandler._maybe_rust_anthropic_messages(**kwargs) + return BaseLLMHTTPHandler._attempt_rust_anthropic_messages(**kwargs) @pytest.mark.asyncio @@ -234,7 +262,8 @@ async def test_gate_invokes_rust_and_marks_response_header(): response = await _gate() - assert response is not None + assert isinstance(response, Handled) + response = response.value assert response["id"] == "msg_123" assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} call = bridge.calls[0] @@ -247,13 +276,13 @@ async def test_gate_invokes_rust_and_marks_response_header(): @pytest.mark.asyncio -async def test_gate_propagates_unknown_native_errors(): +async def test_gate_reports_failure_to_harness(): bridge = RaisingAsyncMessages() litellm.rust(True) rust_messages.set_rust_messages(amessages=bridge) - with pytest.raises(RuntimeError, match="bad request"): - await _gate() + response = await _gate() + assert isinstance(response, NativeFailed) assert bridge.calls == 1 @@ -264,7 +293,7 @@ async def test_gate_skips_rust_when_flag_absent(): response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) - assert response is None + assert isinstance(response, NativeSkipped) assert bridge.calls == 0 @@ -276,7 +305,8 @@ async def test_gate_uses_process_enable_without_request_override(): response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) - assert response is not None + assert isinstance(response, Handled) + response = response.value assert bridge.calls[0]["custom_llm_provider"] == "azure_ai" @@ -288,7 +318,8 @@ async def test_gate_ignores_request_flag_when_process_enabled(): response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) - assert response is not None + assert isinstance(response, Handled) + response = response.value assert len(bridge.calls) == 1 @@ -306,7 +337,8 @@ async def test_gate_invokes_rust_for_native_anthropic_provider(): headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"}, ) - assert response is not None + assert isinstance(response, Handled) + response = response.value assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} assert bridge.calls[0]["custom_llm_provider"] == "anthropic" assert bridge.calls[0]["api_key"] == "sk-ant" @@ -323,7 +355,8 @@ async def test_gate_invokes_rust_when_env_var_set(monkeypatch): litellm_params=GenericLiteLLMParams(api_key="sk-ant"), ) - assert response is not None + assert isinstance(response, Handled) + response = response.value assert bridge.calls[0]["custom_llm_provider"] == "anthropic" @@ -338,7 +371,7 @@ async def test_gate_env_var_falsey_does_not_enable(monkeypatch): litellm_params=GenericLiteLLMParams(api_key="sk-ant"), ) - assert response is None + assert isinstance(response, NativeSkipped) assert bridge.calls == 0 @@ -350,7 +383,7 @@ async def test_gate_skips_rust_for_unsupported_provider(): response = await _gate(custom_llm_provider="openai") - assert response is None + assert isinstance(response, NativeSkipped) assert bridge.calls == 0 @@ -362,7 +395,7 @@ async def test_gate_skips_rust_for_agentic_hook(): response = await _gate(has_agentic_hook=True) - assert response is None + assert isinstance(response, NativeSkipped) assert bridge.calls == 0 @@ -378,7 +411,8 @@ async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): request_body=streaming_body, ) - assert response is not None + assert isinstance(response, Handled) + response = response.value assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} assert "stream" not in bridge.calls[0]["body"] assert bridge.calls[0]["body"] == REQUEST_BODY @@ -411,4 +445,105 @@ async def test_gate_falls_back_when_bridge_unavailable(monkeypatch): response = await _gate() - assert response is None + assert isinstance(response, NativeSkipped) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("selection", ("native", "disabled", "failed", "declined", "upstream")) +async def test_messages_handler_runs_selected_backend_once(selection: str, monkeypatch: pytest.MonkeyPatch) -> None: + from datetime import datetime + from types import SimpleNamespace + + import httpx + + from litellm.exceptions import RateLimitError + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import AnthropicMessagesConfig + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.rust_bridge import bindings + + class Declined(Exception): + pass + + class Upstream(Exception): + pass + + error = ( + Upstream(429, "rate limited") + if selection == "upstream" + else Declined("unsupported") + if selection == "declined" + else RuntimeError("native failed") + if selection == "failed" + else None + ) + + class Native: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, *args: object, **kwargs: object) -> dict[str, object]: + self.calls += 1 + if error is not None: + raise error + return dict(FAKE_MESSAGES_RESPONSE) + + bridge = Native() + monkeypatch.setattr( + bindings, + "get_native_bridge", + lambda: SimpleNamespace( + RustBridgeDeclined=Declined, + RustUpstreamError=Upstream, + ), + ) + rust_messages.set_rust_messages(amessages=bridge) + litellm.rust(selection != "disabled") + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json=FAKE_MESSAGES_RESPONSE) + + logging_obj = Logging( + model=FAKE_MESSAGES_RESPONSE["model"], + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="harness-test", + function_id="harness-test", + ) + client = AsyncHTTPHandler() + await client.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as transport: + client.client = transport + + async def run(): + return await BaseLLMHTTPHandler().async_anthropic_messages_handler( + model=FAKE_MESSAGES_RESPONSE["model"], + messages=[{"role": "user", "content": "hello"}], + anthropic_messages_provider_config=AnthropicMessagesConfig(), + anthropic_messages_optional_request_params={"max_tokens": 10}, + custom_llm_provider="anthropic", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + api_key="sk-test", + api_base="https://example.test", + client=client, + ) + + if selection in ("failed", "upstream"): + with pytest.raises(RateLimitError if selection == "upstream" else RuntimeError) as caught: + await run() + if selection == "upstream": + assert caught.value.__cause__ is error + assert caught.value.llm_provider == "anthropic" + assert caught.value.model == FAKE_MESSAGES_RESPONSE["model"] + else: + assert caught.value is error + else: + response = await run() + assert response["id"] == FAKE_MESSAGES_RESPONSE["id"] + assert len(requests) == (1 if selection in ("disabled", "declined") else 0) + assert bridge.calls == (0 if selection == "disabled" else 1) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 019022ed919..da70f422f44 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -23,7 +23,9 @@ async def test_make_call_passes_logging_obj_to_client_post(): mock_client = AsyncMock() mock_response = MagicMock() mock_response.aiter_lines = MagicMock( - return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']) + return_value=iter( + [b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n'] + ) ) mock_client.post.return_value = mock_response @@ -92,7 +94,9 @@ def test_redacted_thinking_content_block_delta(): "data": "EuoBCoYBGAIiQJ/SxkPAgqxhKok29YrpJHRUJ0OT8ahCHKAwyhmRuUhtdmDX9+mn4gDzKNv3fVpQdB01zEPMzNY3QuTCd+1bdtEqQK6JuKHqdndbwpr81oVWb4wxd1GqF/7Jkw74IlQa27oobX+KuRkopr9Dllt/RDe7Se0sI1IkU7tJIAQCoP46OAwSDF51P09q67xhHlQ3ihoM2aOVlkghq/X0w8NlIjBMNvXYNbjhyrOcIg6kPFn2ed/KK7Cm5prYAtXCwkb4Wr5tUSoSHu9T5hKdJRbr6WsqEc7Lle7FULqMLZGkhqXyc3BA", }, } - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=False, json_mode=False + ) model_response = model_response_iterator.chunk_parser(chunk=chunk) print(f"\n\nmodel_response: {model_response}\n\n") assert model_response.choices[0].delta.thinking_blocks is not None @@ -100,14 +104,19 @@ def test_redacted_thinking_content_block_delta(): print( f"\n\nmodel_response.choices[0].delta.thinking_blocks[0]: {model_response.choices[0].delta.thinking_blocks[0]}\n\n" ) - assert model_response.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking" + assert ( + model_response.choices[0].delta.thinking_blocks[0]["type"] + == "redacted_thinking" + ) assert model_response.choices[0].delta.provider_specific_fields is not None assert "thinking_blocks" in model_response.choices[0].delta.provider_specific_fields def test_streaming_thinking_blocks_are_replayable_after_signature_delta(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) chunks = [ { "type": "content_block_start", @@ -131,12 +140,17 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta(): }, ] - parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks] + parsed_chunks = [ + model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks + ] reasoning_content = "".join( - getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks + getattr(chunk.choices[0].delta, "reasoning_content", None) or "" + for chunk in parsed_chunks ) thinking_blocks = tuple( - block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) + block + for chunk in parsed_chunks + for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) ) expected_delta_blocks = ( {"type": "thinking", "thinking": "Step 1. "}, @@ -150,12 +164,18 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta(): assert reasoning_content == "Step 1. Step 2." assert thinking_blocks == (*expected_delta_blocks, expected_thinking_block) - assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_delta_blocks[0]]} - assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_thinking_block]} + assert parsed_chunks[1].choices[0].delta.provider_specific_fields == { + "thinking_blocks": [expected_delta_blocks[0]] + } + assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == { + "thinking_blocks": [expected_thinking_block] + } def test_streaming_unsigned_thinking_deltas_keep_reasoning_content(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) chunks = [ { "type": "content_block_start", @@ -175,12 +195,17 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content(): {"type": "content_block_stop", "index": 0}, ] - parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks] + parsed_chunks = [ + model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks + ] reasoning_content = "".join( - getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks + getattr(chunk.choices[0].delta, "reasoning_content", None) or "" + for chunk in parsed_chunks ) thinking_blocks = tuple( - block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) + block + for chunk in parsed_chunks + for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) ) assert reasoning_content == "Step 1. Step 2." @@ -191,7 +216,9 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content(): def test_streaming_truncated_thinking_deltas_keep_reasoning_content(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) chunks = [ { "type": "content_block_start", @@ -210,12 +237,17 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content(): }, ] - parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks] + parsed_chunks = [ + model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks + ] reasoning_content = "".join( - getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks + getattr(chunk.choices[0].delta, "reasoning_content", None) or "" + for chunk in parsed_chunks ) thinking_blocks = tuple( - block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) + block + for chunk in parsed_chunks + for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or []) ) assert reasoning_content == "Step 1. Step 2." @@ -226,7 +258,9 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content(): def test_handle_json_mode_chunk_response_format_tool(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=True + ) response_format_tool = ChatCompletionToolCallChunk( id="tool_123", type="function", @@ -237,7 +271,9 @@ def test_handle_json_mode_chunk_response_format_tool(): index=0, ) - text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool) + text, tool_use = model_response_iterator._handle_json_mode_chunk( + "", response_format_tool + ) print(f"\n\nresponse_format_tool text: {text}\n\n") print(f"\n\nresponse_format_tool tool_use: {tool_use}\n\n") @@ -246,11 +282,15 @@ def test_handle_json_mode_chunk_response_format_tool(): def test_handle_json_mode_chunk_regular_tool(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=True + ) regular_tool = ChatCompletionToolCallChunk( id="tool_456", type="function", - function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'), + function=ChatCompletionToolCallFunctionChunk( + name="get_weather", arguments='{"location": "San Francisco, CA"}' + ), index=0, ) @@ -264,13 +304,17 @@ def test_handle_json_mode_chunk_regular_tool(): def test_handle_json_mode_chunk_streaming_response_format_tool(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=True + ) # First chunk: response_format tool with id and name, but no arguments first_chunk = ChatCompletionToolCallChunk( id="tool_123", type="function", - function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments=""), + function=ChatCompletionToolCallFunctionChunk( + name=RESPONSE_FORMAT_TOOL_NAME, arguments="" + ), index=0, ) @@ -278,7 +322,9 @@ def test_handle_json_mode_chunk_streaming_response_format_tool(): second_chunk = ChatCompletionToolCallChunk( id=None, type="function", - function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"question": "What is the weather?"'), + function=ChatCompletionToolCallFunctionChunk( + name=None, arguments='{"question": "What is the weather?"' + ), index=0, ) @@ -286,7 +332,9 @@ def test_handle_json_mode_chunk_streaming_response_format_tool(): third_chunk = ChatCompletionToolCallChunk( id=None, type="function", - function=ChatCompletionToolCallFunctionChunk(name=None, arguments=', "answer": "It is sunny"}'), + function=ChatCompletionToolCallFunctionChunk( + name=None, arguments=', "answer": "It is sunny"}' + ), index=0, ) @@ -317,7 +365,9 @@ def test_handle_json_mode_chunk_streaming_response_format_tool(): def test_handle_json_mode_chunk_streaming_regular_tool(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=True + ) # First chunk: regular tool with id and name, but no arguments first_chunk = ChatCompletionToolCallChunk( @@ -331,7 +381,9 @@ def test_handle_json_mode_chunk_streaming_regular_tool(): second_chunk = ChatCompletionToolCallChunk( id=None, type="function", - function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"location": "San Francisco, CA"}'), + function=ChatCompletionToolCallFunctionChunk( + name=None, arguments='{"location": "San Francisco, CA"}' + ), index=0, ) @@ -356,19 +408,27 @@ def test_handle_json_mode_chunk_streaming_regular_tool(): def test_response_format_tool_finish_reason(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=True + ) # First chunk: response_format tool response_format_tool = ChatCompletionToolCallChunk( id="tool_123", type="function", - function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}'), + function=ChatCompletionToolCallFunctionChunk( + name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}' + ), index=0, ) # Process the tool call (should set converted_response_format_tool flag) - text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool) - print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n") + text, tool_use = model_response_iterator._handle_json_mode_chunk( + "", response_format_tool + ) + print( + f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n" + ) # Simulate message_delta chunk with tool_use stop_reason message_delta_chunk = { @@ -387,19 +447,25 @@ def test_response_format_tool_finish_reason(): def test_regular_tool_finish_reason(): - model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True) + model_response_iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=True + ) # First chunk: regular tool (not response_format) regular_tool = ChatCompletionToolCallChunk( id="tool_456", type="function", - function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'), + function=ChatCompletionToolCallFunctionChunk( + name="get_weather", arguments='{"location": "San Francisco, CA"}' + ), index=0, ) # Process the tool call (should NOT set converted_response_format_tool flag) text, tool_use = model_response_iterator._handle_json_mode_chunk("", regular_tool) - print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n") + print( + f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n" + ) # Simulate message_delta chunk with tool_use stop_reason message_delta_chunk = { @@ -459,7 +525,9 @@ def test_text_only_streaming_has_index_zero(): for chunk in chunks: parsed = iterator.chunk_parser(chunk) if parsed.choices: - assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}" + assert ( + parsed.choices[0].index == 0 + ), f"Expected index=0, got {parsed.choices[0].index}" def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage(): @@ -636,7 +704,9 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin ] self._write_response( content_type="text/event-stream", - body="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode("utf-8"), + body="".join( + f"data: {json.dumps(event)}\n\n" for event in events + ).encode("utf-8"), ) return @@ -717,9 +787,13 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin assert content_chunks == [answer_text] assert stream_usage is not None stream_completion_details = stream_usage["completion_tokens_details"] - assert stream_completion_details["reasoning_tokens"] == non_stream_details.reasoning_tokens + assert ( + stream_completion_details["reasoning_tokens"] + == non_stream_details.reasoning_tokens + ) assert stream_completion_details["text_tokens"] == ( - stream_usage["completion_tokens"] - stream_completion_details["reasoning_tokens"] + stream_usage["completion_tokens"] + - stream_completion_details["reasoning_tokens"] ) assert requests_seen == [ { @@ -811,9 +885,9 @@ def test_text_and_tool_streaming_has_index_zero(): for chunk in chunks: parsed = iterator.chunk_parser(chunk) if parsed.choices: - assert parsed.choices[0].index == 0, ( - f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}" - ) + assert ( + parsed.choices[0].index == 0 + ), f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}" def test_multiple_tools_streaming_has_index_zero(): @@ -866,11 +940,15 @@ def test_multiple_tools_streaming_has_index_zero(): for chunk in chunks: parsed = iterator.chunk_parser(chunk) if parsed.choices: - assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}" + assert ( + parsed.choices[0].index == 0 + ), f"Expected index=0, got {parsed.choices[0].index}" def test_streaming_chunks_have_stable_ids(): - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=False, json_mode=False + ) first_chunk = { "type": "content_block_delta", "index": 0, @@ -895,7 +973,9 @@ def test_partial_json_chunk_accumulation(): This tests the fix for https://github.com/BerriAI/litellm/issues/17473 where network fragmentation can cause SSE data to arrive in partial chunks. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) partial_chunk_1 = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel' partial_chunk_2 = 'lo"}}' @@ -903,21 +983,31 @@ def test_partial_json_chunk_accumulation(): # First partial chunk should return None (still accumulating) result1 = iterator._parse_sse_data(f"data:{partial_chunk_1}") assert result1 is None, "First partial chunk should return None while accumulating" - assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode" - assert iterator.accumulated_json == partial_chunk_1, "Should have accumulated first part" + assert ( + iterator.chunk_type == "accumulated_json" + ), "Should switch to accumulated_json mode" + assert ( + iterator.accumulated_json == partial_chunk_1 + ), "Should have accumulated first part" # Second partial chunk should complete the JSON and return a parsed result result2 = iterator._parse_sse_data(f"data:{partial_chunk_2}") assert result2 is not None, "Second chunk should return parsed result" - assert iterator.accumulated_json == "", "Buffer should be cleared after successful parse" - assert result2.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result2.choices[0].delta.content}'" + assert ( + iterator.accumulated_json == "" + ), "Buffer should be cleared after successful parse" + assert ( + result2.choices[0].delta.content == "Hello" + ), f"Expected 'Hello', got '{result2.choices[0].delta.content}'" def test_complete_json_chunk_no_accumulation(): """ Test that complete JSON chunks are parsed immediately without accumulation. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) complete_chunk = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}' @@ -925,14 +1015,18 @@ def test_complete_json_chunk_no_accumulation(): assert result is not None, "Complete chunk should return parsed result immediately" assert iterator.chunk_type == "valid_json", "Should remain in valid_json mode" assert iterator.accumulated_json == "", "Buffer should remain empty" - assert result.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result.choices[0].delta.content}'" + assert ( + result.choices[0].delta.content == "Hello" + ), f"Expected 'Hello', got '{result.choices[0].delta.content}'" def test_multiple_partial_chunks_accumulation(): """ Test that multiple partial chunks can be accumulated across several iterations. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) # Split a JSON chunk into three parts part1 = '{"type":"content_block_del' @@ -960,11 +1054,17 @@ def test_accumulated_json_partial_fragment_returns_none_without_parsing(): unlike Vertex which already deferred parsing until the buffer could close. A fragment that can't close a JSON value must not trigger a decode attempt. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) iterator.chunk_type = "accumulated_json" - with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: - result = iterator._handle_accumulated_json_chunk('{"type":"content_block_delta","index":0,"delta":') + with patch.object( + json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode + ) as spy: + result = iterator._handle_accumulated_json_chunk( + '{"type":"content_block_delta","index":0,"delta":' + ) assert result is None assert spy.call_count == 0, "incomplete buffer should not be parsed" @@ -976,15 +1076,21 @@ def test_accumulated_json_does_not_reparse_every_fragment(): fragment. """ text = "x" * 200_000 - blob = json.dumps({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}) + blob = json.dumps( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}} + ) fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)] assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) iterator.chunk_type = "accumulated_json" parsed = None - with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: + with patch.object( + json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode + ) as spy: for fragment in fragments: out = iterator._handle_accumulated_json_chunk(fragment) if out is not None: @@ -1008,7 +1114,9 @@ def test_accumulated_json_concatenated_envelopes_do_not_wedge(): and keeps the remainder, so both values surface across two calls. """ obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) iterator.chunk_type = "accumulated_json" first = iterator._handle_accumulated_json_chunk(obj + obj) @@ -1029,7 +1137,9 @@ def test_accumulated_json_heuristic_passes_but_value_still_incomplete(): heuristic must let the parse attempt through, and pop_next_value finding nothing must propagate as None rather than raising. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) iterator.chunk_type = "accumulated_json" result = iterator._handle_accumulated_json_chunk('{"type": {"nested": 1}') @@ -1043,7 +1153,9 @@ def test_accumulated_json_setter_and_sync_end_of_stream_drain(): underlying stream ends, instead of being silently dropped. """ obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' - iterator = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=iter([]), sync_stream=True, json_mode=False + ) iterator.chunk_type = "accumulated_json" iterator.accumulated_json = obj # exercises the setter @@ -1058,7 +1170,9 @@ def test_accumulated_json_async_end_of_stream_drain(): import asyncio obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=False, json_mode=False + ) iterator.chunk_type = "accumulated_json" iterator.accumulated_json = obj mock_async_iterator = MagicMock() @@ -1080,7 +1194,9 @@ def test_web_search_tool_result_no_extra_tool_calls(): The issue was that web_search_tool_result blocks have input_json_delta events with {} that were incorrectly being converted to tool calls. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) # Simulate the streaming sequence: # 1. server_tool_use block starts (web_search) @@ -1155,7 +1271,9 @@ def test_web_search_tool_result_no_extra_tool_calls(): # Should have exactly 2 tool calls: # 1. From content_block_start (server_tool_use) with id and name # 2. From content_block_delta with the actual query - assert len(tool_calls_emitted) == 2, f"Expected 2 tool calls, got {len(tool_calls_emitted)}" + assert ( + len(tool_calls_emitted) == 2 + ), f"Expected 2 tool calls, got {len(tool_calls_emitted)}" # First tool call should have the id and name assert tool_calls_emitted[0]["id"] == "srvtoolu_01ABC123" @@ -1171,7 +1289,9 @@ def test_current_content_block_type_tracking(): """ Test that current_content_block_type is properly tracked and reset. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) # Initially should be None assert iterator.current_content_block_type is None @@ -1224,7 +1344,9 @@ def test_web_search_tool_result_captured_in_provider_specific_fields(): The web_search_tool_result content comes ALL AT ONCE in content_block_start, not in deltas, so we need to capture it there. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) # Simulate the streaming sequence with web_search_tool_result chunks = [ @@ -1295,15 +1417,23 @@ def test_web_search_tool_result_captured_in_provider_specific_fields(): and parsed.choices[0].delta.provider_specific_fields and "web_search_results" in parsed.choices[0].delta.provider_specific_fields ): - web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"] + web_search_results = parsed.choices[0].delta.provider_specific_fields[ + "web_search_results" + ] # Verify web_search_results was captured assert web_search_results is not None, "web_search_results should be captured" assert len(web_search_results) == 1, "Should have 1 web_search_tool_result block" - assert web_search_results[0]["type"] == "web_search_tool_result", "Block type should be web_search_tool_result" - assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match" + assert ( + web_search_results[0]["type"] == "web_search_tool_result" + ), "Block type should be web_search_tool_result" + assert ( + web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123" + ), "tool_use_id should match" assert len(web_search_results[0]["content"]) == 2, "Should have 2 search results" - assert web_search_results[0]["content"][0]["title"] == "Fun Otter Facts", "First result title should match" + assert ( + web_search_results[0]["content"][0]["title"] == "Fun Otter Facts" + ), "First result title should match" def test_web_fetch_tool_result_captured_in_provider_specific_fields(): @@ -1317,7 +1447,9 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields(): The web_fetch_tool_result content comes ALL AT ONCE in content_block_start, not in deltas, so we need to capture it there. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) # Simulate the streaming sequence with web_fetch_tool_result chunks = [ @@ -1388,15 +1520,25 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields(): and parsed.choices[0].delta.provider_specific_fields and "web_search_results" in parsed.choices[0].delta.provider_specific_fields ): - web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"] + web_search_results = parsed.choices[0].delta.provider_specific_fields[ + "web_search_results" + ] # Verify web_fetch_tool_result was captured (stored in web_search_results list) assert web_search_results is not None, "web_search_results should be captured" assert len(web_search_results) == 1, "Should have 1 web_fetch_tool_result block" - assert web_search_results[0]["type"] == "web_fetch_tool_result", "Block type should be web_fetch_tool_result" - assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match" - assert web_search_results[0]["content"]["url"] == "https://example.com", "URL should match" - assert web_search_results[0]["content"]["content"]["title"] == "Example Page", "Title should match" + assert ( + web_search_results[0]["type"] == "web_fetch_tool_result" + ), "Block type should be web_fetch_tool_result" + assert ( + web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123" + ), "tool_use_id should match" + assert ( + web_search_results[0]["content"]["url"] == "https://example.com" + ), "URL should match" + assert ( + web_search_results[0]["content"]["content"]["title"] == "Example Page" + ), "Title should match" def test_web_fetch_tool_result_no_extra_tool_calls(): @@ -1409,7 +1551,9 @@ def test_web_fetch_tool_result_no_extra_tool_calls(): The issue was that web_fetch_tool_result blocks have input_json_delta events with {} that were incorrectly being converted to tool calls. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) # to verify it doesn't emit tool calls chunks = [ @@ -1453,9 +1597,9 @@ def test_web_fetch_tool_result_no_extra_tool_calls(): tool_call_count += 1 # Should have 0 tool calls - web_fetch_tool_result should not emit tool calls - assert tool_call_count == 0, ( - f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls" - ) + assert ( + tool_call_count == 0 + ), f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls" def test_container_in_provider_specific_fields_streaming(): @@ -1465,7 +1609,9 @@ def test_container_in_provider_specific_fields_streaming(): When container with skills is used, the container field should be present in the provider_specific_fields of the message_delta chunk. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) # Simulate streaming chunks chunks = [ @@ -1533,12 +1679,20 @@ def test_container_in_provider_specific_fields_streaming(): and parsed.choices[0].delta.provider_specific_fields and "container" in parsed.choices[0].delta.provider_specific_fields ): - container_field = parsed.choices[0].delta.provider_specific_fields["container"] + container_field = parsed.choices[0].delta.provider_specific_fields[ + "container" + ] # Verify container was captured - assert container_field is not None, "container should be captured in provider_specific_fields" - assert container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p", "container id should match" - assert container_field["expires_at"] == "2025-12-16T04:57:16.913181Z", "expires_at should match" + assert ( + container_field is not None + ), "container should be captured in provider_specific_fields" + assert ( + container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p" + ), "container id should match" + assert ( + container_field["expires_at"] == "2025-12-16T04:57:16.913181Z" + ), "expires_at should match" assert len(container_field["skills"]) == 1, "Should have 1 skill" assert container_field["skills"][0]["skill_id"] == "pptx", "skill_id should be pptx" assert container_field["skills"][0]["version"] == "20251013", "version should match" @@ -1551,7 +1705,9 @@ def test_container_in_provider_specific_fields_non_streaming(): When container with skills is used in non-streaming, the container field should be present in the provider_specific_fields of the response. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=False, json_mode=False + ) # Simulate a message_delta chunk with container (as it would appear in non-streaming) message_delta_chunk = { @@ -1587,13 +1743,21 @@ def test_container_in_provider_specific_fields_non_streaming(): # Verify container is in provider_specific_fields assert model_response.choices[0].delta.provider_specific_fields is not None assert "container" in model_response.choices[0].delta.provider_specific_fields - container_field = model_response.choices[0].delta.provider_specific_fields["container"] + container_field = model_response.choices[0].delta.provider_specific_fields[ + "container" + ] assert container_field["id"] == "container_abc123xyz", "container id should match" - assert container_field["expires_at"] == "2025-12-20T10:30:00.000000Z", "expires_at should match" + assert ( + container_field["expires_at"] == "2025-12-20T10:30:00.000000Z" + ), "expires_at should match" assert len(container_field["skills"]) == 2, "Should have 2 skills" - assert container_field["skills"][0]["skill_id"] == "code_execution", "First skill_id should be code_execution" - assert container_field["skills"][1]["skill_id"] == "pptx", "Second skill_id should be pptx" + assert ( + container_field["skills"][0]["skill_id"] == "code_execution" + ), "First skill_id should be code_execution" + assert ( + container_field["skills"][1]["skill_id"] == "pptx" + ), "Second skill_id should be pptx" def test_container_absent_when_not_provided(): @@ -1602,7 +1766,9 @@ def test_container_absent_when_not_provided(): This ensures we don't add empty or None container fields. """ - iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=False, json_mode=False + ) # message_delta without container message_delta_chunk = { @@ -1621,9 +1787,9 @@ def test_container_absent_when_not_provided(): # Verify container is NOT in provider_specific_fields when not provided if model_response.choices[0].delta.provider_specific_fields: - assert "container" not in model_response.choices[0].delta.provider_specific_fields, ( - "container should not be present when not provided in delta" - ) + assert ( + "container" not in model_response.choices[0].delta.provider_specific_fields + ), "container should not be present when not provided in delta" def test_streaming_code_execution_produces_code_interpreter_results(): @@ -1819,7 +1985,8 @@ def test_streaming_multiple_code_executions_no_duplicates(): # Second (final) emission: cumulative list with BOTH results # This is what stream_chunk_builder will pick as "last value wins" assert len(emissions[1]) == 2, ( - f"Expected final emission to have 2 results, got {len(emissions[1])}. IDs: {[r.id for r in emissions[1]]}" + f"Expected final emission to have 2 results, got {len(emissions[1])}. " + f"IDs: {[r.id for r in emissions[1]]}" ) assert emissions[1][0].id == "srvtoolu_01AAA" assert emissions[1][0].code == "echo first" @@ -1983,7 +2150,9 @@ def test_empty_output_produces_null_outputs(): assert code_results is not None, "No code_interpreter_results emitted" assert len(code_results) == 1 assert code_results[0].id == "srvtoolu_01AAA" - assert code_results[0].outputs is None, f"Expected outputs=None for empty execution, got {code_results[0].outputs}" + assert ( + code_results[0].outputs is None + ), f"Expected outputs=None for empty execution, got {code_results[0].outputs}" def test_non_bash_tool_result_skipped(): @@ -2046,10 +2215,12 @@ def test_non_bash_tool_result_skipped(): code_results = psf["code_interpreter_results"] # code_interpreter_results should be emitted but empty (no bash results) - assert code_results is not None, "Expected code_interpreter_results key to be emitted" - assert len(code_results) == 0, ( - f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}" - ) + assert ( + code_results is not None + ), "Expected code_interpreter_results key to be emitted" + assert ( + len(code_results) == 0 + ), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}" class TestRustChatCompletionsHook: @@ -2086,9 +2257,13 @@ class TestRustChatCompletionsHook: from litellm.rust_bridge import chat_completions as bridge monkeypatch.setenv("LITELLM_RUST", "1") - bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) yield - bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) @staticmethod def _completion_kwargs(**overrides): @@ -2134,19 +2309,8 @@ class TestRustChatCompletionsHook: seen["gate"].append(kwargs) return decline_reason - def native(request, *, options, context): - seen["call"].append( - { - "model": request.model, - "messages": request.messages, - "optional_params": request.optional_params, - "api_key": options.api_key, - "api_base": options.api_base, - "custom_llm_provider": options.custom_llm_provider, - "extra_headers": options.extra_headers, - "timeout_seconds": options.timeout_seconds, - } - ) + def native(**kwargs): + seen["call"].append(kwargs) if sync_error is not None: raise sync_error return dict(sync_result if sync_result is not None else self.RUST_RESPONSE) @@ -2197,7 +2361,9 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion seen = self._inject() - AnthropicChatCompletion().completion(**self._completion_kwargs(optional_params={"max_tokens": 7})) + AnthropicChatCompletion().completion( + **self._completion_kwargs(optional_params={"max_tokens": 7}) + ) assert seen["call"][0]["optional_params"]["max_tokens"] == 7 def test_without_the_opt_in_the_core_is_never_consulted(self, monkeypatch): @@ -2206,14 +2372,15 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.transformation import AnthropicConfig seen = self._inject() - with ( - patch.object( - AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} - ) as transform, - patch.object(AnthropicChatCompletion, "acompletion_function"), + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ) as transform, patch.object( + AnthropicChatCompletion, "acompletion_function" ): try: - AnthropicChatCompletion().completion(**self._completion_kwargs(litellm_params={})) + AnthropicChatCompletion().completion( + **self._completion_kwargs(litellm_params={}) + ) except Exception: # The Python path goes on to make an HTTP call; reaching it is # the assertion, so the network failure below is expected. @@ -2227,7 +2394,9 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.transformation import AnthropicConfig seen = self._inject(decline_reason="unrecognized request parameter") - with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): try: AnthropicChatCompletion().completion(**self._completion_kwargs()) except Exception: @@ -2240,7 +2409,9 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.transformation import AnthropicConfig seen = self._inject() - with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): try: AnthropicChatCompletion().completion( **self._completion_kwargs(optional_params={"max_tokens": 16, "stream": True}) @@ -2254,7 +2425,9 @@ class TestRustChatCompletionsHook: seen = self._inject() logging_obj = MagicMock() - AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) + AnthropicChatCompletion().completion( + **self._completion_kwargs(logging_obj=logging_obj) + ) assert logging_obj.pre_call.call_count == 1 assert len(seen["call"]) == 1 @@ -2268,7 +2441,9 @@ class TestRustChatCompletionsHook: self._inject() logging_obj = MagicMock() - AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) + AnthropicChatCompletion().completion( + **self._completion_kwargs(logging_obj=logging_obj) + ) assert logging_obj.post_call.call_count == 1 logged = logging_obj.post_call.call_args.kwargs["original_response"] @@ -2289,16 +2464,22 @@ class TestRustChatCompletionsHook: RustBridgeDeclined = _Declined RustUpstreamError = type("_Upstream", (Exception,), {}) - def declining_native(request, *, options, context): + def declining_native(**_kwargs): raise _Declined("blank message text") monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()) - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, chat_completions=declining_native + ) logging_obj, calls = self._recording_logging_obj() - with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): try: - AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) + AnthropicChatCompletion().completion( + **self._completion_kwargs(logging_obj=logging_obj) + ) except Exception: # The Python path goes on to make an HTTP call; the log count is # the assertion, so a failure past this point is expected. @@ -2320,18 +2501,24 @@ class TestRustChatCompletionsHook: monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()) - async def declining_native(request, *, options, context): + async def declining_native(**_kwargs): raise _Declined("blank message text") - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=declining_native + ) sentinel = object() async def python_path(**_kwargs): return sentinel - with patch.object(AnthropicChatCompletion, "acompletion_function", side_effect=python_path) as python_call: - result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True)) + with patch.object( + AnthropicChatCompletion, "acompletion_function", side_effect=python_path + ) as python_call: + result = await AnthropicChatCompletion().completion( + **self._completion_kwargs(acompletion=True) + ) assert result is sentinel assert python_call.called, "a failing rust call must re-enter the python path" @@ -2341,18 +2528,23 @@ class TestRustChatCompletionsHook: from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from litellm.rust_bridge import chat_completions as bridge - async def native(request, *, options, context): + async def native(**_kwargs): return dict(self.RUST_RESPONSE) - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=native + ) with patch.object(AnthropicChatCompletion, "acompletion_function") as python_call: - result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True)) + result = await AnthropicChatCompletion().completion( + **self._completion_kwargs(acompletion=True) + ) assert result.choices[0].message.content == "hello from rust" assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} assert not python_call.called + def test_pre_call_logging_fires_once_when_the_sync_rust_call_declines(self, monkeypatch): """One request, one pre_call, on the synchronous path too. Without the suppression the Python path logs a second time for the same attempt.""" @@ -2369,22 +2561,30 @@ class TestRustChatCompletionsHook: monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()) - def declining_native(request, *, options, context): + def declining_native(**_kwargs): raise _Declined("blank message text") - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, chat_completions=declining_native + ) logging_obj, calls = self._recording_logging_obj() - with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): try: - AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) + AnthropicChatCompletion().completion( + **self._completion_kwargs(logging_obj=logging_obj) + ) except Exception: # The Python path goes on to make an HTTP call; the log count is # the assertion, so a failure past this point is expected. pass assert len(calls["pre_call"]) == 1 - assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == ("claude-sonnet-4-5") + assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == ( + "claude-sonnet-4-5" + ) def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch): """The suppression must not swallow the log on the ordinary path.""" @@ -2394,7 +2594,9 @@ class TestRustChatCompletionsHook: self._inject() logging_obj, calls = self._recording_logging_obj() - with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): try: AnthropicChatCompletion().completion( **self._completion_kwargs(litellm_params={}, logging_obj=logging_obj) diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index c614e040d4c..becd6ecb832 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -10,8 +10,8 @@ from unittest.mock import MagicMock, patch import httpx import pytest -from botocore.credentials import Credentials +from botocore.credentials import Credentials from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import chat_completions as bridge @@ -49,9 +49,13 @@ RESOLVED_CREDENTIALS = Credentials( @pytest.fixture(autouse=True) def reset_bridge(monkeypatch): monkeypatch.setenv("LITELLM_RUST", "1") - bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) yield - bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) def _inject(*, decline_reason=None, error: Exception | None = None): @@ -61,20 +65,8 @@ def _inject(*, decline_reason=None, error: Exception | None = None): seen["gate"].append(kwargs) return decline_reason - def native(request, *, options, context): - seen["call"].append( - { - "model": request.model, - "messages": request.messages, - "optional_params": request.optional_params, - "bedrock": options.bedrock, - "api_key": options.api_key, - "api_base": options.api_base, - "custom_llm_provider": options.custom_llm_provider, - "extra_headers": options.extra_headers, - "timeout_seconds": options.timeout_seconds, - } - ) + def native(**kwargs): + seen["call"].append(kwargs) if error is not None: raise error return dict(RUST_RESPONSE) @@ -134,18 +126,20 @@ def test_the_core_receives_the_credentials_this_handler_already_resolved(): seen = _inject() _run() - bedrock = seen["call"][0]["bedrock"] - assert bedrock.aws_access_key_id == "AKIARESOLVED" - assert bedrock.aws_secret_access_key == "resolved-secret" - assert bedrock.aws_session_token == "resolved-token" - assert bedrock.aws_region_name == "us-east-1" + params = seen["call"][0]["optional_params"] + assert params["aws_access_key_id"] == "AKIARESOLVED" + assert params["aws_secret_access_key"] == "resolved-secret" + assert params["aws_session_token"] == "resolved-token" + assert params["aws_region_name"] == "us-east-1" def test_the_core_receives_the_converse_url_this_handler_already_built(): seen = _inject() _run() - assert seen["call"][0]["api_base"].endswith("/model/anthropic.claude-sonnet-4-5-v1%3A0/converse") + assert seen["call"][0]["api_base"].endswith( + "/model/anthropic.claude-sonnet-4-5-v1%3A0/converse" + ) assert "bedrock-runtime.us-east-1.amazonaws.com" in seen["call"][0]["api_base"] @@ -213,10 +207,12 @@ async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch): monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()) - async def declining_native(request, *, options, context): + async def declining_native(**_kwargs): raise _Declined("blank message text") - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=declining_native + ) sentinel = object() @@ -224,10 +220,16 @@ async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch): return sentinel with ( - patch.object(BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS), - patch.object(BedrockConverseLLM, "async_completion", side_effect=python_path) as python_call, + patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ), + patch.object( + BedrockConverseLLM, "async_completion", side_effect=python_path + ) as python_call, ): - result = await BedrockConverseLLM().completion(**_completion_kwargs(acompletion=True)) + result = await BedrockConverseLLM().completion( + **_completion_kwargs(acompletion=True) + ) assert result is sentinel assert python_call.called, "a failing rust call must re-enter the python path" @@ -235,16 +237,22 @@ async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch): @pytest.mark.asyncio async def test_the_async_path_serves_the_rust_response_without_the_fallback(): - async def native(request, *, options, context): + async def native(**_kwargs): return dict(RUST_RESPONSE) - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=native + ) with ( - patch.object(BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS), + patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ), patch.object(BedrockConverseLLM, "async_completion") as python_call, ): - result = await BedrockConverseLLM().completion(**_completion_kwargs(acompletion=True)) + result = await BedrockConverseLLM().completion( + **_completion_kwargs(acompletion=True) + ) assert result.choices[0].message.content == "hello from rust" assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} @@ -263,7 +271,7 @@ async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines(): RustBridgeDeclined = _Declined RustUpstreamError = type("_Upstream", (Exception,), {}) - async def declining_native(request, *, options, context): + async def declining_native(**_kwargs): raise _Declined("blank message text") logging_obj = MagicMock() @@ -275,11 +283,19 @@ async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines(): with ( patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()), - patch.object(BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS), - patch.object(BedrockConverseLLM, "async_completion", side_effect=python_path), + patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ), + patch.object( + BedrockConverseLLM, "async_completion", side_effect=python_path + ), ): - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native) - await BedrockConverseLLM().completion(**_completion_kwargs(acompletion=True, logging_obj=logging_obj)) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=declining_native + ) + await BedrockConverseLLM().completion( + **_completion_kwargs(acompletion=True, logging_obj=logging_obj) + ) assert logging_obj.pre_call.call_count == 1 assert served and served[0]["skip_pre_call_logging"] is True @@ -368,13 +384,15 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines(): RustBridgeDeclined = _Declined RustUpstreamError = type("_Upstream", (Exception,), {}) - def declining_native(request, *, options, context): + def declining_native(**_kwargs): raise _Declined("blank message text") logging_obj = MagicMock() with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()): - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, chat_completions=declining_native + ) response = _run( logging_obj=logging_obj, client=_sync_client_returning_converse_response(), @@ -420,14 +438,20 @@ async def test_post_call_logging_fires_on_the_async_rust_path(): cannot drift apart the way the pre_call suppression once did.""" import json - async def native(request, *, options, context): + async def native(**_kwargs): return dict(RUST_RESPONSE) - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=native + ) logging_obj = MagicMock() - with patch.object(BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS): - await BedrockConverseLLM().completion(**_completion_kwargs(acompletion=True, logging_obj=logging_obj)) + with patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ): + await BedrockConverseLLM().completion( + **_completion_kwargs(acompletion=True, logging_obj=logging_obj) + ) assert logging_obj.post_call.call_count == 1 logged = logging_obj.post_call.call_args.kwargs["original_response"] @@ -446,13 +470,15 @@ def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines(): RustBridgeDeclined = _Declined RustUpstreamError = type("_Upstream", (Exception,), {}) - def declining_native(request, *, options, context): + def declining_native(**_kwargs): raise _Declined("blank message text") logging_obj, calls = _recording_logging_obj() with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()): - bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, chat_completions=declining_native + ) response = _run( logging_obj=logging_obj, client=_sync_client_returning_converse_response(), @@ -486,11 +512,9 @@ def test_the_rust_opt_in_needs_no_sigv4_principal(): response = _run(credentials=None, api_key="bedrock-bearer-token") assert response.choices[0].message.content == "hello from rust" - bedrock = seen["call"][0]["bedrock"] - assert bedrock.aws_access_key_id is None - assert bedrock.aws_secret_access_key is None - assert bedrock.aws_session_token is None - assert bedrock.aws_region_name == "us-east-1" + params = seen["call"][0]["optional_params"] + assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys() + assert params["aws_region_name"] == "us-east-1" assert seen["call"][0]["api_key"] == "bedrock-bearer-token" diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 2a9b314ac85..cd30a0b698f 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -12,13 +12,8 @@ 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, - NativeVertexOptions, - PreparedNativeCall, -) +from litellm.rust_bridge.request import NativeOCRRequest, NativeRequestContext, NativeRequestOptions +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` @@ -52,6 +47,7 @@ class RecordingBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.contexts: list[NativeRequestContext] = [] def __call__( self, @@ -69,10 +65,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) @@ -81,6 +77,7 @@ class RecordingAsyncBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.contexts: list[NativeRequestContext] = [] async def __call__( self, @@ -98,10 +95,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) @@ -156,6 +153,9 @@ class FakeOCRConfig: def get_api_key_env_var(self) -> str: return self.api_key_env_var + def supports_rust_bridge(self) -> bool: + return True + def validate_environment( self, *, @@ -192,7 +192,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, @@ -381,105 +381,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: 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._OCR.override(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", @@ -490,8 +398,13 @@ 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, @@ -503,7 +416,6 @@ 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, } @@ -513,8 +425,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): litellm.rust(True) rust_bridge._OCR.override(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, ) @@ -530,8 +441,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, @@ -552,8 +462,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", @@ -572,8 +481,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): litellm.rust(True) rust_bridge._OCR.override(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", @@ -588,8 +496,11 @@ 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} - assert bridge.calls[0]["vertex"] == NativeVertexOptions(project="project-1", location="us-central1") + assert bridge.calls[0]["optional_params"] == { + "include_image_base64": True, + "vertex_project": "project-1", + "vertex_location": "us-central1", + } def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): @@ -603,8 +514,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", @@ -613,7 +523,8 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana resolve_api_key=_resolver, ) - assert bridge.calls[0]["vertex"] == NativeVertexOptions(project="project-from-secret", location="us-east5") + assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" + assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): @@ -621,8 +532,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): litellm.rust(True) rust_bridge._OCR.override(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", @@ -640,8 +550,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): litellm.rust(True) rust_bridge._OCR.override(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", @@ -662,8 +571,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): litellm.rust(True) rust_bridge._OCR.override(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", @@ -810,7 +718,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 = ( @@ -823,6 +731,25 @@ 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/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index aea5137db66..de8f91b6580 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -4,7 +4,12 @@ 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.request import NativeRequestContext, NativeResponsesWebSocketRequest +from litellm.rust_bridge.request import ( + NativeRequestContext, + NativeRequestOptions, + NativeResponsesWebSocketRequest, +) +from litellm.rust_bridge.runtime import Handled, NativeFailed, NativeSkipped, NativeSkipReason class _FakeNativeConnection: @@ -28,14 +33,17 @@ class _ClosedNativeConnection: class _FakeNativeBridge: + contexts: list[NativeRequestContext] = [] + @classmethod async def connect( cls, request: NativeResponsesWebSocketRequest, *, - options: object, + options: NativeRequestOptions, context: NativeRequestContext, ) -> _FakeNativeConnection: + cls.contexts.append(context) return _FakeNativeConnection() @@ -58,25 +66,22 @@ def test_rust_websocket_bridge_uses_process_enablement() -> None: @pytest.mark.asyncio async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: - adapter = responses_websocket._ConnectionAdapter(_ClosedNativeConnection()) + adapter = responses_websocket.ConnectionAdapter(_ClosedNativeConnection()) with pytest.raises(responses_websocket.ConnectionClosedOK): await adapter.recv() @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 @@ -92,10 +97,13 @@ 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() + assert _FakeNativeBridge.contexts[-1].capabilities.websocket_mode == "native" + assert _FakeNativeBridge.contexts[-1].capabilities.requires_connection is True class _FailingNativeBridge: @@ -104,16 +112,74 @@ class _FailingNativeBridge: cls, request: NativeResponsesWebSocketRequest, *, - options: object, + options: NativeRequestOptions, context: NativeRequestContext, ) -> _FakeNativeConnection: raise RuntimeError("connection failed") +@pytest.mark.asyncio +async def test_connection_failure_is_reported_to_orchestration() -> None: + configuration.rust(True) + responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge) + result = await responses_websocket.connect(url="wss://example.test/responses", headers={}, timeout=None) + assert isinstance(result, NativeFailed) + assert str(result.error) == "connection failed" + + +@pytest.mark.asyncio +async def test_managed_connection_closes_native_socket_on_consumer_failure() -> None: + configuration.rust(True) + socket = _FakeNativeConnection() + + class Bridge: + @classmethod + async def connect( + cls, + request: NativeResponsesWebSocketRequest, + *, + options: NativeRequestOptions, + context: NativeRequestContext, + ) -> _FakeNativeConnection: + return socket + + responses_websocket.set_rust_responses_websocket(connection=Bridge) + result = await responses_websocket.managed_connect(url="wss://example.test/responses", headers={}, timeout=1.0) + assert isinstance(result, Handled) + + async def use_connection() -> None: + async with result.value as connection: + await connection.send("hello") + raise ValueError("consumer failed") + + with pytest.raises(ValueError, match="consumer failed"): + await use_connection() + assert socket.sent == ["hello"] + assert socket.closed + + @pytest.mark.asyncio async def test_connection_failure_does_not_authorize_python_fallback() -> None: + from contextlib import AbstractAsyncContextManager + + from litellm.rust_bridge.dispatch import anative_context, provider_errors + configuration.rust(True) responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge) + @anative_context( + native=lambda: responses_websocket.managed_connect( + url="wss://example.test/responses", headers={}, timeout=None + ), + route="responses_websocket", + errors=lambda: provider_errors("openai", "responses websocket"), + ) + def execute() -> AbstractAsyncContextManager[object]: + pytest.fail("unknown native failures must not open a Python connection") + + async def run() -> None: + async with execute(): + pytest.fail("connection must fail before entering its body") + with pytest.raises(RuntimeError, match="connection failed"): - await responses_websocket.connect(url="wss://example.test/responses", headers={}, timeout=None) + await run() diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index f81b7050134..31163ebaa05 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -12,12 +12,8 @@ import pytest import litellm from litellm.rust_bridge import bindings, configuration from litellm.rust_bridge import chat_completions as bridge -from litellm.rust_bridge.request import ( - NativeBedrockOptions, - NativeRequestCapabilities, - NativeRequestContext, - anthropic_options, -) +from litellm.rust_bridge.request import NativeChatCompletionsRequest, NativeRequestContext, NativeRequestOptions +from litellm.rust_bridge.runtime import Handled, NativeFailed, NativeSkipped from litellm.types.utils import ModelResponse RUST_RESPONSE = { @@ -100,16 +96,40 @@ class _RecordingCall: self.result = result if result is not None else dict(RUST_RESPONSE) self.error = error self.calls: list[dict] = [] + self.contexts: list[NativeRequestContext] = [] - def __call__(self, request, *, options, context): - self.calls.append({"request": request, "options": options, "context": context}) + def __call__( + self, + request: NativeChatCompletionsRequest, + *, + options: NativeRequestOptions, + context: NativeRequestContext, + ): + kwargs = { + "model": request.model, + "messages": request.messages, + "optional_params": request.optional_params, + "api_key": options.api_key, + "api_base": options.api_base, + "custom_llm_provider": options.custom_llm_provider, + "extra_headers": options.extra_headers, + "timeout_seconds": options.timeout_seconds, + } + self.calls.append(kwargs) + self.contexts.append(context) if self.error is not None: raise self.error return self.result class _RecordingAsyncCall(_RecordingCall): - async def __call__(self, request, *, options, context): + async def __call__( + self, + request: NativeChatCompletionsRequest, + *, + options: NativeRequestOptions, + context: NativeRequestContext, + ): return _RecordingCall.__call__(self, request, options=options, context=context) @@ -256,7 +276,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" @@ -270,16 +291,28 @@ class TestSyncCall: native = _RecordingCall() bridge.set_rust_chat_completions(chat_completions=native) bridge.chat_completions(**_call_kwargs(ModelResponse())) - assert native.calls[0]["options"].timeout_seconds == 30.0 + assert native.calls[0]["timeout_seconds"] == 30.0 - def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch): + def test_preserves_execution_and_client_capabilities(self): + native = _RecordingCall() + bridge.set_rust_chat_completions(chat_completions=native) + bridge.chat_completions( + **_call_kwargs(ModelResponse()), + stream=True, + has_custom_client=True, + ) + assert native.contexts[0].capabilities.execution_mode == "sync" + assert native.contexts[0].capabilities.stream is True + assert native.contexts[0].capabilities.has_custom_client is True + + 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: @@ -287,187 +320,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 RateLimitError - - bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited"))) - with pytest.raises(RateLimitError) 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 InternalServerError - - 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(InternalServerError): - 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_does_not_authorize_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")), - ) - - with pytest.raises(RuntimeError, match="connection failed"): - bridge.chat_completions(**_call_kwargs(ModelResponse())) - - async def fallback(): - pytest.fail("unknown failure must not retry through Python") - - with pytest.raises(RuntimeError, match="connection failed"): - await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) - - -def test_provider_credentials_are_separate_from_chat_body_params(): - native = _RecordingCall() - bridge.set_rust_chat_completions(chat_completions=native) - configuration.rust(True) - kwargs = _call_kwargs(ModelResponse()) - kwargs["optional_params"] = { - "max_tokens": 32, - } - kwargs["bedrock"] = NativeBedrockOptions( - aws_access_key_id="test-access-key", - aws_secret_access_key="test-secret-key", - ) - bridge.chat_completions(**kwargs) - request = native.calls[0]["request"] - options = native.calls[0]["options"] - assert request.optional_params == {"max_tokens": 32} - assert options.bedrock.aws_access_key_id == "test-access-key" - assert options.bedrock.aws_secret_access_key == "test-secret-key" - - -def test_provider_payload_extensions_cross_the_boundary_without_partitioning(): - native = _RecordingCall() - bridge.set_rust_chat_completions(chat_completions=native) - configuration.rust(True) - extensions = { - "vendor_object": {"nested": None}, - "vendor_array": [1, "two", False], - "vendor_scalar": 0.25, - "extra_body": {"temperature": 0.2, "config": {"replacement": True}}, - } - - kwargs = _call_kwargs(ModelResponse()) - kwargs["optional_params"] = extensions - bridge.chat_completions(**kwargs) - - assert native.calls[0]["request"].optional_params == extensions - - -def test_typed_capability_and_provider_metadata_facts_are_isolated(): - context = NativeRequestContext( - capabilities=NativeRequestCapabilities( - stream=True, - has_agentic_hook=True, - has_custom_client=True, - request_format="native", - ) - ) - anthropic = anthropic_options({"metadata": {"user_id": "user-123", "ignored": object()}}) - - assert context.capabilities.request_format == "native" - assert context.capabilities.has_agentic_hook is True - assert anthropic.user_id == "user-123" + assert isinstance(await bridge.achat_completions(**_call_kwargs(ModelResponse())), NativeFailed) diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index 57c936c170b..6b58859155a 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -4,11 +4,8 @@ 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.request import NativeRequestContext, NativeRequestOptions, NativeTranscriptionRequest +from litellm.rust_bridge.runtime import Handled rust_bridge = importlib.import_module("litellm.rust_bridge.transcription") @@ -16,6 +13,7 @@ rust_bridge = importlib.import_module("litellm.rust_bridge.transcription") class SyncBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.contexts: list[NativeRequestContext] = [] def __call__( self, @@ -25,13 +23,9 @@ class SyncBridge: context: NativeRequestContext, ) -> dict[str, object]: self.calls.append( - { - "model": request.model, - "audio": request.audio, - "optional_params": request.optional_params, - "bedrock": options.bedrock, - } + {"model": request.model, "audio": request.audio, "optional_params": request.optional_params} ) + self.contexts.append(context) return {"text": "hello"} @@ -48,7 +42,7 @@ class AsyncBridge: def test_enabled_sync_bridge_receives_audio() -> None: bridge = SyncBridge() - rust_bridge.configure_rust_transcription(True, transcription=bridge) + rust_bridge.configure_rust_transcription(transcription=bridge) result = rust_bridge.transcription( model="mistral.voxtral-mini-3b-2507", audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, @@ -58,14 +52,22 @@ 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 result == {"text": "hello"} + assert isinstance(result, Handled) + assert result.value == {"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 async def test_enabled_async_bridge() -> None: - rust_bridge.configure_rust_transcription(True, atranscription=AsyncBridge()) + rust_bridge.configure_rust_transcription(atranscription=AsyncBridge()) result = await rust_bridge.atranscription( model="mistral.voxtral-mini-3b-2507", audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, @@ -76,7 +78,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: @@ -87,7 +89,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( @@ -104,10 +107,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( @@ -124,7 +125,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 request, *, options, context: {"text": "rust"}, + transcription=lambda *_args, **_: {"text": "rust"}, atranscription=None, ) try: @@ -140,9 +141,7 @@ 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( - request: NativeTranscriptionRequest, *, options: object, context: NativeRequestContext - ) -> dict[str, object]: + async def rust_response(*_args: object, **_: object) -> dict[str, object]: return {"text": "rust"} rust_bridge.configure_rust_transcription(transcription=None, atranscription=rust_response)