scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

fix(native): preserve typed capability context

Browse source
This commit is contained in:
Yujong Lee 2026-09-05 21:35:34 -07:00 • committed by yujonglee
parent 221ddb3ab6
commit fba6af26ad
19 changed files with 1595 additions and 1323 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm-rust/crates/core/src/request_context.rs
View file

@ -7,10 +7,15 @@ pub struct RequestAttribution {
#[derive(Clone, Debug, Default, PartialEq)]
pub struct RequestCapabilities {
pub execution_mode: Option<String>,
pub stream: bool,
pub has_agentic_hook: bool,
pub has_custom_client: bool,
pub request_format: Option<String>,
pub input_source_kind: Option<String>,
pub native_response_format: bool,
pub websocket_mode: Option<String>,
pub requires_connection: bool,
}
#[derive(Clone, Debug, Default, PartialEq)]

15
litellm-rust/crates/python-bridge/src/marshal.rs
View file

@ -118,10 +118,15 @@ pub(crate) struct NativeRequestContext {
#[derive(FromPyObject)]
struct NativeRequestCapabilities {
execution_mode: Option<String>,
stream: bool,
has_agentic_hook: bool,
has_custom_client: bool,
request_format: Option<String>,
input_source_kind: Option<String>,
native_response_format: bool,
websocket_mode: Option<String>,
requires_connection: bool,
}
impl From<NativeRequestContext> for litellm_core::request_context::LiteLlmRequestContext {
@ -136,10 +141,15 @@ impl From<NativeRequestContext> 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:

200
litellm/llms/anthropic/chat/handler.py
View file

@ -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

318
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -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()

7
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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:

282
litellm/ocr/main.py
View file

@ -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,

173
litellm/rust_bridge/chat_completions.py
View file

@ -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),
)

83
litellm/rust_bridge/messages.py
View file

@ -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),
)

308
litellm/rust_bridge/ocr.py
View file

@ -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),
)

5
litellm/rust_bridge/request.py
View file

@ -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)

73
litellm/rust_bridge/responses_websocket.py
View file

@ -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)

86
litellm/rust_bridge/transcription.py
View file

@ -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),
)

191
tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
View file

@ -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)

486
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
View file

@ -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)

124
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -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"

175
tests/test_litellm/ocr/test_rust_bridge.py
View file

@ -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,

96
tests/test_litellm/responses/test_rust_bridge_websocket.py
View file

@ -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()

244
tests/test_litellm/rust_bridge/test_chat_completions.py
View file

@ -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)

47
tests/test_litellm/test_audio_transcription_rust_bridge.py
View file

@ -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)