scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 134

Merge pull request #41479 from BerriAI/litellm_rust_bridge_declarative_route_catalog

Browse source 
refactor(rust_bridge): declarative route catalog and shared runtime selection
This commit is contained in:
yujonglee 2026-09-17 11:18:36 -07:00 • committed by GitHub
commit d5b8400aa9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
95 changed files with 4268 additions and 4076 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/test-unit.yml vendored
View file

@ -100,6 +100,7 @@ jobs:
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/chat_completions
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
@ -109,6 +110,7 @@ jobs:
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/messages
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag

51
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -157,11 +157,6 @@ mod tests {
let module = PyModule::new(py, "routes").expect("module should be created");
crate::routes::register(&module).expect("routes should register");
let routes = [
(
"ocr",
"aocr",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, input_sources=None, timeout_seconds=None)",
),
(
"transcription",
"atranscription",
@ -244,24 +239,22 @@ mod tests {
kwargs
.set_item("extra_headers", &invalid_headers)
.expect("kwargs should accept extra_headers");
let document = PyDict::new(py);
let audio = PyDict::new(py);
for (sync_name, async_name) in [("ocr", "aocr"), ("transcription", "atranscription")] {
let sync_error = module
.getattr(sync_name)
.and_then(|function| function.call(("model", &document), Some(&kwargs)))
.expect_err("sync route should reject non-dict extra_headers");
let async_error = module
.getattr(async_name)
.and_then(|function| function.call(("model", &document), Some(&kwargs)))
.expect_err("async route should reject non-dict extra_headers");
let sync_error = module
.getattr("transcription")
.and_then(|function| function.call(("model", &audio), Some(&kwargs)))
.expect_err("sync route should reject non-dict extra_headers");
let async_error = module
.getattr("atranscription")
.and_then(|function| function.call(("model", &audio), Some(&kwargs)))
.expect_err("async route should reject non-dict extra_headers");
assert_eq!(
sync_error.to_string(),
"ValueError: extra_headers must be a dict"
);
assert_eq!(async_error.to_string(), sync_error.to_string());
}
assert_eq!(
sync_error.to_string(),
"ValueError: extra_headers must be a dict"
);
assert_eq!(async_error.to_string(), sync_error.to_string());
});
}
@ -312,15 +305,13 @@ mod tests {
let invalid_payload =
PyModule::new(py, "invalid_payload").expect("invalid payload should be created");
for name in ["ocr", "transcription"] {
let error = module
.getattr(name)
.and_then(|function| {
function.call(("model", &invalid_payload), Some(&headers_kwargs))
})
.expect_err("payload should be validated before headers");
assert!(!error.to_string().contains("extra_headers"));
}
let error = module
.getattr("transcription")
.and_then(|function| {
function.call(("model", &invalid_payload), Some(&headers_kwargs))
})
.expect_err("payload should be validated before headers");
assert!(!error.to_string().contains("extra_headers"));
});
}

6
litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs
View file

@ -159,8 +159,8 @@ fn redact(
}
pub(super) fn response(py: Python<'_>, response: &LiteLLMOcrResponse) -> PyResult<Py<PyAny>> {
py.import("litellm.rust_bridge.ocr")?
.getattr("_response")?
py.import("litellm.rust_bridge.ocr.callbacks")?
.getattr("response")?
.call1((to_py(py, response)?,))
.map(Bound::unbind)
}
@ -172,7 +172,7 @@ pub(super) fn map_failure(
provider: &str,
) -> PyResult<Py<PyBaseException>> {
Ok(py
.import("litellm.rust_bridge.ocr_lifecycle")?
.import("litellm.rust_bridge.ocr.callbacks")?
.getattr("map_failure")?
.call1((error, request, provider))?
.extract()?)

28
litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs
View file

@ -297,8 +297,7 @@ impl litellm_core::ocr::hooks::OcrHooks for BridgeOcrHooks {
}
}
#[pyfunction]
fn _ocr_lifecycle(
fn run_ocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
@ -328,6 +327,27 @@ fn _ocr_lifecycle(
run_call(py, call, host)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(_ocr_lifecycle, module)?)
#[pyfunction]
fn ocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_ocr(py, request, args, kwargs, false)
}
#[pyfunction]
fn aocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_ocr(py, request, args, kwargs, true)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(ocr, module)?)?;
module.add_function(wrap_pyfunction!(aocr, module)?)
}

2
litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs
View file

@ -3,11 +3,9 @@ mod document;
mod errors;
mod lifecycle;
mod project;
mod value;
use pyo3::prelude::*;
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
value::register(module)?;
lifecycle::register(module)
}

80
litellm-rust/crates/python-bridge/src/routes/ocr/value.rs
View file

@ -1,80 +0,0 @@
use litellm_core::ocr::Error;
use std::future::Future;
use litellm_core::ocr::wire::{OcrWireRequest, decode_request};
use pyo3::prelude::*;
use serde_json::Value;
use super::errors::to_pyerr as ocr_error_to_pyerr;
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
fn prepare_ocr(
inputs: OcrInputs,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
let document = inputs.document;
let options = RouteOptions::from_python(RouteOptionsInputs {
model: inputs.model,
api_key: inputs.api_key,
api_base: inputs.api_base,
custom_llm_provider: inputs.custom_llm_provider,
extra_headers: inputs.extra_headers,
timeout_seconds: inputs.timeout_seconds,
})?;
let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
let input_sources = inputs
.input_sources
.map(serde_json::from_value)
.transpose()
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?
.unwrap_or_default();
Ok(async move {
let RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout,
} = options;
let request = decode_request(OcrWireRequest {
model,
document,
api_key,
api_base,
custom_llm_provider,
extra_headers,
optional_params,
input_sources,
timeout_seconds: timeout.map(|value| value.as_secs_f64()),
})?;
litellm_core::ocr::ocr(request)
.await
.map(|response| response.into_json())
})
}
bridge_route! {
sync = ocr,
asynchronous = aocr,
inputs = OcrInputs,
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
document: serde_json::Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
input_sources: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_ocr,
errors = ocr_error_to_pyerr,
}

21
litellm/__init__.py
View file

@ -1406,8 +1406,22 @@ from .images.main import *
from .videos.main import *
from .batch_completion.main import *
from .rerank_api.main import *
from .llms.anthropic.experimental_pass_through.messages.handler import *
from .responses.main import *
from .messages.dispatch import *
from .responses.dispatch import *
from .responses.main import (
acancel_responses,
acompact_responses,
adelete_responses,
aget_responses,
alist_input_items,
aresponses_api_with_mcp,
cancel_responses,
compact_responses,
delete_responses,
get_responses,
list_input_items,
mock_responses_api_response,
)
# Interactions API is available as litellm.interactions module
# Usage: litellm.interactions.create(), litellm.interactions.get(), etc.
@ -1435,7 +1449,8 @@ from .skills.main import (
adelete_skill,
)
from .containers.main import *
from .ocr.main import *
from .ocr.dispatch import *
from .chat_completions.dispatch import *
from .rust_bridge import rust
from .rag.main import *
from .sandbox.main import *

4
litellm/anthropic_interface/messages/__init__.py
View file

@ -13,10 +13,10 @@ This is an __init__.py file to allow the following interface
from collections.abc import AsyncIterator, Coroutine, Iterator
from typing import Any
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
from litellm.messages import (
anthropic_messages as _async_anthropic_messages,
)
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
from litellm.messages import (
anthropic_messages_handler as _sync_anthropic_messages,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (

3
litellm/chat_completions/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .dispatch import acompletion, completion
__all__ = ("acompletion", "completion")

126
litellm/chat_completions/dispatch.py Normal file
View file

@ -0,0 +1,126 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm import main
from litellm.rust_bridge.catalog import Context, Delivery, Route
from litellm.rust_bridge.chat_completions.entrypoints import (
NATIVE_ACOMPLETION,
NATIVE_COMPLETION,
LiteLLMChatCompletionsRequest,
)
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
from litellm.rust_bridge.public_call import (
bind,
optional_bool,
optional_mapping,
optional_sequence,
optional_str,
signature,
)
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
__all__ = ("acompletion", "completion")
ChatResult: TypeAlias = ModelResponse | CustomStreamWrapper
PythonCompletion: TypeAlias = Callable[..., ChatResult | Coroutine[object, object, ChatResult]]
PythonAcompletion: TypeAlias = Callable[..., Awaitable[ChatResult]]
def _python_completion() -> PythonCompletion:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonCompletion,
main.completion, # noqa: TID251 # dispatch boundary owns this Python fallback
)
def _python_acompletion() -> PythonAcompletion:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonAcompletion,
main.acompletion, # noqa: TID251 # dispatch boundary owns this Python fallback
)
_PYTHON_COMPLETION: Final = _python_completion()
_COMPLETION: Final = signature(_PYTHON_COMPLETION)
_PYTHON_ACOMPLETION: Final = _python_acompletion()
_ACOMPLETION: Final = signature(_PYTHON_ACOMPLETION)
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMChatCompletionsRequest | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
model: Final = fields.get("model")
messages: Final = optional_sequence(fields.get("messages"))
extra: Final = optional_mapping(fields.get("kwargs")) or MappingProxyType({})
if not isinstance(model, str) or messages is None:
return None
return LiteLLMChatCompletionsRequest(
model=model,
messages=messages,
stream=optional_bool(fields.get("stream")),
api_key=optional_str(fields.get("api_key")),
api_base=optional_str(extra.get("api_base")) or optional_str(fields.get("base_url")),
custom_llm_provider=optional_str(extra.get("custom_llm_provider")),
extra_headers=optional_mapping(fields.get("extra_headers")),
kwargs=extra,
)
def _context(request: LiteLLMChatCompletionsRequest) -> Context:
return Context(
Route.CHAT_COMPLETIONS,
provider=request.custom_llm_provider,
model=request.model,
delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
)
_DISPATCH: Final = PublicDispatch(
route=Route.CHAT_COMPLETIONS,
request=lambda args, kwargs: _public_request(_COMPLETION, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("acompletion") is True,
)
_ADISPATCH: Final = PublicDispatch(
route=Route.CHAT_COMPLETIONS,
request=lambda args, kwargs: _public_request(_ACOMPLETION, args, kwargs),
context=_context,
)
def completion(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public chat completions call shape
) -> ChatResult | Coroutine[object, object, ChatResult]:
python: Final = _PYTHON_COMPLETION
return _DISPATCH.run(
args,
kwargs,
python=python,
binding=NATIVE_COMPLETION,
native=call_hook,
)
async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-ok: preserve the public call shape
python: Final = _PYTHON_ACOMPLETION
return await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=NATIVE_ACOMPLETION,
native=call_hook,
)
completion.__doc__ = _PYTHON_COMPLETION.__doc__
completion.__wrapped__ = _PYTHON_COMPLETION # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature
acompletion.__doc__ = _PYTHON_ACOMPLETION.__doc__
acompletion.__wrapped__ = _PYTHON_ACOMPLETION # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature

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

@ -25,8 +25,6 @@ from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
)
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.types.llms.anthropic import (
ContentBlockDelta,
ContentBlockStart,
@ -375,24 +373,22 @@ class AnthropicChatCompletion(BaseLLM):
"""Filter beta headers and emit pre_call, returning `(headers, data)`.
The pair stays mutable because the streaming path rewrites it in
place (`data["stream"] = True`) before sending. A Rust attempt that
declined already emitted pre_call for this request, so skip it there.
place (`data["stream"] = True`) before sending.
"""
request_headers, data = update_request_with_filtered_beta(
headers=headers,
request_data=request_data,
provider=custom_llm_provider,
)
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": request_headers,
},
)
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": request_headers,
},
)
print_verbose(f"_is_function_call: {_is_function_call}")
return request_headers, data
@ -456,68 +452,6 @@ class AnthropicChatCompletion(BaseLLM):
timeout=timeout,
)
# 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
# `max_tokens` among them) that `transform_request` would have applied.
rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy
**AnthropicConfig.get_config(model=model),
**optional_params,
}
serves_via_rust: Final = rust_chat_completions_accepts(
model=model,
messages=messages,
optional_params=rust_optional_params,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
stream=stream,
)
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:
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=acompletion_dispatch,
)
rust_response: Final = rust_chat_completions_bridge.chat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
)
if rust_response is not None:
return rust_response
if acompletion is True:
return acompletion_dispatch()
else:

2
litellm/llms/anthropic/experimental_pass_through/messages/handler.py
View file

@ -40,6 +40,8 @@ from ..utils import is_reasoning_auto_summary_enabled
from .interceptors import get_messages_interceptors
from .utils import AnthropicMessagesRequestUtils, mock_response
__all__ = ("anthropic_messages", "anthropic_messages_handler")
# Providers that are routed directly to the OpenAI Responses API instead of
# going through chat/completions.
_RESPONSES_API_PROVIDERS: Final = frozenset({"openai"})

4
litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py
View file

@ -414,9 +414,7 @@ async def _call_messages_handler(
Using the public function (decorated with @client) ensures logging, retries,
and provider resolution all work correctly, identical to a direct user call.
"""
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
anthropic_messages,
)
from litellm.messages import anthropic_messages
return await anthropic_messages(
model=model,

82
litellm/llms/bedrock/audio_transcription/__init__.py
View file

@ -1,13 +1,29 @@
import base64
from typing import Final
from typing import Final, NoReturn
import httpx
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.rust_bridge import transcription as rust_transcription_bridge
from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Context, Route
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.rust_bridge.transcription.native import (
NATIVE_ATRANSCRIPTION,
NATIVE_TRANSCRIPTION,
RustAtranscription,
RustTranscription,
)
from litellm.types.utils import FileTypes, TranscriptionResponse
def _no_python_implementation() -> NoReturn:
raise NotImplementedError("Bedrock audio transcription is implemented in Rust only")
async def _no_async_python_implementation() -> NoReturn:
_no_python_implementation()
class BedrockAudioTranscriptionRustDispatch:
@staticmethod
def _audio_payload(audio_file: FileTypes) -> dict[str, object]:
@ -43,19 +59,26 @@ class BedrockAudioTranscriptionRustDispatch:
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> TranscriptionResponse:
rust_response: Final = rust_transcription_bridge.transcription(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout=timeout,
def native(rust: RustTranscription) -> TranscriptionResponse:
return TranscriptionResponse(
**rust(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
)
)
return runtime.run(
Context(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model),
binding=NATIVE_TRANSCRIPTION,
native=native,
python=_no_python_implementation,
)
if rust_response is None:
raise RuntimeError("Rust audio transcription bridge is unavailable")
return TranscriptionResponse(**rust_response)
async def async_audio_transcriptions(
self,
@ -69,16 +92,23 @@ class BedrockAudioTranscriptionRustDispatch:
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> TranscriptionResponse:
rust_response: Final = await rust_transcription_bridge.atranscription(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout=timeout,
async def native(rust: RustAtranscription) -> TranscriptionResponse:
return TranscriptionResponse(
**await rust(
model=model,
audio=self._audio_payload(audio_file),
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
)
)
return await runtime.arun(
Context(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model),
binding=NATIVE_ATRANSCRIPTION,
native=native,
python=_no_async_python_implementation,
)
if rust_response is None:
raise RuntimeError("Rust audio transcription bridge is unavailable")
return TranscriptionResponse(**rust_response)

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

@ -1,6 +1,4 @@
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final
import httpx
@ -16,8 +14,6 @@ from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
)
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.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
@ -26,22 +22,6 @@ from ..common_utils import BedrockError, _get_all_bedrock_regions, error_respons
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
def _sigv4_principal(credentials: Credentials | None) -> Mapping[str, str]:
if credentials is None:
return MappingProxyType({})
return MappingProxyType(
{
key: value
for key, value in (
("aws_access_key_id", credentials.access_key),
("aws_secret_access_key", credentials.secret_key),
("aws_session_token", credentials.token),
)
if value is not None
}
)
def make_sync_call(
client: HTTPHandler | None,
api_base: str,
@ -401,87 +381,6 @@ class BedrockConverseLLM(BaseAWSLLM):
# Filter beta headers in HTTP headers before making the request
headers = update_headers_with_filtered_beta(headers=headers, provider="bedrock_converse")
# The Rust core owns the whole call for the subset it accepts. Ask
# before transforming so whichever path runs emits pre_call once, and
# hand down the credentials, region and endpoint this handler already
# 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 = { # 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,
optional_params=rust_optional_params,
custom_llm_provider="bedrock",
litellm_params=litellm_params,
stream=stream,
)
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,
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(
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,
)
if rust_response is not None:
return rust_response
### ROUTING (ASYNC, STREAMING, SYNC)
if acompletion:
if isinstance(client, HTTPHandler):
@ -548,21 +447,15 @@ class BedrockConverseLLM(BaseAWSLLM):
)
## 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,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": proxy_endpoint_url,
"headers": prepped.headers,
},
)
logging_obj.pre_call(
input=messages,
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:

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

@ -166,9 +166,11 @@ from litellm.utils import (
def _rust_responses_websocket_enabled(
custom_llm_provider: str | None,
) -> bool:
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.catalog import Context, Delivery, Route, decision
from litellm.rust_bridge.configuration import Decision
return custom_llm_provider == "openai" and rust_enabled()
context: Final = Context(Route.RESPONSES, provider=custom_llm_provider, delivery=Delivery.WEBSOCKET)
return decision(context) is not Decision.PYTHON
from .http_handler import get_shared_realtime_ssl_context
@ -183,9 +185,6 @@ if TYPE_CHECKING:
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
AnthropicMessagesStreamingResponse,
)
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from litellm.types.llms.openai_evals import (
CancelEvalResponse,
@ -2283,36 +2282,6 @@ class BaseLLMHTTPHandler:
},
)
rust_messages_response: Final = await self._maybe_rust_anthropic_messages(
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
has_agentic_hook=self._has_agentic_completion_hook(logging_obj),
model=model,
api_key=api_key,
api_base=api_base,
headers=headers,
request_body=request_body,
timeout=self._resolve_anthropic_messages_timeout(
litellm_params=litellm_params,
stream=stream or False,
custom_llm_provider=custom_llm_provider,
),
)
if rust_messages_response is not None:
if stream:
return self._rust_anthropic_messages_fake_stream(rust_messages_response)
return await self._finalize_anthropic_messages_response(
initial_response=rust_messages_response,
model=model,
messages=messages,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
kwargs=kwargs,
)
response: Final = await self._async_post_anthropic_messages_with_http_error_retry(
async_httpx_client=async_httpx_client,
request_url=request_url,
@ -2441,73 +2410,6 @@ class BaseLLMHTTPHandler:
"anthropic_messages",
)
@staticmethod
async def _maybe_rust_anthropic_messages(
*,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
has_agentic_hook: bool,
model: str,
api_key: str | None,
api_base: str | None,
headers: dict,
request_body: dict,
timeout: float | httpx.Timeout | None,
) -> AnthropicMessagesResponse | None:
if custom_llm_provider not in ("azure_ai", "anthropic"):
return None
from litellm.rust_bridge.configuration import rust_enabled
if not rust_enabled():
return None
if has_agentic_hook:
return None
from litellm.rust_bridge import messages as rust_messages_bridge
upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"}
try:
rust_response: Final = await rust_messages_bridge.amessages(
model=model,
body=upstream_body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
)
except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path
verbose_logger.debug(
"Rust Anthropic messages bridge raised %s; falling back to Python path",
type(rust_error).__name__,
)
return None
if rust_response is None:
return None
response_obj: Final = cast(AnthropicMessagesResponse, dict(rust_response))
response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}}
return response_obj
@staticmethod
def _rust_anthropic_messages_fake_stream(
rust_response: AnthropicMessagesResponse,
) -> "AnthropicMessagesStreamingResponse":
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
AnthropicMessagesStreamHiddenParams,
AnthropicMessagesStreamingResponse,
)
completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response))
hidden_params: Final = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"})
return AnthropicMessagesStreamingResponse(
completion_stream=completion_stream,
hidden_params=hidden_params,
)
def anthropic_messages_handler(
self,
model: str,
@ -6658,7 +6560,7 @@ class BaseLLMHTTPHandler:
@asynccontextmanager
async def _backend_connection():
if _rust_responses_websocket_enabled(custom_llm_provider):
from litellm.rust_bridge import responses_websocket as rust_responses_websocket
from litellm.rust_bridge.responses import websocket as rust_responses_websocket
rust_backend: Final = await rust_responses_websocket.connect(
url=ws_url,

4
litellm/main.py
View file

@ -5968,7 +5968,7 @@ def responses_with_retries(*args, **kwargs):
except Exception as e:
raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}")
from litellm.responses.main import responses
from litellm.responses.dispatch import responses
num_retries: Final = kwargs.pop("num_retries", 3)
# reset retries in .responses()
@ -5998,7 +5998,7 @@ async def aresponses_with_retries(*args, **kwargs):
except Exception as e:
raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}")
from litellm.responses.main import aresponses
from litellm.responses.dispatch import aresponses
num_retries: Final = kwargs.pop("num_retries", 3)
kwargs["max_retries"] = 0

3
litellm/messages/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .dispatch import anthropic_messages, anthropic_messages_handler
__all__ = ("anthropic_messages", "anthropic_messages_handler")

125
litellm/messages/dispatch.py Normal file
View file

@ -0,0 +1,125 @@
import inspect
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm.llms.anthropic.experimental_pass_through.messages import handler as main
from litellm.rust_bridge.catalog import Context, Delivery, Route
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
from litellm.rust_bridge.messages.entrypoints import (
NATIVE_AMESSAGES,
NATIVE_MESSAGES,
LiteLLMMessagesRequest,
)
from litellm.rust_bridge.public_call import (
bind,
optional_bool,
optional_mapping,
optional_sequence,
optional_str,
signature,
)
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
__all__ = ("anthropic_messages", "anthropic_messages_handler")
MessagesResult: TypeAlias = AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]
PythonMessages: TypeAlias = Callable[..., MessagesResult | Coroutine[object, object, MessagesResult]]
PythonAmessages: TypeAlias = Callable[..., Awaitable[MessagesResult]]
def _python_messages() -> PythonMessages:
return cast( # cast-ok: forward the original call shape through the legacy handler
PythonMessages,
main.anthropic_messages_handler, # noqa: TID251 # dispatch boundary owns this Python fallback
)
def _python_amessages() -> PythonAmessages:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonAmessages,
main.anthropic_messages, # noqa: TID251 # dispatch boundary owns this Python fallback
)
_PYTHON_MESSAGES: Final = _python_messages()
_MESSAGES: Final = signature(_PYTHON_MESSAGES)
_PYTHON_AMESSAGES: Final = _python_amessages()
_AMESSAGES: Final = signature(_PYTHON_AMESSAGES)
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMMessagesRequest | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
model: Final = fields.get("model")
messages: Final = optional_sequence(fields.get("messages"))
max_tokens: Final = fields.get("max_tokens")
if not isinstance(model, str) or messages is None or not isinstance(max_tokens, int):
return None
return LiteLLMMessagesRequest(
model=model,
messages=messages,
max_tokens=max_tokens,
stream=optional_bool(fields.get("stream")),
api_key=optional_str(fields.get("api_key")),
api_base=optional_str(fields.get("api_base")),
custom_llm_provider=optional_str(fields.get("custom_llm_provider")),
kwargs=optional_mapping(fields.get("kwargs")) or MappingProxyType({}),
)
def _context(request: LiteLLMMessagesRequest) -> Context:
return Context(
Route.MESSAGES,
provider=request.custom_llm_provider,
model=request.model,
delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
)
_DISPATCH: Final = PublicDispatch(
route=Route.MESSAGES,
request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("is_async") is True,
)
_ADISPATCH: Final = PublicDispatch(
route=Route.MESSAGES,
request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs),
context=_context,
)
def anthropic_messages_handler(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public Anthropic Messages call shape
) -> MessagesResult | Coroutine[object, object, MessagesResult]:
python: Final = _PYTHON_MESSAGES
return _DISPATCH.run(
args,
kwargs,
python=python,
binding=NATIVE_MESSAGES,
native=call_hook,
)
async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult: # kwargs-ok: public call shape
python: Final = _PYTHON_AMESSAGES
return await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=NATIVE_AMESSAGES,
native=call_hook,
)
anthropic_messages_handler.__doc__ = _PYTHON_MESSAGES.__doc__
anthropic_messages_handler.__wrapped__ = _PYTHON_MESSAGES # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature
anthropic_messages.__doc__ = _PYTHON_AMESSAGES.__doc__
anthropic_messages.__wrapped__ = _PYTHON_AMESSAGES # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature

2
litellm/ocr/__init__.py
View file

@ -1,5 +1,5 @@
"""OCR module for LiteLLM."""
from .main import aocr, ocr
from .dispatch import aocr, ocr
__all__ = ["aocr", "ocr"]

93
litellm/ocr/dispatch.py Normal file
View file

@ -0,0 +1,93 @@
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from typing import Final, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
import httpx
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import main
from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type
from litellm.rust_bridge.catalog import Context, Route
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR, LiteLLMOcrRequest
__all__ = ("aocr", "convert_file_document_to_url_document", "get_mime_type", "ocr")
def _bind_request(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> LiteLLMOcrRequest:
return LiteLLMOcrRequest(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
def _public_request(name: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> LiteLLMOcrRequest:
try:
return _bind_request(*args, **kwargs) # pyright: ignore[reportArgumentType] # Python binds the public arguments before native validation
except TypeError as error:
raise TypeError(str(error).replace("_bind_request()", f"{name}()")) from None
_PYTHON_OCR: Final = cast( # cast-ok: forward the original call shape through the Python @client decorator
Callable[..., OCRResponse | Coroutine[object, object, OCRResponse]],
main.ocr, # noqa: TID251 # dispatch boundary owns this Python fallback
)
_PYTHON_AOCR: Final = cast( # cast-ok: forward the original call shape through the Python @client decorator
Callable[..., Awaitable[OCRResponse]],
main.aocr, # noqa: TID251 # dispatch boundary owns this Python fallback
)
def _context(request: LiteLLMOcrRequest) -> Context:
return Context(Route.OCR, provider=request.custom_llm_provider, model=request.model)
_DISPATCH: Final = PublicDispatch(
route=Route.OCR,
request=lambda args, kwargs: _public_request("ocr", args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("aocr") is True,
)
_ADISPATCH: Final = PublicDispatch(
route=Route.OCR,
request=lambda args, kwargs: _public_request("aocr", args, kwargs),
context=_context,
)
def ocr(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
return _DISPATCH.run(
args,
kwargs,
python=_PYTHON_OCR,
binding=NATIVE_OCR,
native=call_hook,
)
async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape
return await _ADISPATCH.arun(
args,
kwargs,
python=_PYTHON_AOCR,
binding=NATIVE_AOCR,
native=call_hook,
)

416
litellm/ocr/legacy.py
View file

@ -1,416 +0,0 @@
"""
Main OCR function for LiteLLM.
"""
import asyncio
import base64
import mimetypes
import os
import re
from collections.abc import Coroutine, Mapping
from dataclasses import dataclass
from io import IOBase
from types import MappingProxyType
from typing import Final, Protocol, cast # noqa: TID251 # adapters preserve the legacy untyped contracts
import httpx
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.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.types.router import GenericLiteLLMParams
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.utils import ProviderConfigManager, client
base_llm_http_handler: Final = BaseLLMHTTPHandler()
class FileReader(Protocol):
def read(self) -> bytes | str: ...
@dataclass(frozen=True, slots=True)
class _PreparedOCRRequest:
model: str
document: Mapping[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
def _prepare_ocr_request(
model: str,
document: Mapping[str, object],
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
kwargs: dict[str, object],
) -> _PreparedOCRRequest:
litellm_logging_obj: Final = cast( # cast-ok: @client supplies the logging object; preserve legacy failure behavior
LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")
)
litellm_call_id: Final = cast( # cast-ok: @client supplies the call id without coercion
str | None, kwargs.get("litellm_call_id", None)
)
if not isinstance(document, dict):
raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}")
doc_type = document.get("type")
if doc_type == "file":
document = convert_file_document_to_url_document(document)
doc_type = document.get("type")
if doc_type not in ["document_url", "image_url"]:
raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'")
(
model,
custom_llm_provider,
dynamic_api_key,
dynamic_api_base,
) = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
ocr_provider_config: Final = ProviderConfigManager.get_provider_ocr_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
if ocr_provider_config is None:
raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}")
resolved_api_key, resolved_api_base = ocr_provider_config.resolve_connection_params(
api_key=api_key,
api_base=api_base,
dynamic_api_key=dynamic_api_key,
dynamic_api_base=dynamic_api_base,
)
verbose_logger.debug("OCR call - model: %s, provider: %s", model, custom_llm_provider)
litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs)
supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model)
requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM)
if requested_format is not None:
try:
parsed_format: Final = parse_ocr_request_format(requested_format)
except ValueError as e:
raise litellm.exceptions.UnsupportedParamsError(
message=f"{e}", model=model, llm_provider=custom_llm_provider
) from e
if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native":
raise litellm.exceptions.UnsupportedParamsError(
message=(
f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, "
f"model: {model}"
),
model=model,
llm_provider=custom_llm_provider,
)
non_default_params: Final = {}
for param in supported_params:
if param in kwargs:
non_default_params[param] = kwargs.pop(param)
optional_params: Final = ocr_provider_config.map_ocr_params(
non_default_params=non_default_params,
optional_params={},
model=model,
)
verbose_logger.debug("OCR optional_params after mapping: %s", optional_params)
effective_timeout: Final = timeout or request_timeout
litellm_logging_obj.update_from_kwargs(
kwargs=kwargs,
model=model,
optional_params=optional_params,
litellm_params={
"litellm_call_id": litellm_call_id,
"api_base": resolved_api_base,
**litellm_params.model_dump(include=frozenset(CustomPricingLiteLLMParams.model_fields), exclude_none=True),
},
custom_llm_provider=custom_llm_provider,
)
return _PreparedOCRRequest(
model=model,
document=document,
api_key=resolved_api_key,
api_base=resolved_api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
provider_config=ocr_provider_config,
optional_params=cast(
dict[str, object], optional_params
), # cast-ok: provider configs return heterogeneous OCR options
litellm_params=dict(litellm_params),
effective_timeout=effective_timeout,
litellm_logging_obj=litellm_logging_obj,
)
def _error_provider(model: str, custom_llm_provider: str | None) -> str | None:
if custom_llm_provider is not None:
return custom_llm_provider
prefix: Final = model.partition("/")[0]
if prefix in {"mistral", "azure_ai", "vertex_ai"}:
return prefix
return "mistral" if model.startswith("mistral-ocr") else None
@client
async def aocr(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> OCRResponse:
completion_kwargs: Final[dict[str, object]] = {
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"timeout": timeout,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"kwargs": kwargs,
}
try:
prepared: Final = _prepare_ocr_request(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
model = prepared.model
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
response = base_llm_http_handler.ocr(
model=prepared.model,
document=cast( # cast-ok: preserve legacy document fields for provider validation
dict[str, str], 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,
)
if asyncio.iscoroutine(response):
response = await response
if response is None:
raise ValueError(f"Got an unexpected None response from the OCR API: {response}")
return response
except Exception as e:
error_provider: Final = _error_provider(model, custom_llm_provider)
error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model
raise litellm.exception_type(
model=error_model,
custom_llm_provider=error_provider,
original_exception=e,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)
_MIME_PATTERN: Final = re.compile(r"^[\w.+-]+/[\w.+-]+$")
_MIME_TYPE_MAP: Final = MappingProxyType(
{
".pdf": "application/pdf",
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
".tiff": "image/tiff",
".tif": "image/tiff",
".bmp": "image/bmp",
}
)
def get_mime_type(file_path: str) -> str:
ext: Final = os.path.splitext(file_path)[1].lower()
mime: Final = _MIME_TYPE_MAP.get(ext)
if mime:
return mime
guessed, _ = mimetypes.guess_type(file_path)
return guessed or "application/octet-stream"
def _read_file(file_input: object) -> tuple[bytes, str, str | None]:
if isinstance(file_input, str):
raise ValueError(
"OCR file input does not accept bare str values. Pass bytes, "
"a pathlib.Path, or a file-like object. To OCR a local file "
"from a path, call open(path, 'rb') yourself."
)
if isinstance(file_input, os.PathLike):
file_path: Final = str(cast(object, file_input)) # cast-ok: preserve staging's str(PathLike) conversion
if not os.path.isfile(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
mime_type: Final = get_mime_type(file_path)
with open(file_path, "rb") as stream:
return stream.read(), mime_type, os.path.basename(file_path)
if isinstance(file_input, bytes):
return file_input, "application/octet-stream", None
if isinstance(file_input, IOBase) or hasattr(file_input, "read"):
file_name: Final = cast( # cast-ok: retain legacy validation and errors for file-like metadata
str | None, getattr(file_input, "name", None)
)
inferred_mime: Final = get_mime_type(file_name) if file_name else "application/octet-stream"
reader: Final = cast(FileReader, file_input) # cast-ok: legacy accepts duck-typed file readers
content: Final = reader.read()
return content.encode("utf-8") if isinstance(content, str) else content, inferred_mime, file_name
raise ValueError(
f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object."
)
def convert_file_document_to_url_document(document: Mapping[str, object]) -> dict[str, str]:
file_input: Final = document.get("file")
if file_input is None:
raise ValueError(
"document with type='file' must include a 'file' field containing "
"a pathlib.Path, file-like object, or bytes"
)
file_bytes, inferred_mime, file_name = _read_file(file_input)
if not file_bytes:
raise ValueError("File is empty or could not be read")
mime_type: Final = cast( # cast-ok: keep staging's MIME validation errors
str, document.get("mime_type", inferred_mime)
)
if not _MIME_PATTERN.match(mime_type):
raise ValueError(f"Invalid MIME type: {mime_type}")
base64_data: Final = base64.b64encode(file_bytes).decode("utf-8")
data_uri: Final = f"data:{mime_type};base64,{base64_data}"
if mime_type.startswith("image/"):
verbose_logger.debug(
"OCR file input: Converted file to image_url data URI (mime=%s, size=%s bytes, name=%s)",
mime_type,
len(file_bytes),
file_name,
)
return {"type": "image_url", "image_url": data_uri}
verbose_logger.debug(
"OCR file input: Converted file to document_url data URI (mime=%s, size=%s bytes, name=%s)",
mime_type,
len(file_bytes),
file_name,
)
return {"type": "document_url", "document_url": data_uri}
@client
def ocr(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
completion_kwargs: Final[dict[str, object]] = {
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"timeout": timeout,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"kwargs": kwargs,
}
try:
_is_async: Final = kwargs.pop("aocr", False) is True
completion_kwargs["aocr"] = _is_async
prepared: Final = _prepare_ocr_request(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
kwargs=kwargs,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout=timeout,
)
model = prepared.model
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
response: Final = base_llm_http_handler.ocr(
model=prepared.model,
document=cast( # cast-ok: preserve legacy document fields for provider validation
dict[str, str], 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 response
except Exception as e:
error_provider: Final = _error_provider(model, custom_llm_provider)
error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model
raise litellm.exception_type(
model=error_model,
custom_llm_provider=error_provider,
original_exception=e,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)

457
litellm/ocr/main.py
View file

@ -1,20 +1,191 @@
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from typing import Final, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
"""
Main OCR function for LiteLLM.
"""
import asyncio
import base64
import mimetypes
import os
import re
from collections.abc import Coroutine, Mapping
from dataclasses import dataclass
from io import IOBase
from types import MappingProxyType
from typing import Final, Protocol, cast # noqa: TID251 # adapters preserve the legacy untyped contracts
import httpx
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import legacy
from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
from litellm.rust_bridge.bindings import native_exception_types
from litellm.rust_bridge.configuration import rust_ocr_enabled
from litellm.rust_bridge.ocr import LiteLLMOcrRequest
from litellm.rust_bridge.ocr_lifecycle import select
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.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.types.router import GenericLiteLLMParams
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.utils import ProviderConfigManager, client
__all__ = ("aocr", "convert_file_document_to_url_document", "get_mime_type", "ocr")
base_llm_http_handler: Final = BaseLLMHTTPHandler()
def _bind_request(
class FileReader(Protocol):
def read(self) -> bytes | str: ...
@dataclass(frozen=True, slots=True)
class _PreparedOCRRequest:
model: str
document: Mapping[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
def _prepare_ocr_request(
model: str,
document: Mapping[str, object],
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
kwargs: dict[str, object],
) -> _PreparedOCRRequest:
litellm_logging_obj: Final = cast( # cast-ok: @client supplies the logging object; preserve legacy failure behavior
LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")
)
litellm_call_id: Final = cast( # cast-ok: @client supplies the call id without coercion
str | None, kwargs.get("litellm_call_id", None)
)
if not isinstance(document, dict):
raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}")
doc_type = document.get("type")
if doc_type == "file":
document = convert_file_document_to_url_document(document)
doc_type = document.get("type")
if doc_type not in ["document_url", "image_url"]:
raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'")
(
model,
custom_llm_provider,
dynamic_api_key,
dynamic_api_base,
) = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
ocr_provider_config: Final = ProviderConfigManager.get_provider_ocr_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
if ocr_provider_config is None:
raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}")
resolved_api_key, resolved_api_base = ocr_provider_config.resolve_connection_params(
api_key=api_key,
api_base=api_base,
dynamic_api_key=dynamic_api_key,
dynamic_api_base=dynamic_api_base,
)
verbose_logger.debug("OCR call - model: %s, provider: %s", model, custom_llm_provider)
litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs)
supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model)
requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM)
if requested_format is not None:
try:
parsed_format: Final = parse_ocr_request_format(requested_format)
except ValueError as e:
raise litellm.exceptions.UnsupportedParamsError(
message=f"{e}", model=model, llm_provider=custom_llm_provider
) from e
if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native":
raise litellm.exceptions.UnsupportedParamsError(
message=(
f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, "
f"model: {model}"
),
model=model,
llm_provider=custom_llm_provider,
)
non_default_params: Final = {}
for param in supported_params:
if param in kwargs:
non_default_params[param] = kwargs.pop(param)
optional_params: Final = ocr_provider_config.map_ocr_params(
non_default_params=non_default_params,
optional_params={},
model=model,
)
verbose_logger.debug("OCR optional_params after mapping: %s", optional_params)
effective_timeout: Final = timeout or request_timeout
litellm_logging_obj.update_from_kwargs(
kwargs=kwargs,
model=model,
optional_params=optional_params,
litellm_params={
"litellm_call_id": litellm_call_id,
"api_base": resolved_api_base,
**litellm_params.model_dump(include=frozenset(CustomPricingLiteLLMParams.model_fields), exclude_none=True),
},
custom_llm_provider=custom_llm_provider,
)
return _PreparedOCRRequest(
model=model,
document=document,
api_key=resolved_api_key,
api_base=resolved_api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
provider_config=ocr_provider_config,
optional_params=cast(
dict[str, object], optional_params
), # cast-ok: provider configs return heterogeneous OCR options
litellm_params=dict(litellm_params),
effective_timeout=effective_timeout,
litellm_logging_obj=litellm_logging_obj,
)
def _error_provider(model: str, custom_llm_provider: str | None) -> str | None:
if custom_llm_provider is not None:
return custom_llm_provider
prefix: Final = model.partition("/")[0]
if prefix in {"mistral", "azure_ai", "vertex_ai"}:
return prefix
return "mistral" if model.startswith("mistral-ocr") else None
@client
async def aocr(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
@ -23,61 +194,223 @@ def _bind_request(
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> LiteLLMOcrRequest:
return LiteLLMOcrRequest(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
def _public_request(name: str, args: tuple[object, ...], kwargs: dict[str, object]) -> LiteLLMOcrRequest:
) -> OCRResponse:
completion_kwargs: Final[dict[str, object]] = {
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"timeout": timeout,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"kwargs": kwargs,
}
try:
return _bind_request(*args, **kwargs) # pyright: ignore[reportArgumentType] # Python binds the public arguments before native validation
except TypeError as error:
raise TypeError(str(error).replace("_bind_request()", f"{name}()")) from None
prepared: Final = _prepare_ocr_request(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
model = prepared.model
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
response = base_llm_http_handler.ocr(
model=prepared.model,
document=cast( # cast-ok: preserve legacy document fields for provider validation
dict[str, str], 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,
)
if asyncio.iscoroutine(response):
response = await response
if response is None:
raise ValueError(f"Got an unexpected None response from the OCR API: {response}")
return response
except Exception as e:
error_provider: Final = _error_provider(model, custom_llm_provider)
error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model
raise litellm.exception_type(
model=error_model,
custom_llm_provider=error_provider,
original_exception=e,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)
_MIME_PATTERN: Final = re.compile(r"^[\w.+-]+/[\w.+-]+$")
_MIME_TYPE_MAP: Final = MappingProxyType(
{
".pdf": "application/pdf",
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
".tiff": "image/tiff",
".tif": "image/tiff",
".bmp": "image/bmp",
}
)
def get_mime_type(file_path: str) -> str:
ext: Final = os.path.splitext(file_path)[1].lower()
mime: Final = _MIME_TYPE_MAP.get(ext)
if mime:
return mime
guessed, _ = mimetypes.guess_type(file_path)
return guessed or "application/octet-stream"
def _read_file(file_input: object) -> tuple[bytes, str, str | None]:
if isinstance(file_input, str):
raise ValueError(
"OCR file input does not accept bare str values. Pass bytes, "
"a pathlib.Path, or a file-like object. To OCR a local file "
"from a path, call open(path, 'rb') yourself."
)
if isinstance(file_input, os.PathLike):
file_path: Final = str(cast(object, file_input)) # cast-ok: preserve staging's str(PathLike) conversion
if not os.path.isfile(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
mime_type: Final = get_mime_type(file_path)
with open(file_path, "rb") as stream:
return stream.read(), mime_type, os.path.basename(file_path)
if isinstance(file_input, bytes):
return file_input, "application/octet-stream", None
if isinstance(file_input, IOBase) or hasattr(file_input, "read"):
file_name: Final = cast( # cast-ok: retain legacy validation and errors for file-like metadata
str | None, getattr(file_input, "name", None)
)
inferred_mime: Final = get_mime_type(file_name) if file_name else "application/octet-stream"
reader: Final = cast(FileReader, file_input) # cast-ok: legacy accepts duck-typed file readers
content: Final = reader.read()
return content.encode("utf-8") if isinstance(content, str) else content, inferred_mime, file_name
raise ValueError(
f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object."
)
def convert_file_document_to_url_document(document: Mapping[str, object]) -> dict[str, str]:
file_input: Final = document.get("file")
if file_input is None:
raise ValueError(
"document with type='file' must include a 'file' field containing "
"a pathlib.Path, file-like object, or bytes"
)
file_bytes, inferred_mime, file_name = _read_file(file_input)
if not file_bytes:
raise ValueError("File is empty or could not be read")
mime_type: Final = cast( # cast-ok: keep staging's MIME validation errors
str, document.get("mime_type", inferred_mime)
)
if not _MIME_PATTERN.match(mime_type):
raise ValueError(f"Invalid MIME type: {mime_type}")
base64_data: Final = base64.b64encode(file_bytes).decode("utf-8")
data_uri: Final = f"data:{mime_type};base64,{base64_data}"
if mime_type.startswith("image/"):
verbose_logger.debug(
"OCR file input: Converted file to image_url data URI (mime=%s, size=%s bytes, name=%s)",
mime_type,
len(file_bytes),
file_name,
)
return {"type": "image_url", "image_url": data_uri}
verbose_logger.debug(
"OCR file input: Converted file to document_url data URI (mime=%s, size=%s bytes, name=%s)",
mime_type,
len(file_bytes),
file_name,
)
return {"type": "document_url", "document_url": data_uri}
@client
def ocr(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
model: str,
document: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
request: Final = _public_request("ocr", args, kwargs)
native: Final = select(request) if rust_ocr_enabled() else None
if native is not None:
try:
return cast( # cast-ok: False selects the synchronous result
OCRResponse, native(request, args, kwargs, False)
)
except _decline_types():
pass
fallback: Final = cast( # cast-ok: forward the original call shape through the legacy @client decorator
Callable[..., OCRResponse | Coroutine[object, object, OCRResponse]], legacy.ocr
)
return fallback(*args, **kwargs)
completion_kwargs: Final[dict[str, object]] = {
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"timeout": timeout,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"kwargs": kwargs,
}
try:
_is_async: Final = kwargs.pop("aocr", False) is True
completion_kwargs["aocr"] = _is_async
prepared: Final = _prepare_ocr_request(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
kwargs=kwargs,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout=timeout,
)
model = prepared.model
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
response: Final = base_llm_http_handler.ocr(
model=prepared.model,
document=cast( # cast-ok: preserve legacy document fields for provider validation
dict[str, str], 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,
)
async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape
request: Final = _public_request("aocr", args, kwargs)
native: Final = select(request) if rust_ocr_enabled() else None
if native is not None:
try:
return await cast( # cast-ok: True selects the asynchronous result
Awaitable[OCRResponse], native(request, args, kwargs, True)
)
except _decline_types():
pass
fallback: Final = cast( # cast-ok: forward the original call shape through the legacy @client decorator
Callable[..., Awaitable[OCRResponse]], legacy.aocr
)
return await fallback(*args, **kwargs)
def _decline_types() -> tuple[type[BaseException], ...]:
exception_types: Final = native_exception_types()
return (exception_types[0],) if exception_types is not None else ()
return response
except Exception as e:
error_provider: Final = _error_provider(model, custom_llm_provider)
error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model
raise litellm.exception_type(
model=error_model,
custom_llm_provider=error_provider,
original_exception=e,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)

2
litellm/proxy/ocr_endpoints/endpoints.py
View file

@ -15,7 +15,7 @@ from litellm.llms.base_llm.ocr.transformation import (
OCRResponse,
parse_ocr_request_format,
)
from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing

118
litellm/responses/dispatch.py Normal file
View file

@ -0,0 +1,118 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm.responses import main
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.rust_bridge.catalog import Context, Delivery, Route
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
from litellm.rust_bridge.public_call import bind, optional_bool, optional_mapping, optional_str, signature
from litellm.rust_bridge.responses.entrypoints import (
NATIVE_ARESPONSES,
NATIVE_RESPONSES,
LiteLLMResponsesRequest,
)
from litellm.types.llms.openai import ResponsesAPIResponse
__all__ = ("aresponses", "responses")
ResponsesResult: TypeAlias = ResponsesAPIResponse | BaseResponsesAPIStreamingIterator
PythonResponses: TypeAlias = Callable[..., ResponsesResult | Coroutine[object, object, ResponsesResult]]
PythonAresponses: TypeAlias = Callable[..., Awaitable[ResponsesResult]]
def _python_responses() -> PythonResponses:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonResponses,
main.responses, # noqa: TID251 # dispatch boundary owns this Python fallback
)
def _python_aresponses() -> PythonAresponses:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonAresponses,
main.aresponses, # noqa: TID251 # dispatch boundary owns this Python fallback
)
_PYTHON_RESPONSES: Final = _python_responses()
_RESPONSES: Final = signature(_PYTHON_RESPONSES)
_PYTHON_ARESPONSES: Final = _python_aresponses()
_ARESPONSES: Final = signature(_PYTHON_ARESPONSES)
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMResponsesRequest | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
model: Final = fields.get("model")
extra: Final = optional_mapping(fields.get("kwargs")) or MappingProxyType({})
if not isinstance(model, str):
return None
return LiteLLMResponsesRequest(
model=model,
input=fields.get("input"),
stream=optional_bool(fields.get("stream")),
api_key=optional_str(extra.get("api_key")),
api_base=optional_str(extra.get("api_base")) or optional_str(extra.get("base_url")),
custom_llm_provider=optional_str(fields.get("custom_llm_provider")),
extra_headers=optional_mapping(fields.get("extra_headers")),
kwargs=extra,
)
def _context(request: LiteLLMResponsesRequest) -> Context:
return Context(
Route.RESPONSES,
provider=request.custom_llm_provider,
model=request.model,
delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
)
_DISPATCH: Final = PublicDispatch(
route=Route.RESPONSES,
request=lambda args, kwargs: _public_request(_RESPONSES, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("aresponses") is True,
)
_ADISPATCH: Final = PublicDispatch(
route=Route.RESPONSES,
request=lambda args, kwargs: _public_request(_ARESPONSES, args, kwargs),
context=_context,
)
def responses(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public Responses call shape
) -> ResponsesResult | Coroutine[object, object, ResponsesResult]:
python: Final = _PYTHON_RESPONSES
return _DISPATCH.run(
args,
kwargs,
python=python,
binding=NATIVE_RESPONSES,
native=call_hook,
)
async def aresponses(*args: object, **kwargs: object) -> ResponsesResult: # kwargs-ok: preserve the public call shape
python: Final = _PYTHON_ARESPONSES
return await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=NATIVE_ARESPONSES,
native=call_hook,
)
responses.__doc__ = _PYTHON_RESPONSES.__doc__
responses.__wrapped__ = _PYTHON_RESPONSES # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature
aresponses.__doc__ = _PYTHON_ARESPONSES.__doc__
aresponses.__wrapped__ = _PYTHON_ARESPONSES # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature

2
litellm/responses/file_search/emulated_handler.py
View file

@ -390,7 +390,7 @@ def _synthesize_responses_api_response(
async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover – thin wrapper for patching in tests
from litellm.responses.main import aresponses
from litellm.responses.main import aresponses # noqa: TID251 # inner call must not re-enter file-search emulation
return await aresponses(input=input, model=model, tools=tools, **kwargs)

17
litellm/responses/main.py
View file

@ -67,6 +67,23 @@ else:
from .streaming_iterator import BaseResponsesAPIStreamingIterator
__all__ = (
"acancel_responses",
"acompact_responses",
"adelete_responses",
"aget_responses",
"alist_input_items",
"aresponses",
"aresponses_api_with_mcp",
"cancel_responses",
"compact_responses",
"delete_responses",
"get_responses",
"list_input_items",
"mock_responses_api_response",
"responses",
)
####### ENVIRONMENT VARIABLES ###################
# Initialize any necessary instances or variables here
base_llm_http_handler = BaseLLMHTTPHandler()

2
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -16,7 +16,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
split_server_prefix_from_name,
strip_known_server_prefix,
)
from litellm.responses.main import aresponses
from litellm.responses.main import aresponses # noqa: TID251 # inner call must skip the MCP gateway that invoked it
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.types.llms.openai import (
ResponseInputParam,

4
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -609,7 +609,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
"""Create the initial response iterator by making the first LLM call"""
try:
# Import the core aresponses function that doesn't have MCP logic
from litellm.responses.main import aresponses
from litellm.responses.main import aresponses # noqa: TID251 # core call without MCP logic
# Make the initial response API call - but avoid the MCP wrapper
params: Final[dict[str, object]] = self.original_request_params.copy()
@ -773,7 +773,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.base_iterator = None
return
from litellm.responses.main import aresponses
from litellm.responses.main import aresponses # noqa: TID251 # follow-up call without MCP logic
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)

38
litellm/rust_bridge/_native.pyi
View file

@ -1,44 +1,23 @@
from asyncio import Future
from collections.abc import Coroutine, Mapping, Sequence
from typing import Literal, Never, TypeAlias, final
from typing import Never, final
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.ocr import LiteLLMOcrRequest
_InputSource: TypeAlias = Literal["request", "deployment", "environment"]
from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest
class RustBridgeDeclined(Exception): ...
class RustUpstreamError(Exception): ...
def ocr(
model: str,
document: object,
api_key: str | None = None,
api_base: str | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, object] | None = None,
optional_params: Mapping[str, object] | None = None,
input_sources: Mapping[str, _InputSource] | None = None,
timeout_seconds: float | None = None,
) -> dict[str, object]: ...
def aocr(
model: str,
document: object,
api_key: str | None = None,
api_base: str | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, object] | None = None,
optional_params: Mapping[str, object] | None = None,
input_sources: Mapping[str, _InputSource] | None = None,
timeout_seconds: float | None = None,
) -> Future[dict[str, object]]: ...
def _ocr_lifecycle(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: dict[str, object],
asynchronous: bool,
) -> OCRResponse | Coroutine[object, object, OCRResponse]: ...
) -> OCRResponse: ...
def aocr(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: dict[str, object],
) -> Coroutine[object, object, OCRResponse]: ...
def transcription(
model: str,
audio: object,
@ -134,7 +113,6 @@ __all__ = [
"RustBridgeDeclined",
"RustUpstreamError",
"TokenCounter",
"_ocr_lifecycle",
"achat_completions",
"amessages",
"aocr",

71
litellm/rust_bridge/catalog.py Normal file
View file

@ -0,0 +1,71 @@
"""Declarative Rust/Python selection for routes with Rust integration.
Rules are static data matched top to bottom; the first match wins and a
context with no matching rule stays on Python. Whether the Rust core can serve
a specific request body is not decided here: that is Rust admission, which
signals ``RustBridgeDeclined`` before any provider I/O.
"""
from __future__ import annotations
from dataclasses import dataclass
from enum import Enum, auto
from typing import Final, TypeAlias
from litellm.rust_bridge.configuration import Decision, Rollout
from litellm.rust_bridge.configuration import decision as _decision
class Route(str, Enum):
CHAT_COMPLETIONS = "chat_completions"
MESSAGES = "messages"
RESPONSES = "responses"
TRANSCRIPTION = "transcription"
OCR = "ocr"
class Delivery(Enum):
COMPLETED = auto()
STREAMING = auto()
WEBSOCKET = auto()
@dataclass(frozen=True, slots=True)
class Context:
route: Route
provider: str | None = None
model: str | None = None
delivery: Delivery = Delivery.COMPLETED
@dataclass(frozen=True, slots=True)
class Rule:
route: Route
rollout: Rollout
providers: frozenset[str] | None = None
models: frozenset[str] | None = None
deliveries: frozenset[Delivery] | None = None
def matches(self, context: Context) -> bool:
return (
context.route is self.route
and (self.providers is None or context.provider in self.providers)
and (self.models is None or context.model in self.models)
and (self.deliveries is None or context.delivery in self.deliveries)
)
Rules: TypeAlias = tuple[Rule, ...]
RULES: Final[Rules] = (
Rule(Route.OCR, Rollout.RUST_OPT_OUT),
Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})),
)
def rollout(context: Context, rules: Rules = RULES) -> Rollout:
return next((rule.rollout for rule in rules if rule.matches(context)), Rollout.PYTHON_ONLY)
def decision(context: Context, rules: Rules = RULES) -> Decision:
return _decision(rollout(context, rules))

446
litellm/rust_bridge/chat_completions.py
View file

@ -1,446 +0,0 @@
"""Thin Python wrapper for the native Rust chat completions bridge.
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 dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.exceptions import APIError
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_model_response_object,
)
from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
# Providers whose `/chat/completions` deployments the Rust core can serve. A
# provider outside this set never reaches the bridge.
RUST_CHAT_COMPLETIONS_PROVIDERS: Final = frozenset({"anthropic", "bedrock"})
# `litellm_params` values are `object`, so validate the one this module reads
# rather than narrowing an unparameterized `Mapping` and typing the result Any.
_LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
RUST_RESPONSE_HEADER: Final = "x-litellm-rust"
class RustChatCompletions(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout_seconds: float | None,
) -> Mapping[str, object]:
raise NotImplementedError
class RustAchatCompletions(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout_seconds: float | None,
) -> Awaitable[Mapping[str, object]]:
raise NotImplementedError
class RustChatCompletionsDecline(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
custom_llm_provider: str | None,
) -> str | None:
raise NotImplementedError
class ResponseObserver(Protocol):
"""Invoked with the payload the core returned, on success only.
Lets the caller emit its own `post_call` on whichever path served the
request. Both entry points call it, so the synchronous and asynchronous
paths cannot drift apart the way the pre_call suppression once did.
"""
def __call__(self, rust_response: Mapping[str, object], /) -> None:
raise NotImplementedError
def response_logger(
*,
logging_obj: LiteLLMLoggingObj,
messages: Sequence[object],
api_key: str,
additional_args: Mapping[str, object],
) -> ResponseObserver:
"""A `ResponseObserver` that emits the caller's `post_call` for a Rust-served
request.
The core owns the provider call, so the Python transform that normally
raises this event never runs; without it every `post_call` callback goes
silent on a Rust-served request and `original_response` stays unset. The
payload is the core's normalized response rather than the provider's wire
body, which is the closest thing that crosses the bridge.
"""
def log(rust_response: Mapping[str, object], /) -> None:
logging_obj.post_call(
input=messages,
api_key=api_key,
original_response=json.dumps(rust_response),
additional_args=additional_args,
)
return log
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustChatCompletionsState:
chat_completions: RustChatCompletions | None = None
achat_completions: RustAchatCompletions | None = None
decline: RustChatCompletionsDecline | None = None
_STATE: Final[_RustChatCompletionsState] = _RustChatCompletionsState()
def set_rust_chat_completions(
*,
chat_completions: RustChatCompletions | None | _Unset = _UNSET,
achat_completions: RustAchatCompletions | None | _Unset = _UNSET,
decline: RustChatCompletionsDecline | None | _Unset = _UNSET,
) -> None:
"""Inject the native callables, so tests can supply a double instead of
patching module attributes."""
if not isinstance(chat_completions, _Unset):
_STATE.chat_completions = chat_completions
if not isinstance(achat_completions, _Unset):
_STATE.achat_completions = achat_completions
if not isinstance(decline, _Unset):
_STATE.decline = decline
def load_rust_chat_completions() -> RustChatCompletions | None:
if _STATE.chat_completions is not None:
return _STATE.chat_completions
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustChatCompletions | None = getattr(native_bridge, "chat_completions", None)
return loaded
def load_rust_achat_completions() -> RustAchatCompletions | None:
if _STATE.achat_completions is not None:
return _STATE.achat_completions
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustAchatCompletions | None = getattr(native_bridge, "achat_completions", None)
return loaded
def _load_rust_decline() -> RustChatCompletionsDecline | None:
if _STATE.decline is not None:
return _STATE.decline
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustChatCompletionsDecline | None = getattr(native_bridge, "chat_completions_decline", None)
return loaded
def _anthropic_user_id_reaches_the_body(litellm_params: Mapping[str, object] | None) -> bool:
metadata: Final = litellm_params.get("metadata") if litellm_params is not None else None
try:
entries: Final = _LITELLM_METADATA_ADAPTER.validate_python(metadata)
except ValidationError:
return False
return entries.get("user_id") is not None
def _litellm_metadata_reaches_the_provider(
custom_llm_provider: str | None, litellm_params: Mapping[str, object] | None
) -> bool:
"""Whether the Python transform would promote proxy-owned attribution into the
provider request, below this gate and inside the function the Rust route replaces.
`AnthropicConfig.transform_request` promotes a valid `metadata["user_id"]`
into the Messages body, so the core never sees the key and would send the
request to Anthropic with the abuse-detection attribution missing.
`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the
Converse body whenever the operator armed `bedrock_request_metadata_fields`.
Owning that field also means evicting a caller-supplied one, which the core
cannot do either, so ownership alone is the condition rather than whether
anything resolved.
Deliberately a superset of Python's condition in both cases: declining a
request Python would not have attributed anyway costs only the Rust path,
while missing one loses the attribution silently.
"""
match custom_llm_provider:
case "anthropic":
return _anthropic_user_id_reaches_the_body(litellm_params)
case "bedrock":
return bedrock_request_metadata_is_owned()
case _:
return False
def rust_chat_completions_accepts(
*,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
custom_llm_provider: str | None,
litellm_params: Mapping[str, object] | None,
stream: object,
) -> bool:
"""Whether the Rust path will serve this request.
Asked before the caller commits to either path, so pre-call logging is
emitted exactly once, on whichever path actually runs. The core's own
capability gate answers the second half; it resolves no credentials and
performs no I/O.
"""
if custom_llm_provider not in RUST_CHAT_COMPLETIONS_PROVIDERS:
return False
if stream:
return False
if not rust_enabled():
return False
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
decline: Final = _load_rust_decline()
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 rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path
verbose_logger.debug(
"Rust chat completions gate raised %s; staying on the Python path",
type(rust_error).__name__,
)
return False
if reason is not None:
verbose_logger.debug("Rust chat completions declined (%s); using the Python path", reason)
return False
return True
def _rust_bridge_exceptions() -> tuple[type[BaseException], type[BaseException]] | None:
"""`(declined, upstream_failed)` from the native module, or None when absent."""
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
declined: Final = getattr(native_bridge, "RustBridgeDeclined", None)
upstream: Final = getattr(native_bridge, "RustUpstreamError", None)
if declined is None or upstream is None:
return None
return declined, upstream
def _reraise_or_decline(
rust_error: BaseException,
*,
model: str,
custom_llm_provider: str | None,
) -> None:
"""Re-raise a failure the provider already saw, or return so the caller declines.
A request that never reached the provider is safe to serve on the Python
path. One that did is not: the provider has already done the work, so a
second attempt bills for it twice. Those surface as an `APIError` carrying
the upstream status, which LiteLLM's exception mapping already understands.
"""
exceptions: Final = _rust_bridge_exceptions()
if exceptions is None:
verbose_logger.debug(
"Rust chat completions bridge raised %s; falling back to Python path",
type(rust_error).__name__,
)
return
declined, upstream_failed = exceptions
if isinstance(rust_error, upstream_failed):
args: Final = rust_error.args
status: Final = args[0] if args else 0
message: Final = args[1] if len(args) > 1 else ""
raise APIError(
status_code=int(status) or 500,
message=f"litellm rust chat completions: {message}",
llm_provider=custom_llm_provider or "",
model=model,
)
if not isinstance(rust_error, declined):
raise rust_error
verbose_logger.debug(
"Rust chat completions declined before calling the provider (%s); using the Python path",
rust_error,
)
def _build_model_response(
rust_response: Mapping[str, object],
model_response: ModelResponse,
) -> ModelResponse:
built: Final = convert_to_model_response_object(
response_object=dict(rust_response), # mutable-ok: the converter takes a real dict and rewrites it
model_response_object=model_response,
hidden_params={"additional_headers": {RUST_RESPONSE_HEADER: "true"}}, # mutable-ok: rewritten by the converter
)
if not isinstance(built, ModelResponse):
raise TypeError(f"expected a ModelResponse from the rust path, got {type(built).__name__}")
return built
def chat_completions(
*,
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,
) -> ModelResponse | None:
rust_chat_completions: Final = load_rust_chat_completions()
if rust_chat_completions is None:
return None
try:
rust_response: Final = rust_chat_completions(
model=model,
messages=messages,
optional_params=optional_params,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)
except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw
_reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider)
return None
on_response(rust_response)
return _build_model_response(rust_response, model_response)
async def achat_completions(
*,
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,
) -> ModelResponse | None:
rust_achat_completions: Final = load_rust_achat_completions()
if rust_achat_completions is None:
return None
try:
rust_response: Final = await rust_achat_completions(
model=model,
messages=messages,
optional_params=optional_params,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)
except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw
_reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider)
return None
on_response(rust_response)
return _build_model_response(rust_response, model_response)
async def achat_completions_or_fallback(
*,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
model_response: ModelResponse,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
python_fallback: Callable[[], Awaitable[object]],
) -> object:
"""Await the Rust path, falling back to the caller's own Python path when
the bridge is unavailable or the call fails.
The caller supplies the fallback, so the bridge stays free of provider
dispatch. This exists because a caller that dispatches asynchronously has
already returned a coroutine by the time a Rust failure surfaces, and so
cannot fall back on its own.
"""
response: Final = await achat_completions(
model=model,
messages=messages,
optional_params=optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout=timeout,
on_response=on_response,
)
if response is not None:
return response
return await python_fallback()

0
litellm/rust_bridge/chat_completions/__init__.py Normal file
View file

19
litellm/rust_bridge/chat_completions/callbacks.py Normal file
View file

@ -0,0 +1,19 @@
from __future__ import annotations
from collections.abc import Mapping
from litellm.rust_bridge import failures
from litellm.rust_bridge.chat_completions.entrypoints import LiteLLMChatCompletionsRequest
from litellm.types.utils import ModelResponse
def response(value: Mapping[str, object]) -> ModelResponse:
return ModelResponse(**value)
def arguments(request: LiteLLMChatCompletionsRequest) -> Mapping[str, object]:
return request.kwargs
def map_failure(error: Exception, request: LiteLLMChatCompletionsRequest, request_provider: str) -> Exception:
return failures.map_failure(error, request.model, request_provider, arguments(request))

54
litellm/rust_bridge/chat_completions/entrypoints.py Normal file
View file

@ -0,0 +1,54 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from litellm.rust_bridge.bindings import NativeBinding
from litellm.types.utils import ModelResponse
@dataclass(frozen=True, slots=True)
class LiteLLMChatCompletionsRequest:
model: str
messages: Sequence[object]
stream: bool | None
api_key: str | None
api_base: str | None
custom_llm_provider: str | None
extra_headers: Mapping[str, object] | None
kwargs: Mapping[str, object]
class NativeCompletion(Protocol):
def __call__(
self,
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ModelResponse: ...
class NativeAcompletion(Protocol):
def __call__(
self,
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[ModelResponse]: ...
def _completion_binding(value: object) -> NativeCompletion | None:
if not callable(value):
return None
return cast("NativeCompletion", value) # cast-ok: callable validated at the native binding boundary
def _acompletion_binding(value: object) -> NativeAcompletion | None:
if not callable(value):
return None
return cast("NativeAcompletion", value) # cast-ok: callable validated at the native binding boundary
NATIVE_COMPLETION: Final = NativeBinding("completion", validate=_completion_binding)
NATIVE_ACOMPLETION: Final = NativeBinding("acompletion", validate=_acompletion_binding)

73
litellm/rust_bridge/configuration.py
View file

@ -1,11 +1,27 @@
from __future__ import annotations
import os
from enum import Enum, auto
from typing import Final
DEFAULT_RUST_ENABLED: Final = False
_TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"})
from pydantic import TypeAdapter, ValidationError
from typing_extensions import assert_never
_GLOBAL_ENV_NAME: Final = "LITELLM_RUST"
_ENV_BOOL: Final = TypeAdapter(bool)
class Rollout(Enum):
PYTHON_ONLY = auto()
RUST_OPT_IN = auto()
RUST_OPT_OUT = auto()
RUST_REQUIRED = auto()
class Decision(Enum):
PYTHON = auto()
RUST_WITH_FALLBACK = auto()
RUST_REQUIRED = auto()
class _RustConfiguration:
@ -19,47 +35,56 @@ _CONFIGURATION: Final = _RustConfiguration()
def _parse_env_bool(value: str | None) -> bool | None:
if value is None:
return None
return value.strip().lower() in _TRUE_ENV_VALUES
try:
return _ENV_BOOL.validate_python(value.strip())
except ValidationError:
return None
def resolve_rust_enabled(
def decide(
rollout: Rollout,
*,
process_override: bool | None,
environment_override: bool | None,
release_default: bool = DEFAULT_RUST_ENABLED,
) -> bool:
if process_override is not None:
return process_override
if environment_override is not None:
return environment_override
return release_default
) -> Decision:
match rollout:
case Rollout.PYTHON_ONLY:
return Decision.PYTHON
case Rollout.RUST_REQUIRED:
return Decision.RUST_REQUIRED
case Rollout.RUST_OPT_IN | Rollout.RUST_OPT_OUT:
switch: Final = (
environment_override
if environment_override is not None
else process_override
if process_override is not None
else rollout is Rollout.RUST_OPT_OUT
)
return Decision.RUST_WITH_FALLBACK if switch else Decision.PYTHON
case _:
assert_never(rollout)
def rust_enabled() -> bool:
return resolve_rust_enabled(
def decision(rollout: Rollout) -> Decision:
return decide(
rollout,
process_override=_CONFIGURATION.override,
environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)),
)
def rust_ocr_enabled() -> bool:
environment: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME))
if environment is False:
return False
return resolve_rust_enabled(
process_override=_CONFIGURATION.override,
environment_override=environment,
release_default=True,
)
def rust_enabled() -> bool:
return decision(Rollout.RUST_OPT_IN) is not Decision.PYTHON
def reset_rust_configuration() -> None:
_CONFIGURATION.override = None
def rust(enabled: bool) -> None:
def rust(enabled: bool | None) -> None:
"""Set the process override for optional Rust paths.
Rust-only paths, including Bedrock transcription, are not controlled by this switch.
``PYTHON_ONLY`` and ``RUST_REQUIRED`` routes in the catalog ignore this switch,
and an explicit ``LITELLM_RUST`` environment value wins over it.
"""
_CONFIGURATION.override = enabled

93
litellm/rust_bridge/dispatch.py Normal file
View file

@ -0,0 +1,93 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from typing import Final, Generic, TypeVar
from litellm.rust_bridge import catalog, runtime
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Context, Route, Rules
from litellm.rust_bridge.configuration import Decision
from litellm.rust_bridge.configuration import decision as rollout_decision
RequestT = TypeVar("RequestT")
NativeT = TypeVar("NativeT")
ResultT = TypeVar("ResultT")
NativeHook = Callable[[RequestT, tuple[object, ...], Mapping[str, object]], ResultT]
def call_hook(
hook: NativeHook[RequestT, ResultT],
request: RequestT,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResultT:
return hook(request, args, kwargs)
@dataclass(frozen=True, slots=True)
class PublicDispatch(Generic[RequestT]):
route: Route
request: Callable[[tuple[object, ...], Mapping[str, object]], RequestT | None]
context: Callable[[RequestT], Context]
bypass: Callable[[RequestT], bool] | None = None
def _requires_projection(self, rules: Rules) -> bool:
for rule in rules:
if rule.route is not self.route:
continue
if rule.providers is not None or rule.models is not None or rule.deliveries is not None:
if rollout_decision(rule.rollout) is not Decision.PYTHON:
return True
continue
return rollout_decision(rule.rollout) is not Decision.PYTHON
return False
def run(
self,
args: tuple[object, ...],
kwargs: Mapping[str, object],
*,
python: Callable[..., ResultT],
binding: NativeBinding[NativeT],
native: Callable[[NativeT, RequestT, tuple[object, ...], Mapping[str, object]], ResultT],
rules: Rules | None = None,
) -> ResultT:
selected_rules: Final = catalog.RULES if rules is None else rules
if not self._requires_projection(selected_rules):
return python(*args, **kwargs)
request: Final = self.request(args, kwargs)
if request is None or (self.bypass is not None and self.bypass(request)):
return python(*args, **kwargs)
return runtime.run(
self.context(request),
binding=binding,
native=lambda hook: native(hook, request, args, kwargs),
python=lambda: python(*args, **kwargs),
rules=selected_rules,
)
async def arun(
self,
args: tuple[object, ...],
kwargs: Mapping[str, object],
*,
python: Callable[..., Awaitable[ResultT]],
binding: NativeBinding[NativeT],
native: Callable[[NativeT, RequestT, tuple[object, ...], Mapping[str, object]], Awaitable[ResultT]],
rules: Rules | None = None,
) -> ResultT:
selected_rules: Final = catalog.RULES if rules is None else rules
if not self._requires_projection(selected_rules):
return await python(*args, **kwargs)
request: Final = self.request(args, kwargs)
if request is None or (self.bypass is not None and self.bypass(request)):
return await python(*args, **kwargs)
return await runtime.arun(
self.context(request),
binding=binding,
native=lambda hook: native(hook, request, args, kwargs),
python=lambda: python(*args, **kwargs),
rules=selected_rules,
)

37
litellm/rust_bridge/failures.py Normal file
View file

@ -0,0 +1,37 @@
"""Map a native failure onto LiteLLM's public exception contract."""
from __future__ import annotations
from collections.abc import Mapping
from typing import Final, Protocol, cast # noqa: TID251 # adapts the public exception mapper
import litellm
class ExceptionMapper(Protocol):
def __call__(
self,
*,
model: str,
custom_llm_provider: str | None,
original_exception: Exception,
completion_kwargs: dict[str, object], # mutable-ok: the legacy public exception mapper mutates its kwargs
extra_kwargs: dict[str, object], # mutable-ok: the legacy public exception mapper mutates its kwargs
) -> Exception: ...
def map_failure(error: Exception, model: str, request_provider: str, kwargs: Mapping[str, object]) -> Exception:
mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper
ExceptionMapper, litellm.exception_type
)
try:
return mapper(
model=model.removeprefix(f"{request_provider}/"),
custom_llm_provider=request_provider,
original_exception=error,
completion_kwargs=dict(kwargs), # mutable-ok: exception mapper requires owned kwargs
extra_kwargs=dict(kwargs), # mutable-ok: exception mapper requires owned kwargs
)
except Exception as public_error:
public_error.__context__ = error
return public_error

136
litellm/rust_bridge/messages.py
View file

@ -1,136 +0,0 @@
"""Thin Python wrapper for the native Rust Anthropic Messages bridge."""
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, Protocol, cast
import httpx
from litellm.rust_bridge.timeouts import timeout_to_seconds
class RustMessages(Protocol):
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAmessages(Protocol):
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]:
raise NotImplementedError
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustMessagesState:
messages: RustMessages | None = None
amessages: RustAmessages | None = None
_STATE: Final[_RustMessagesState] = _RustMessagesState()
def set_rust_messages(
*,
messages: RustMessages | None | _Unset = _UNSET,
amessages: RustAmessages | None | _Unset = _UNSET,
) -> None:
if not isinstance(messages, _Unset):
_STATE.messages = messages
if not isinstance(amessages, _Unset):
_STATE.amessages = amessages
def load_rust_messages() -> RustMessages | None:
if _STATE.messages is not None:
return _STATE.messages
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustMessages, getattr(native_bridge, "messages", None))
def load_rust_amessages() -> RustAmessages | None:
if _STATE.amessages is not None:
return _STATE.amessages
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustAmessages, getattr(native_bridge, "amessages", None))
def messages(
*,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_messages: Final = load_rust_messages()
if rust_messages is None:
return None
return rust_messages(
model=model,
body=body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)
async def amessages(
*,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_amessages: Final = load_rust_amessages()
if rust_amessages is None:
return None
return await rust_amessages(
model=model,
body=body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)

0
litellm/rust_bridge/messages/__init__.py Normal file
View file

23
litellm/rust_bridge/messages/callbacks.py Normal file
View file

@ -0,0 +1,23 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import cast # noqa: TID251 # narrows the normalized native payload to the public TypedDict
from litellm.rust_bridge import failures
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
def response(value: Mapping[str, object]) -> AnthropicMessagesResponse:
return cast( # cast-ok: AnthropicMessagesResponse is a TypedDict over the normalized native payload
AnthropicMessagesResponse,
dict(value), # mutable-ok: the public Messages response is a TypedDict the caller may annotate in place
)
def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]:
return request.kwargs
def map_failure(error: Exception, request: LiteLLMMessagesRequest, request_provider: str) -> Exception:
return failures.map_failure(error, request.model, request_provider, arguments(request))

54
litellm/rust_bridge/messages/entrypoints.py Normal file
View file

@ -0,0 +1,54 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from litellm.rust_bridge.bindings import NativeBinding
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
@dataclass(frozen=True, slots=True)
class LiteLLMMessagesRequest:
model: str
messages: Sequence[object]
max_tokens: int
stream: bool | None
api_key: str | None
api_base: str | None
custom_llm_provider: str | None
kwargs: Mapping[str, object]
class NativeMessages(Protocol):
def __call__(
self,
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse: ...
class NativeAmessages(Protocol):
def __call__(
self,
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[AnthropicMessagesResponse]: ...
def _messages_binding(value: object) -> NativeMessages | None:
if not callable(value):
return None
return cast("NativeMessages", value) # cast-ok: callable validated at the native binding boundary
def _amessages_binding(value: object) -> NativeAmessages | None:
if not callable(value):
return None
return cast("NativeAmessages", value) # cast-ok: callable validated at the native binding boundary
NATIVE_MESSAGES: Final = NativeBinding("anthropic_messages_handler", validate=_messages_binding)
NATIVE_AMESSAGES: Final = NativeBinding("anthropic_messages", validate=_amessages_binding)

145
litellm/rust_bridge/ocr.py
View file

@ -1,145 +0,0 @@
"""Thin Python wrapper for the native Rust OCR bridge."""
from __future__ import annotations
from collections.abc import Awaitable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables
import httpx
from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
@dataclass(frozen=True, slots=True)
class LiteLLMOcrRequest:
model: str
document: Mapping[str, object]
api_key: str | None
api_base: str | None
timeout: float | httpx.Timeout | None
custom_llm_provider: str | None
extra_headers: dict[str, object] | None
kwargs: Mapping[str, object]
input_sources: Mapping[str, str] | None = None
class RustOcr(Protocol):
def __call__(
self,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
input_sources: dict[str, str],
timeout_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAocr(Protocol):
def __call__(
self,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
input_sources: dict[str, str],
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]:
raise NotImplementedError
def _as_ocr(value: object) -> RustOcr | None:
return cast(RustOcr, value) if callable(value) else None
def _as_aocr(value: object) -> RustAocr | None:
return cast(RustAocr, value) if callable(value) else None
_OCR: Final = NativeBinding("ocr", validate=_as_ocr)
_AOCR: Final = NativeBinding("aocr", validate=_as_aocr)
def load_rust_ocr() -> RustOcr | None:
return _OCR.load()
def load_rust_aocr() -> RustAocr | None:
return _AOCR.load()
def _response(response: Mapping[str, object]) -> OCRResponse:
provider_native_response: Final = response.get(PROVIDER_NATIVE_RESPONSE_KEY)
normalized: Final = OCRResponse.model_validate(
MappingProxyType({key: value for key, value in response.items() if key != PROVIDER_NATIVE_RESPONSE_KEY})
)
if isinstance(provider_native_response, Mapping):
normalized.set_provider_native_response(provider_native_response)
return normalized
def ocr(
*,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
input_sources: Mapping[str, str] | None = None,
) -> dict[str, object] | None:
rust_ocr: Final = load_rust_ocr()
if rust_ocr is None:
return None
return rust_ocr(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict
timeout_seconds=_timeout_to_seconds(timeout),
)
async def aocr(
*,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
input_sources: Mapping[str, str] | None = None,
) -> dict[str, object] | None:
rust_aocr: Final = load_rust_aocr()
if rust_aocr is None:
return None
return await rust_aocr(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict
timeout_seconds=_timeout_to_seconds(timeout),
)

0
litellm/rust_bridge/ocr/__init__.py Normal file
View file

31
litellm/rust_bridge/ocr/callbacks.py Normal file
View file

@ -0,0 +1,31 @@
from __future__ import annotations
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter
from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse
from litellm.rust_bridge import failures
from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest
_RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object])
def response(value: Mapping[str, object]) -> OCRResponse:
provider_native_response: Final = value.get(PROVIDER_NATIVE_RESPONSE_KEY)
normalized: Final = OCRResponse.model_validate(
MappingProxyType({key: item for key, item in value.items() if key != PROVIDER_NATIVE_RESPONSE_KEY})
)
if isinstance(provider_native_response, Mapping):
normalized.set_provider_native_response(_RESPONSE_ADAPTER.validate_python(provider_native_response))
return normalized
def arguments(request: LiteLLMOcrRequest) -> Mapping[str, object]:
return request.kwargs
def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: str) -> Exception:
return failures.map_failure(error, request.model, request_provider, arguments(request))

57
litellm/rust_bridge/ocr/entrypoints.py Normal file
View file

@ -0,0 +1,57 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
import httpx
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.bindings import NativeBinding
@dataclass(frozen=True, slots=True)
class LiteLLMOcrRequest:
model: str
document: Mapping[str, object]
api_key: str | None
api_base: str | None
timeout: float | httpx.Timeout | None
custom_llm_provider: str | None
extra_headers: dict[str, object] | None
kwargs: Mapping[str, object]
input_sources: Mapping[str, str] | None = None
class NativeOcr(Protocol):
def __call__(
self,
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse: ...
class NativeAocr(Protocol):
def __call__(
self,
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[OCRResponse]: ...
def _ocr_binding(value: object) -> NativeOcr | None:
if not callable(value):
return None
return cast("NativeOcr", value) # cast-ok: callable validated at the native binding boundary
def _aocr_binding(value: object) -> NativeAocr | None:
if not callable(value):
return None
return cast("NativeAocr", value) # cast-ok: callable validated at the native binding boundary
NATIVE_OCR: Final = NativeBinding("ocr", validate=_ocr_binding)
NATIVE_AOCR: Final = NativeBinding("aocr", validate=_aocr_binding)

67
litellm/rust_bridge/ocr_lifecycle.py
View file

@ -1,67 +0,0 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.ocr import LiteLLMOcrRequest
class NativeOcrLifecycle(Protocol):
def __call__(
self,
request: LiteLLMOcrRequest,
args: Sequence[object],
kwargs: Mapping[str, object],
asynchronous: bool,
) -> OCRResponse | Awaitable[OCRResponse]: ...
class ExceptionMapper(Protocol):
def __call__(
self,
*,
model: str,
custom_llm_provider: str | None,
original_exception: Exception,
completion_kwargs: dict[str, object],
extra_kwargs: dict[str, object],
) -> Exception: ...
def _binding(value: object) -> NativeOcrLifecycle | None:
if not callable(value):
return None
return cast("NativeOcrLifecycle", value) # cast-ok: callable validated at the native binding boundary
NATIVE_OCR_LIFECYCLE: Final = NativeBinding("_ocr_lifecycle", validate=_binding)
def select(request: LiteLLMOcrRequest) -> NativeOcrLifecycle | None:
if request.kwargs.get("aocr"):
return None
return NATIVE_OCR_LIFECYCLE.load()
def arguments(request: LiteLLMOcrRequest) -> Mapping[str, object]:
return request.kwargs
def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: str) -> Exception:
mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper
ExceptionMapper, litellm.exception_type
)
try:
return mapper(
model=request.model.removeprefix(f"{request_provider}/"),
custom_llm_provider=request_provider,
original_exception=error,
completion_kwargs=dict(arguments(request)), # mutable-ok: exception mapper requires owned kwargs
extra_kwargs=dict(request.kwargs), # mutable-ok: exception mapper requires owned kwargs
)
except Exception as public_error:
public_error.__context__ = error
return public_error

42
litellm/rust_bridge/public_call.py Normal file
View file

@ -0,0 +1,42 @@
"""Bind a public LiteLLM call to its legacy Python signature without running it."""
from __future__ import annotations
import inspect
from collections.abc import Callable, Mapping, Sequence
from typing import Final, cast # noqa: TID251 # narrows caller-owned containers without copying them
def signature(legacy: Callable[..., object]) -> inspect.Signature:
return inspect.signature(legacy)
def bind(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> Mapping[str, object] | None:
try:
bound: Final = legacy.bind(*args, **kwargs)
except TypeError:
return None
bound.apply_defaults()
return bound.arguments
def optional_str(value: object) -> str | None:
return value if isinstance(value, str) else None
def optional_bool(value: object) -> bool | None:
return value if isinstance(value, bool) else None
def optional_mapping(value: object) -> Mapping[str, object] | None:
if not isinstance(value, Mapping):
return None
return cast("Mapping[str, object]", value) # cast-ok: the same caller-owned object is handed on unchanged
def optional_sequence(value: object) -> Sequence[object] | None:
if isinstance(value, str | bytes) or not isinstance(value, Sequence):
return None
return cast("Sequence[object]", value) # cast-ok: the same caller-owned object is handed on unchanged

12
litellm/rust_bridge/response_metadata.py Normal file
View file

@ -0,0 +1,12 @@
from typing import TypeVar
from litellm.router_utils.add_retry_fallback_headers import (
_add_headers_to_response, # pyright: ignore[reportPrivateUsage] # reuse the proxy's identity-preserving response metadata writer
)
ResultT = TypeVar("ResultT")
def mark_rust_response(response: ResultT) -> ResultT:
_add_headers_to_response(response, {"x-litellm-rust": "true"})
return response

0
litellm/rust_bridge/responses/__init__.py Normal file
View file

19
litellm/rust_bridge/responses/callbacks.py Normal file
View file

@ -0,0 +1,19 @@
from __future__ import annotations
from collections.abc import Mapping
from litellm.rust_bridge import failures
from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest
from litellm.types.llms.openai import ResponsesAPIResponse
def response(value: Mapping[str, object]) -> ResponsesAPIResponse:
return ResponsesAPIResponse.model_validate(value)
def arguments(request: LiteLLMResponsesRequest) -> Mapping[str, object]:
return request.kwargs
def map_failure(error: Exception, request: LiteLLMResponsesRequest, request_provider: str) -> Exception:
return failures.map_failure(error, request.model, request_provider, arguments(request))

54
litellm/rust_bridge/responses/entrypoints.py Normal file
View file

@ -0,0 +1,54 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from litellm.rust_bridge.bindings import NativeBinding
from litellm.types.llms.openai import ResponsesAPIResponse
@dataclass(frozen=True, slots=True)
class LiteLLMResponsesRequest:
model: str
input: object
stream: bool | None
api_key: str | None
api_base: str | None
custom_llm_provider: str | None
extra_headers: Mapping[str, object] | None
kwargs: Mapping[str, object]
class NativeResponses(Protocol):
def __call__(
self,
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse: ...
class NativeAresponses(Protocol):
def __call__(
self,
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[ResponsesAPIResponse]: ...
def _responses_binding(value: object) -> NativeResponses | None:
if not callable(value):
return None
return cast("NativeResponses", value) # cast-ok: callable validated at the native binding boundary
def _aresponses_binding(value: object) -> NativeAresponses | None:
if not callable(value):
return None
return cast("NativeAresponses", value) # cast-ok: callable validated at the native binding boundary
NATIVE_RESPONSES: Final = NativeBinding("responses", validate=_responses_binding)
NATIVE_ARESPONSES: Final = NativeBinding("aresponses", validate=_aresponses_binding)

0
litellm/rust_bridge/responses_websocket.py → litellm/rust_bridge/responses/websocket.py
View file

93
litellm/rust_bridge/runtime.py
View file

@ -2,21 +2,20 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from enum import Enum
from typing import Final, Generic, NoReturn, TypeAlias, TypeVar
from typing_extensions import assert_never
from litellm.exceptions import APIError
from litellm.rust_bridge.bindings import native_exception_types
from litellm.rust_bridge.bindings import NativeBinding, native_exception_types
from litellm.rust_bridge.catalog import RULES, Context, Rules, decision
from litellm.rust_bridge.configuration import Decision
from litellm.rust_bridge.response_metadata import mark_rust_response
NativeT = TypeVar("NativeT")
ResultT = TypeVar("ResultT")
class FallbackMode(Enum):
PYTHON = "python"
RUST_REQUIRED = "rust_required"
@dataclass(frozen=True, slots=True)
class RustHandled(Generic[ResultT]):
value: ResultT
@ -42,36 +41,68 @@ class BridgeErrorContext:
model: str
def invoke(
def run(
context: Context,
*,
native_call: Callable[[], NativeT] | None,
fallback: Callable[[], ResultT],
adapt: Callable[[NativeT], ResultT],
mode: FallbackMode,
context: BridgeErrorContext,
binding: NativeBinding[NativeT],
native: Callable[[NativeT], ResultT],
python: Callable[[], ResultT],
rules: Rules | None = None,
) -> ResultT:
result: Final = attempt(native_call=native_call, adapt=adapt, context=context)
if isinstance(result, RustHandled):
return result.value
if mode is FallbackMode.PYTHON:
return fallback()
_raise_required(result, context)
selected: Final = decision(context, RULES if rules is None else rules)
match selected:
case Decision.PYTHON:
return python()
case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED:
loaded: Final = binding.load()
result: Final = attempt(
native_call=None if loaded is None else lambda: native(loaded),
adapt=_identity,
context=_error_context(context),
)
if isinstance(result, RustHandled):
return mark_rust_response(result.value)
if selected is Decision.RUST_REQUIRED:
_raise_required(result, _error_context(context))
return python()
case _:
assert_never(selected)
async def ainvoke(
async def arun(
context: Context,
*,
native_call: Callable[[], Awaitable[NativeT]] | None,
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
mode: FallbackMode,
context: BridgeErrorContext,
binding: NativeBinding[NativeT],
native: Callable[[NativeT], Awaitable[ResultT]],
python: Callable[[], Awaitable[ResultT]],
rules: Rules | None = None,
) -> ResultT:
result: Final = await aattempt(native_call=native_call, adapt=adapt, context=context)
if isinstance(result, RustHandled):
return result.value
if mode is FallbackMode.PYTHON:
return await fallback()
_raise_required(result, context)
selected: Final = decision(context, RULES if rules is None else rules)
match selected:
case Decision.PYTHON:
return await python()
case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED:
loaded: Final = binding.load()
result: Final = await aattempt(
native_call=None if loaded is None else lambda: native(loaded),
adapt=_identity,
context=_error_context(context),
)
if isinstance(result, RustHandled):
return mark_rust_response(result.value)
if selected is Decision.RUST_REQUIRED:
_raise_required(result, _error_context(context))
return await python()
case _:
assert_never(selected)
def _identity(value: ResultT) -> ResultT:
return value
def _error_context(context: Context) -> BridgeErrorContext:
return BridgeErrorContext(route=context.route.value, provider=context.provider or "", model=context.model or "")
def attempt(

148
litellm/rust_bridge/transcription.py
View file

@ -1,148 +0,0 @@
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, Protocol, cast
import httpx
from litellm.rust_bridge.timeouts import timeout_to_seconds
class RustTranscription(Protocol):
def __call__(
self,
model: str,
audio: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAtranscription(Protocol):
def __call__(
self,
model: str,
audio: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]:
raise NotImplementedError
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass
class _RustTranscriptionState:
transcription: RustTranscription | None = None
atranscription: RustAtranscription | None = None
_STATE: Final = _RustTranscriptionState()
def configure_rust_transcription(
*,
transcription: RustTranscription | None | _Unset = _UNSET,
atranscription: RustAtranscription | None | _Unset = _UNSET,
) -> None:
if not isinstance(transcription, _Unset):
_STATE.transcription = transcription
if not isinstance(atranscription, _Unset):
_STATE.atranscription = atranscription
def load_rust_transcription() -> RustTranscription | None:
if _STATE.transcription is not None:
return _STATE.transcription
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
return (
None
if native_bridge is None
else cast( # cast-ok: native extension protocol is runtime-defined
RustTranscription, getattr(native_bridge, "transcription", None)
)
)
def load_rust_atranscription() -> RustAtranscription | None:
if _STATE.atranscription is not None:
return _STATE.atranscription
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
return (
None
if native_bridge is None
else cast( # cast-ok: native extension protocol is runtime-defined
RustAtranscription, getattr(native_bridge, "atranscription", None)
)
)
def transcription(
*,
model: str,
audio: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_transcription: Final = load_rust_transcription()
if rust_transcription is None:
return None
return rust_transcription(
model=model,
audio=audio,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
)
async def atranscription(
*,
model: str,
audio: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_atranscription: Final = load_rust_atranscription()
if rust_atranscription is None:
return None
return await rust_atranscription(
model=model,
audio=audio,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
optional_params=optional_params,
timeout_seconds=timeout_to_seconds(timeout),
)

0
litellm/rust_bridge/transcription/__init__.py Normal file
View file

52
litellm/rust_bridge/transcription/native.py Normal file
View file

@ -0,0 +1,52 @@
from __future__ import annotations
from collections.abc import Awaitable
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from litellm.rust_bridge.bindings import NativeBinding
class RustTranscription(Protocol):
def __call__(
self,
model: str,
audio: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAtranscription(Protocol):
def __call__(
self,
model: str,
audio: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]:
raise NotImplementedError
def _sync_binding(value: object) -> RustTranscription | None:
if not callable(value):
return None
return cast("RustTranscription", value) # cast-ok: callable validated at the native binding boundary
def _async_binding(value: object) -> RustAtranscription | None:
if not callable(value):
return None
return cast("RustAtranscription", value) # cast-ok: callable validated at the native binding boundary
NATIVE_TRANSCRIPTION: Final = NativeBinding("transcription", validate=_sync_binding)
NATIVE_ATRANSCRIPTION: Final = NativeBinding("atranscription", validate=_async_binding)

12
ruff-strict.toml
View file

@ -57,3 +57,15 @@ max-args = 5
"typing_extensions.TypeGuard".msg = "Same as typing.TypeGuard."
"typing.TypeIs".msg = "Unverified narrowing (the body is trusted). Parse into a concrete type instead."
"typing_extensions.TypeIs".msg = "Same as typing.TypeIs."
# Dispatched public entry points: import them from their dispatch module so every
# supported call path selects Rust or Python in one place. Only the dispatch
# modules and internal recursive calls may reach the Python implementation
# directly, each with a `# noqa: TID251 # <reason>`.
"litellm.responses.main.responses".msg = "Import litellm.responses.dispatch.responses so the call routes through dispatch."
"litellm.responses.main.aresponses".msg = "Import litellm.responses.dispatch.aresponses so the call routes through dispatch."
"litellm.llms.anthropic.experimental_pass_through.messages.handler.anthropic_messages".msg = "Import litellm.messages.anthropic_messages so the call routes through dispatch."
"litellm.llms.anthropic.experimental_pass_through.messages.handler.anthropic_messages_handler".msg = "Import litellm.messages.anthropic_messages_handler so the call routes through dispatch."
"litellm.ocr.main.ocr".msg = "Import litellm.ocr.dispatch.ocr so the call routes through dispatch."
"litellm.ocr.main.aocr".msg = "Import litellm.ocr.dispatch.aocr so the call routes through dispatch."
"litellm.main.completion".msg = "Import litellm.completion so the call routes through dispatch."
"litellm.main.acompletion".msg = "Import litellm.acompletion so the call routes through dispatch."

27
tests/e2e/access_control/test_model_access_group_e2e.py
View file

@ -23,7 +23,7 @@ from access_control_client import (
MODEL_ACCESS_DENIED_MARKER,
TEAM_MODEL_ACCESS_DENIED_MARKER,
)
from e2e_config import unique_marker
from e2e_config import settle_propagation, unique_marker
from lifecycle import ResourceManager
from models import (
ChatResponse,
@ -31,6 +31,7 @@ from models import (
LiteLLMParamsBody,
ModelInfoBody,
ModelNewBody,
TeamInfoResponse,
)
pytestmark = pytest.mark.e2e
@ -111,20 +112,6 @@ def _await_group_members(client: AccessControlClient, access_group: str, expecte
)
def _await_team_allowlist(client: AccessControlClient, team_id: str, access_group: str) -> None:
deadline = time.monotonic() + client.proxy.poll_timeout
listed: list[str] | None = None
while time.monotonic() < deadline:
listed = client.team_models(team_id)
if listed == [access_group]:
return
time.sleep(client.proxy.poll_interval)
pytest.fail(
f"/team/info never settled the team's allow-list to [{access_group!r}] after the team-scoped "
f"deployment was registered; last read {listed}"
)
@pytest.fixture(scope="module")
def grouped(client: AccessControlClient) -> Iterator[GroupedDeployments]:
marker: Final = unique_marker()
@ -168,9 +155,15 @@ def team_grant(client: AccessControlClient) -> Iterator[TeamGrant]:
),
listed_for=key,
)
client.set_team_models(team_id, team_alias, [access_group])
try:
_await_team_allowlist(client, team_id, access_group)
client.set_team_models(team_id, team_alias, [access_group])
written_at: Final = time.monotonic()
_ = client.proxy.read_body_back_everywhere(
f"/team/info?team_id={team_id}",
TeamInfoResponse,
settled=lambda response: response.team_id == team_id and response.team_info.models == [access_group],
)
settle_propagation(written_at)
yield TeamGrant(access_group=access_group, team_id=team_id, key=key)
finally:
client.proxy.delete_model(model_id)

2
tests/e2e/e2e_config.py
View file

@ -101,8 +101,6 @@ SLOW_PROVIDER_TIMEOUT_SECONDS = float(os.environ.get("E2E_SLOW_PROVIDER_TIMEOUT"
# fresh connection and the next call re-rolls. See ProxyClient._await_model_servable.
PROPAGATION_TIMEOUT = float(os.environ.get("E2E_PROPAGATION_TIMEOUT", "15"))
EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes")
# Record/replay fixture selection (see fixture_mode.py and provider_edge.py).
# The raw mode value is parsed and validated there; "live" (the default, also
# for empty values) means the harness behaves exactly as before this knob

10
tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py
View file

@ -11,8 +11,7 @@ sent in the request.
from __future__ import annotations
import pytest
from e2e_config import EXPECT_RUST, unique_marker
from e2e_config import unique_marker
from e2e_http import StreamingResponse, require_successful_call, unwrap
from endpoints_client import EndpointsClient
from lifecycle import ResourceManager
@ -50,13 +49,6 @@ def _assert_streamed_ok(result: StreamingResponse) -> None:
assert any("message_stop" in event for event in result.stream_events), (
"stream never reached message_stop"
)
if EXPECT_RUST:
assert result.headers.get("x-litellm-rust") == "true", (
"E2E_EXPECT_RUST is set, so this gateway must serve /v1/messages through the "
"Rust path, but the response carried no x-litellm-rust marker. The request "
"still succeeded, which is exactly the failure mode: a gateway whose native "
f"extension is unavailable falls back to Python silently. headers={result.headers}"
)
class TestAzureFoundryMessages:

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

@ -1,394 +0,0 @@
"""Tests for the optional Rust-backed Anthropic Messages path."""
import importlib
from typing import cast
import httpx
import pytest
import litellm
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import configuration
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.router import GenericLiteLLMParams
rust_messages = importlib.import_module("litellm.rust_bridge.messages")
rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
FAKE_MESSAGES_RESPONSE: dict[str, object] = {
"id": "msg_123",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5-20250929",
"content": [{"type": "text", "text": "hello world"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 5, "output_tokens": 3},
}
REQUEST_BODY: dict[str, object] = {
"model": "claude-sonnet-4-5",
"max_tokens": 64,
"messages": [{"role": "user", "content": "hi"}],
}
class RecordingMessages:
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> dict[str, object]:
self.calls.append(
{
"model": model,
"body": body,
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"timeout_seconds": timeout_seconds,
}
)
return dict(FAKE_MESSAGES_RESPONSE)
class RecordingAsyncMessages:
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
async def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> dict[str, object]:
self.calls.append(
{
"model": model,
"body": body,
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"timeout_seconds": timeout_seconds,
}
)
return dict(FAKE_MESSAGES_RESPONSE)
class ExplodingAsyncMessages:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, **kwargs: object) -> dict[str, object]:
self.calls += 1
raise AssertionError("bridge must not be called")
class RaisingAsyncMessages:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, **kwargs: object) -> dict[str, object]:
self.calls += 1
raise RuntimeError("upstream request failed with status 400: bad request")
@pytest.fixture(autouse=True)
def _reset_rust_flag():
rust_messages.set_rust_messages(messages=None, amessages=None)
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
yield
rust_messages.set_rust_messages(messages=None, amessages=None)
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
def test_load_rust_messages_returns_injected_impl():
bridge = RecordingMessages()
litellm.rust(True)
rust_messages.set_rust_messages(messages=bridge)
assert rust_messages.load_rust_messages() is bridge
def test_load_rust_amessages_returns_injected_impl():
bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
assert rust_messages.load_rust_amessages() is bridge
def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
"get_native_bridge",
lambda: None,
)
litellm.rust(True)
assert rust_messages.load_rust_messages() is None
result = rust_messages.messages(
model="claude",
body=REQUEST_BODY,
api_key="k",
api_base="b",
custom_llm_provider="azure_ai",
extra_headers={},
timeout=30.0,
)
assert result is None
def test_messages_wrapper_forwards_args_and_converts_timeout():
bridge = RecordingMessages()
litellm.rust(True)
rust_messages.set_rust_messages(messages=bridge)
response = rust_messages.messages(
model="claude-sonnet-4-5",
body=REQUEST_BODY,
api_key="sk-azure",
api_base="https://resource.services.ai.azure.com/anthropic",
custom_llm_provider="azure_ai",
extra_headers={"anthropic-beta": "token-efficient-tools-2025-02-19"},
timeout=httpx.Timeout(600.0, read=42.0),
)
assert response == FAKE_MESSAGES_RESPONSE
assert bridge.calls[0] == {
"model": "claude-sonnet-4-5",
"body": REQUEST_BODY,
"api_key": "sk-azure",
"api_base": "https://resource.services.ai.azure.com/anthropic",
"custom_llm_provider": "azure_ai",
"extra_headers": {"anthropic-beta": "token-efficient-tools-2025-02-19"},
"timeout_seconds": 42.0,
}
@pytest.mark.asyncio
async def test_amessages_wrapper_forwards_args():
bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await rust_messages.amessages(
model="claude-sonnet-4-5",
body=REQUEST_BODY,
api_key="sk-azure",
api_base="https://resource.services.ai.azure.com/anthropic",
custom_llm_provider="azure_ai",
extra_headers=None,
timeout=12.5,
)
assert response == FAKE_MESSAGES_RESPONSE
assert bridge.calls[0]["model"] == "claude-sonnet-4-5"
assert bridge.calls[0]["timeout_seconds"] == 12.5
def _gate(**overrides):
kwargs = {
"custom_llm_provider": "azure_ai",
"litellm_params": GenericLiteLLMParams(api_key="sk-azure"),
"has_agentic_hook": False,
"model": "claude-sonnet-4-5",
"api_key": "sk-azure",
"api_base": "https://resource.services.ai.azure.com/anthropic",
"headers": {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"},
"request_body": dict(REQUEST_BODY),
"timeout": 30.0,
}
kwargs.update(overrides)
return BaseLLMHTTPHandler._maybe_rust_anthropic_messages(**kwargs)
@pytest.mark.asyncio
async def test_gate_invokes_rust_and_marks_response_header():
bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate()
assert response is not None
assert response["id"] == "msg_123"
assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
call = bridge.calls[0]
assert call["model"] == "claude-sonnet-4-5"
assert call["body"] == REQUEST_BODY
assert call["api_key"] == "sk-azure"
assert call["api_base"] == "https://resource.services.ai.azure.com/anthropic"
assert call["extra_headers"] == {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"}
assert call["timeout_seconds"] == 30.0
@pytest.mark.asyncio
async def test_gate_falls_back_to_python_when_bridge_raises():
bridge = RaisingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate()
assert response is None
assert bridge.calls == 1
@pytest.mark.asyncio
async def test_gate_skips_rust_when_flag_absent():
bridge = ExplodingAsyncMessages()
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure"))
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_gate_uses_process_enable_without_request_override():
bridge = RecordingAsyncMessages()
rust_messages.set_rust_messages(amessages=bridge)
litellm.rust(True)
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure"))
assert response is not None
assert bridge.calls[0]["custom_llm_provider"] == "azure_ai"
@pytest.mark.asyncio
async def test_gate_invokes_rust_for_native_anthropic_provider():
bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate(
custom_llm_provider="anthropic",
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
api_key="sk-ant",
api_base="https://api.anthropic.com",
headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"},
)
assert response is not None
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"
@pytest.mark.asyncio
async def test_gate_invokes_rust_when_env_var_set(monkeypatch):
bridge = RecordingAsyncMessages()
rust_messages.set_rust_messages(amessages=bridge)
monkeypatch.setenv("LITELLM_RUST", "1")
response = await _gate(
custom_llm_provider="anthropic",
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
)
assert response is not None
assert bridge.calls[0]["custom_llm_provider"] == "anthropic"
@pytest.mark.asyncio
async def test_gate_env_var_falsey_does_not_enable(monkeypatch):
bridge = ExplodingAsyncMessages()
rust_messages.set_rust_messages(amessages=bridge)
monkeypatch.setenv("LITELLM_RUST", "0")
response = await _gate(
custom_llm_provider="anthropic",
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
)
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_gate_skips_rust_for_unsupported_provider():
bridge = ExplodingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate(custom_llm_provider="openai")
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_gate_skips_rust_for_agentic_hook():
bridge = ExplodingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate(has_agentic_hook=True)
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag():
bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
streaming_body = {**REQUEST_BODY, "stream": True}
response = await _gate(
has_agentic_hook=False,
request_body=streaming_body,
)
assert response is not None
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
@pytest.mark.asyncio
async def test_fake_stream_wraps_rust_response_as_anthropic_sse():
response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE))
stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response)
assert stream._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
chunks = [chunk async for chunk in stream]
joined = b"".join(chunks)
assert b"event: message_start" in joined
assert b"event: content_block_delta" in joined
assert b"hello world" in joined
assert b"event: message_stop" in joined
@pytest.mark.asyncio
async def test_gate_falls_back_when_bridge_unavailable(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
"get_native_bridge",
lambda: None,
)
litellm.rust(True)
response = await _gate()
assert response is None

0
tests/test_litellm/chat_completions/__init__.py Normal file
View file

271
tests/test_litellm/chat_completions/test_dispatch.py Normal file
View file

@ -0,0 +1,271 @@
import inspect
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect
import pytest
import litellm
from litellm import main as python_chat
from litellm.chat_completions.dispatch import (
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
)
from litellm.rust_bridge import catalog
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Route, Rule
from litellm.rust_bridge.chat_completions.entrypoints import (
NATIVE_ACOMPLETION,
NATIVE_COMPLETION,
LiteLLMChatCompletionsRequest,
NativeAcompletion,
NativeCompletion,
)
from litellm.rust_bridge.configuration import Rollout
from litellm.types.utils import ModelResponse
MESSAGES: Final = [{"role": "user", "content": "hi"}]
PYTHON_RULES: Final = ()
RUST_RULES: Final = (Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
def completion_binding(native: NativeCompletion | None) -> NativeBinding[NativeCompletion]:
binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None)
binding.override(native)
return binding
def acompletion_binding(native: NativeAcompletion | None) -> NativeBinding[NativeAcompletion]:
binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None)
binding.override(native)
return binding
def test_public_signature_is_the_legacy_signature() -> None:
public_completion: Final = cast(Callable[..., object], litellm.completion)
legacy_completion: Final = cast(Callable[..., object], python_chat.completion)
public_acompletion: Final = cast(Callable[..., object], litellm.acompletion)
legacy_acompletion: Final = cast(Callable[..., object], python_chat.acompletion)
assert inspect.signature(public_completion) == inspect.signature(legacy_completion)
assert inspect.signature(public_acompletion) == inspect.signature(legacy_acompletion)
def test_python_route_forwards_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = ("gpt-4o", MESSAGES)
kwargs: Final[Mapping[str, object]] = {"temperature": 0.1, "metadata": metadata}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
response: Final = ModelResponse()
def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: records public call shape
captured.append((call_args, call_kwargs))
return response
def native(
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> ModelResponse:
pytest.fail("Python-only dispatch must not call native")
assert (
_DISPATCH.run(
args,
kwargs,
python=python,
binding=completion_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=PYTHON_RULES,
)
is response
)
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[1] is MESSAGES
assert call_kwargs == kwargs
assert call_kwargs["metadata"] is metadata
assert kwargs == {"temperature": 0.1, "metadata": metadata}
@pytest.mark.asyncio
async def test_async_python_route_forwards_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = ("gpt-4o", MESSAGES)
kwargs: Final[Mapping[str, object]] = {"temperature": 0.1, "metadata": metadata}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
response: Final = ModelResponse()
async def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: records call shape
captured.append((call_args, call_kwargs))
return response
async def native(
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> ModelResponse:
pytest.fail("Python-only dispatch must not call native")
result: Final = await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=acompletion_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=PYTHON_RULES,
)
assert result is response
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[1] is MESSAGES
assert call_kwargs == kwargs
assert call_kwargs["metadata"] is metadata
assert kwargs == {"temperature": 0.1, "metadata": metadata}
def test_native_receives_bound_request_and_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
kwargs: Final[Mapping[str, object]] = {
"stream": True,
"api_key": "sk-test",
"base_url": "https://example.invalid",
"extra_headers": {"x-test": "1"},
"custom_llm_provider": "anthropic",
"metadata": metadata,
}
captured: Final[
list[tuple[LiteLLMChatCompletionsRequest, tuple[object, ...], Mapping[str, object]]]
] = []
def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: rejected Rust fallback
pytest.fail("Required Rust dispatch must not call Python")
def native(
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> ModelResponse:
captured.append((request, args, kwargs))
return ModelResponse()
args: Final[tuple[object, ...]] = ("anthropic/claude-sonnet-4-5", MESSAGES)
_DISPATCH.run(
args,
kwargs,
python=python,
binding=completion_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
request, call_args, call_kwargs = captured[0]
assert request.model == "anthropic/claude-sonnet-4-5"
assert request.messages is MESSAGES
assert request.stream is True
assert request.api_key == "sk-test"
assert request.api_base == "https://example.invalid"
assert request.custom_llm_provider == "anthropic"
assert request.extra_headers == {"x-test": "1"}
assert request.kwargs == {"custom_llm_provider": "anthropic", "metadata": metadata}
assert call_args == args
assert call_kwargs == kwargs
assert call_kwargs["metadata"] is metadata
def test_internal_async_marker_bypasses_native() -> None:
response: Final = ModelResponse()
called: Final[list[bool]] = []
def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: records public call shape
called.append(True)
return response
def native(
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> ModelResponse:
pytest.fail("acompletion's inner completion call must stay on Python")
result: Final = _DISPATCH.run(
("gpt-4o", MESSAGES),
{"acompletion": True},
python=python,
binding=completion_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
assert result is response
assert called == [True]
@pytest.mark.parametrize(
("args", "kwargs"),
(
(("gpt-4o", MESSAGES), {"model": "duplicate"}),
((), {}),
),
)
def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Mapping[str, object]) -> None:
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
response: Final = ModelResponse()
def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: records invalid call shape
captured.append((call_args, call_kwargs))
return response
def native(
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> ModelResponse:
pytest.fail("Binding failures must be delegated to Python")
assert (
_DISPATCH.run(
args,
kwargs,
python=python,
binding=completion_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
is response
)
assert captured == [(args, kwargs)]
def test_public_completion_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
captured: Final[list[LiteLLMChatCompletionsRequest]] = []
expected: Final = ModelResponse()
def native(
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ModelResponse:
captured.append(request)
return expected
NATIVE_COMPLETION.override(native)
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
public_completion: Final = cast(Callable[..., ModelResponse], litellm.completion)
try:
result: Final = public_completion(model="gpt-4o", messages=MESSAGES)
finally:
NATIVE_COMPLETION.reset()
assert result is expected
assert [request.model for request in captured] == ["gpt-4o"]
@pytest.mark.asyncio
async def test_public_acompletion_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
captured: Final[list[LiteLLMChatCompletionsRequest]] = []
expected: Final = ModelResponse()
async def native(
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ModelResponse:
captured.append(request)
return expected
NATIVE_ACOMPLETION.override(native)
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
public_acompletion: Final = cast(Callable[..., Awaitable[ModelResponse]], litellm.acompletion)
try:
result: Final = await public_acompletion(model="gpt-4o", messages=MESSAGES)
finally:
NATIVE_ACOMPLETION.reset()
assert result is expected
assert [request.model for request in captured] == ["gpt-4o"]

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

@ -8,9 +8,9 @@ import httpx
import pytest
import litellm
from litellm._uuid import uuid
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call
from litellm._uuid import uuid
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
@ -2333,48 +2333,7 @@ def test_non_bash_tool_result_skipped():
), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"
class TestRustChatCompletionsHook:
"""The `rust: true` opt-in on `/chat/completions` for the Anthropic provider.
The native callables are dependency-injected, so these run without the
compiled extension.
"""
RUST_RESPONSE = {
"created": 1_700_000_000,
"model": "claude-sonnet-4-5-20260101",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hello from rust"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 11,
"completion_tokens": 4,
"total_tokens": 15,
"prompt_tokens_details": {
"cached_tokens": 0,
"cache_creation_tokens": 0,
"text_tokens": 11,
},
},
}
@pytest.fixture(autouse=True)
def _reset_bridge(self, monkeypatch):
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
)
yield
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
class TestAnthropicChatCompletionPreCallLogging:
@staticmethod
def _completion_kwargs(**overrides):
from litellm.types.utils import ModelResponse
@ -2400,319 +2359,22 @@ class TestRustChatCompletionsHook:
kwargs.update(overrides)
return kwargs
@staticmethod
def _recording_logging_obj():
"""A logging object that keeps each hook's payload in a real list, so a
test can assert which path logged and what it carried."""
calls = {"pre_call": [], "post_call": []}
logging_obj = MagicMock()
logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs)
logging_obj.post_call.side_effect = lambda **kwargs: calls["post_call"].append(kwargs)
return logging_obj, calls
def _inject(self, *, decline_reason=None, sync_result=None, sync_error=None):
from litellm.rust_bridge import chat_completions as bridge
seen = {"gate": [], "call": []}
def gate(**kwargs):
seen["gate"].append(kwargs)
return decline_reason
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)
bridge.set_rust_chat_completions(decline=gate, chat_completions=native)
return seen
def test_rust_true_serves_the_call_and_stamps_the_header(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
seen = self._inject()
response = AnthropicChatCompletion().completion(**self._completion_kwargs())
assert response.choices[0].message.content == "hello from rust"
assert response._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
assert len(seen["call"]) == 1
def test_the_core_receives_the_untranslated_openai_messages(self):
"""Rust owns the translation, so the handler must not pre-translate."""
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
seen = self._inject()
AnthropicChatCompletion().completion(
**self._completion_kwargs(
messages=[
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"},
]
)
)
assert seen["call"][0]["messages"] == [
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"},
]
def test_the_anthropic_max_tokens_default_is_merged_in_before_the_gate(self):
"""`transform_request` applies `AnthropicConfig.get_config`; the Rust
path skips it, so the handler has to merge it or Anthropic 400s on a
request that omits `max_tokens`."""
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
seen = self._inject()
AnthropicChatCompletion().completion(**self._completion_kwargs(optional_params={}))
assert "max_tokens" in seen["gate"][0]["optional_params"]
assert seen["call"][0]["optional_params"]["max_tokens"] > 0
def test_a_caller_supplied_max_tokens_outranks_the_default(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
seen = self._inject()
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):
monkeypatch.setenv("LITELLM_RUST", "0")
def test_pre_call_logging_fires_once_on_the_python_path(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject()
calls = {"pre_call": []}
logging_obj = MagicMock()
logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs)
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(logging_obj=logging_obj))
except Exception:
# The Python path goes on to make an HTTP call; reaching it is
# the assertion, so the network failure below is expected.
pass
assert seen["gate"] == []
assert seen["call"] == []
assert transform.called
def test_a_declined_request_never_reaches_the_native_call(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
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": []}
):
try:
AnthropicChatCompletion().completion(**self._completion_kwargs())
except Exception:
pass
assert len(seen["gate"]) == 1
assert seen["call"] == []
def test_streaming_stays_on_the_python_path(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(optional_params={"max_tokens": 16, "stream": True})
)
except Exception:
pass
assert seen["gate"] == []
def test_pre_call_logging_fires_exactly_once_on_the_rust_path(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
seen = self._inject()
logging_obj = MagicMock()
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
assert logging_obj.pre_call.call_count == 1
assert len(seen["call"]) == 1
def test_post_call_logging_fires_on_the_rust_path(self):
"""The Rust core owns the provider call, so the Python transform that
normally raises `post_call` never runs. Without the bridge hook every
post_call callback goes silent and `original_response` stays unset."""
import json
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
self._inject()
logging_obj = MagicMock()
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"]
assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust"
def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines(self, monkeypatch):
"""A decline never reached the provider, so the Python path serves the
request and owns the only post_call. Firing the hook there too would
double every post_call callback for one request."""
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.rust_bridge import chat_completions as bridge
class _Declined(Exception):
pass
class _FakeNative:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
def declining_native(**_kwargs):
raise _Declined("blank message text")
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
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": []}
):
try:
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 calls["post_call"] == []
@pytest.mark.asyncio
async def test_the_async_path_falls_back_when_the_core_declines(self, monkeypatch):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.rust_bridge import chat_completions as bridge
class _Declined(Exception):
pass
class _FakeNative:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
async def declining_native(**_kwargs):
raise _Declined("blank message text")
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)
)
assert result is sentinel
assert python_call.called, "a failing rust call must re-enter the python path"
@pytest.mark.asyncio
async def test_the_async_path_serves_the_rust_response_without_the_fallback(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.rust_bridge import chat_completions as bridge
async def native(**_kwargs):
return dict(self.RUST_RESPONSE)
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)
)
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."""
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.rust_bridge import chat_completions as bridge
class _Declined(Exception):
pass
class _FakeNative:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
def declining_native(**_kwargs):
raise _Declined("blank message text")
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": []}
):
try:
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"
)
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."""
monkeypatch.setenv("LITELLM_RUST", "0")
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
self._inject()
logging_obj, calls = self._recording_logging_obj()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(litellm_params={}, logging_obj=logging_obj)
)
except Exception:
pass
assert len(calls["pre_call"]) == 1
assert calls["pre_call"][0]["additional_args"]["complete_input_dict"] == {

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

@ -1,49 +1,27 @@
"""Tests for `BedrockConverseLLM.completion`'s Rust chat completions hook.
"""Tests for `BedrockConverseLLM.completion`.
The native callables are dependency-injected, so these run without the compiled
extension, and AWS credential resolution is stubbed so nothing reaches STS.
AWS credential resolution is stubbed so nothing reaches STS.
"""
from __future__ import annotations
import asyncio
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import MagicMock, patch
import boto3
import httpx
import pytest
from botocore.credentials import Credentials
from botocore.exceptions import ClientError
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
from litellm.rust_bridge import configuration
from litellm.types.utils import ModelResponse
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
RUST_RESPONSE = {
"created": 1_700_000_000,
"model": "anthropic.claude-sonnet-4-5-v1:0",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hello from rust"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 11,
"completion_tokens": 4,
"total_tokens": 15,
"prompt_tokens_details": {
"cached_tokens": 0,
"cache_creation_tokens": 0,
"text_tokens": 11,
},
},
}
RESOLVED_CREDENTIALS = Credentials(
access_key="AKIARESOLVED",
secret_key="resolved-secret",
@ -52,32 +30,11 @@ RESOLVED_CREDENTIALS = Credentials(
@pytest.fixture(autouse=True)
def reset_bridge(monkeypatch):
def reset_rust_configuration(monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
configuration.reset_rust_configuration()
yield
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
def _inject(*, decline_reason=None, error: Exception | None = None):
seen: dict[str, list[dict]] = {"gate": [], "call": []}
def gate(**kwargs):
seen["gate"].append(kwargs)
return decline_reason
def native(**kwargs):
seen["call"].append(kwargs)
if error is not None:
raise error
return dict(RUST_RESPONSE)
bridge.set_rust_chat_completions(decline=gate, chat_completions=native)
return seen
configuration.reset_rust_configuration()
def _completion_kwargs(**overrides):
@ -106,206 +63,6 @@ def _run(*, credentials: Credentials | None = RESOLVED_CREDENTIALS, **overrides)
return BedrockConverseLLM().completion(**_completion_kwargs(**overrides))
def _recording_logging_obj():
"""A logging object that keeps each hook's payload in a real list, so a test
can assert which path logged and what it carried."""
calls = {"pre_call": [], "post_call": []}
logging_obj = MagicMock()
logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs)
logging_obj.post_call.side_effect = lambda **kwargs: calls["post_call"].append(kwargs)
return logging_obj, calls
def test_rust_true_serves_the_call_and_stamps_the_header():
seen = _inject()
response = _run()
assert response.choices[0].message.content == "hello from rust"
assert response._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
assert len(seen["call"]) == 1
def test_the_core_receives_the_credentials_this_handler_already_resolved():
"""Both paths must sign as the same principal, so the resolved credentials
are handed down rather than re-derived from ambient AWS state."""
seen = _inject()
_run()
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 "bedrock-runtime.us-east-1.amazonaws.com" in seen["call"][0]["api_base"]
def test_the_core_receives_the_untranslated_openai_messages():
seen = _inject()
_run(
messages=[
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"},
]
)
assert seen["call"][0]["messages"] == [
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"},
]
def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "0")
seen = _inject()
try:
_run(litellm_params={})
except Exception:
# The Python path goes on to make an HTTP call; not reaching the gate
# is the assertion, so a failure past this point is expected.
pass
assert seen["gate"] == []
assert seen["call"] == []
def test_streaming_stays_on_the_python_path():
seen = _inject()
try:
_run(optional_params={"maxTokens": 16, "stream": True})
except Exception:
pass
assert seen["gate"] == []
def test_a_declined_request_never_reaches_the_native_call():
seen = _inject(decline_reason="unrecognized request parameter")
try:
_run()
except Exception:
pass
assert len(seen["gate"]) == 1
assert seen["call"] == []
def test_pre_call_logging_fires_exactly_once_on_the_rust_path():
_inject()
logging_obj = MagicMock()
_run(logging_obj=logging_obj)
assert logging_obj.pre_call.call_count == 1
@pytest.mark.asyncio
async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch):
class _Declined(Exception):
pass
class _FakeNative:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
async def declining_native(**_kwargs):
raise _Declined("blank message text")
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(
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)
)
assert result is sentinel
assert python_call.called, "a failing rust call must re-enter the python path"
@pytest.mark.asyncio
async def test_the_async_path_serves_the_rust_response_without_the_fallback():
async def native(**_kwargs):
return dict(RUST_RESPONSE)
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, "async_completion") as python_call,
):
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"}
assert not python_call.called
@pytest.mark.asyncio
async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines():
"""One request, one pre_call. Without the suppression the Python fallback
logs a second one and non-idempotent callbacks run twice."""
class _Declined(Exception):
pass
class _FakeNative:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
async def declining_native(**_kwargs):
raise _Declined("blank message text")
logging_obj = MagicMock()
served = []
async def python_path(**kwargs):
served.append(kwargs)
return ModelResponse()
with (
patch.object(bridge, "get_native_bridge", lambda: _FakeNative()),
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)
)
assert logging_obj.pre_call.call_count == 1
assert served and served[0]["skip_pre_call_logging"] is True
CONVERSE_RESPONSE = {
"output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}},
"stopReason": "end_turn",
@ -314,7 +71,11 @@ CONVERSE_RESPONSE = {
async def _drive_async_completion(
*, skip_pre_call_logging: bool, logging_obj, credentials: Credentials = RESOLVED_CREDENTIALS
*,
skip_pre_call_logging: bool,
logging_obj,
credentials: Credentials = RESOLVED_CREDENTIALS,
outer_dispatch: bool = False,
):
"""Run the real `async_completion` with a stubbed transport."""
import httpx as _httpx
@ -331,6 +92,9 @@ async def _drive_async_completion(
client.post = post
client.__class__ = AsyncHTTPHandler
if outer_dispatch:
return await _run(credentials=credentials, acompletion=True, client=client, logging_obj=logging_obj)
return await BedrockConverseLLM().async_completion(
model="anthropic.claude-sonnet-4-5-v1:0",
messages=[{"role": "user", "content": "hi"}],
@ -381,6 +145,26 @@ async def test_async_completion_signs_off_the_event_loop(monkeypatch):
assert probe.served_during_refresh is True
@pytest.mark.asyncio
@pytest.mark.parametrize("rust_enabled", (False, True))
async def test_python_only_async_dispatch_refreshes_credentials_off_the_event_loop(
monkeypatch: pytest.MonkeyPatch, rust_enabled: bool
) -> None:
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.setenv("LITELLM_RUST", "1" if rust_enabled else "0")
configuration.rust(rust_enabled)
probe: Final = EventLoopProbe()
release: Final = asyncio.create_task(probe.release_refresh_from_the_loop())
response: Final = await _drive_async_completion(
skip_pre_call_logging=False, logging_obj=MagicMock(), credentials=probe.credentials(), outer_dispatch=True
)
await release
assert response.choices[0].message.content == "hi"
assert probe.served_during_refresh is True
def _sync_client_returning_converse_response():
client = MagicMock()
client.post.side_effect = lambda **_kwargs: httpx.Response(
@ -392,48 +176,10 @@ def _sync_client_returning_converse_response():
return client
def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines():
"""One request, one pre_call, on the synchronous path too.
The gate accepts and logs, then the native call declines before the
provider is reached, so execution continues into the Python path below.
That is the same attempt continuing; without the suppression it logs a
second pre_call and non-idempotent callbacks run twice for one request.
"""
class _Declined(Exception):
pass
class _FakeNative:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
def declining_native(**_kwargs):
raise _Declined("blank message text")
logging_obj = MagicMock()
with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()):
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(),
)
assert response.choices[0].message.content == "hi"
assert logging_obj.pre_call.call_count == 1
def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch):
"""The suppression must not swallow the log on a request the gate declined,
so a deployment with no `rust` flag keeps exactly the log it always had."""
monkeypatch.setenv("LITELLM_RUST", "0")
def test_the_sync_python_path_logs_pre_call_once():
logging_obj = MagicMock()
response = _run(
logging_obj=logging_obj,
litellm_params={},
client=_sync_client_returning_converse_response(),
)
@ -441,83 +187,10 @@ def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch
assert logging_obj.pre_call.call_count == 1
def test_post_call_logging_fires_on_the_sync_rust_path():
"""The Rust core owns the provider call, so the Converse transform that
normally raises `post_call` never runs. Without the bridge hook every
post_call callback goes silent and `original_response` stays unset."""
import json
_inject()
logging_obj = MagicMock()
_run(logging_obj=logging_obj)
assert logging_obj.post_call.call_count == 1
logged = logging_obj.post_call.call_args.kwargs["original_response"]
assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust"
@pytest.mark.asyncio
async def test_post_call_logging_fires_on_the_async_rust_path():
"""The asynchronous path runs through the same hook, so the two paths
cannot drift apart the way the pre_call suppression once did."""
import json
async def native(**_kwargs):
return dict(RUST_RESPONSE)
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)
)
assert logging_obj.post_call.call_count == 1
logged = logging_obj.post_call.call_args.kwargs["original_response"]
assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust"
def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines():
"""A decline never reached the provider, so the Python path serves the
request and owns the only post_call. Firing the hook there too would double
every post_call callback for one request."""
class _Declined(Exception):
pass
class _FakeNative:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
def declining_native(**_kwargs):
raise _Declined("blank message text")
logging_obj, calls = _recording_logging_obj()
with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()):
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(),
)
assert response.choices[0].message.content == "hi"
assert len(calls["post_call"]) == 1
assert "hi" in calls["post_call"][0]["original_response"]
def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monkeypatch):
"""With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no
credentials at all. Preparing the Rust handoff must not dereference that
None: the bearer token signs the request on its own."""
monkeypatch.setenv("LITELLM_RUST", "0")
credentials at all. The handler must not dereference that None: the bearer
token signs the request on its own."""
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
client = _sync_client_returning_converse_response()
@ -528,26 +201,11 @@ def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monke
assert sent_headers["Authorization"] == "Bearer bedrock-bearer-token"
def test_the_rust_opt_in_needs_no_sigv4_principal():
"""The core resolves the bearer token itself, so a bearer-only deployment
keeps its opt-in and the gate sees no aws_* credential keys to sign with."""
seen = _inject()
response = _run(credentials=None, api_key="bedrock-bearer-token")
assert response.choices[0].message.content == "hello from rust"
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"
@pytest.mark.parametrize("configured_through", ["env_var", "api_key"])
def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, configured_through):
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
raises; a bearer-token deployment must still serve the request, since the
bearer token alone signs it."""
monkeypatch.setenv("LITELLM_RUST", "0")
if configured_through == "env_var":
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
else:
@ -569,7 +227,6 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co
def test_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch):
"""The tagged STS session signs the Converse call and the tags never reach the request body (#34069)."""
monkeypatch.setenv("LITELLM_RUST", "0")
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False)
monkeypatch.delenv("AWS_ROLE_ARN", raising=False)

14
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -2912,19 +2912,13 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h
assert "".join([chunk.choices[0].delta.content or "" for chunk in collected]) == "hi"
@pytest.mark.parametrize(
"custom_llm_provider, enabled, expected",
[("openai", True, True), ("openai", False, False), ("azure", True, False),
("hosted_vllm", True, False), (None, True, False)],
)
def test_the_rust_responses_websocket_needs_openai_and_process_enablement(
custom_llm_provider, enabled, expected, monkeypatch
):
@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure", "hosted_vllm", None])
def test_the_rust_responses_websocket_stays_on_python_with_the_switch_on(custom_llm_provider, monkeypatch):
from litellm.rust_bridge import configuration
configuration.reset_rust_configuration()
monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
assert _rust_responses_websocket_enabled(custom_llm_provider) is expected
monkeypatch.setenv("LITELLM_RUST", "1")
assert _rust_responses_websocket_enabled(custom_llm_provider) is False
def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch):

0
tests/test_litellm/messages/__init__.py Normal file
View file

287
tests/test_litellm/messages/test_dispatch.py Normal file
View file

@ -0,0 +1,287 @@
import inspect
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect
import pytest
import litellm
from litellm.llms.anthropic.experimental_pass_through.messages import handler as python_messages
from litellm.messages.dispatch import (
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
)
from litellm.rust_bridge import catalog
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Route, Rule, Rules
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.messages.entrypoints import (
NATIVE_AMESSAGES,
NATIVE_MESSAGES,
LiteLLMMessagesRequest,
NativeAmessages,
NativeMessages,
)
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
MESSAGES: Final = [{"role": "user", "content": "hi"}]
PYTHON_RULES: Final[Rules] = ()
RUST_RULES: Final[Rules] = (Rule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
def messages_binding(native: NativeMessages | None) -> NativeBinding[NativeMessages]:
binding: Final[NativeBinding[NativeMessages]] = NativeBinding(
"anthropic_messages_handler", validate=lambda _: None
)
binding.override(native)
return binding
def amessages_binding(native: NativeAmessages | None) -> NativeBinding[NativeAmessages]:
binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("anthropic_messages", validate=lambda _: None)
binding.override(native)
return binding
def response(model: str = "claude-sonnet-4-5") -> AnthropicMessagesResponse:
return AnthropicMessagesResponse(id="msg_test", type="message", role="assistant", model=model, content=[])
def test_public_signature_is_the_legacy_signature() -> None:
public_messages: Final = cast(Callable[..., object], litellm.anthropic_messages_handler)
legacy_messages: Final = cast(Callable[..., object], python_messages.anthropic_messages_handler)
public_amessages: Final = cast(Callable[..., object], litellm.anthropic_messages)
legacy_amessages: Final = cast(Callable[..., object], python_messages.anthropic_messages)
assert inspect.signature(public_messages) == inspect.signature(legacy_messages)
assert inspect.signature(public_amessages) == inspect.signature(legacy_amessages)
def test_python_route_forwards_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = (16, MESSAGES, "claude-sonnet-4-5")
kwargs: Final[Mapping[str, object]] = {"temperature": 0.1, "litellm_metadata": metadata}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape
captured.append((call_args, call_kwargs))
return expected
def native(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse:
pytest.fail("Python-only dispatch must not call native")
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=messages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=PYTHON_RULES,
)
assert result is expected
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[1] is MESSAGES
assert call_kwargs == kwargs
assert call_kwargs["litellm_metadata"] is metadata
assert kwargs == {"temperature": 0.1, "litellm_metadata": metadata}
@pytest.mark.asyncio
async def test_async_python_route_forwards_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = (16, MESSAGES, "claude-sonnet-4-5")
kwargs: Final[Mapping[str, object]] = {"temperature": 0.1, "litellm_metadata": metadata}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
async def python(
*call_args: object, **call_kwargs: object # kwargs-ok: records call shape
) -> AnthropicMessagesResponse:
captured.append((call_args, call_kwargs))
return expected
async def native(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse:
pytest.fail("Python-only dispatch must not call native")
result: Final = await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=amessages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=PYTHON_RULES,
)
assert result is expected
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[1] is MESSAGES
assert call_kwargs == kwargs
assert call_kwargs["litellm_metadata"] is metadata
assert kwargs == {"temperature": 0.1, "litellm_metadata": metadata}
def test_native_receives_normalized_request_and_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = (16, MESSAGES, "anthropic/claude-sonnet-4-5")
kwargs: Final[Mapping[str, object]] = {
"stream": True,
"api_key": "sk-test",
"api_base": "https://example.invalid",
"custom_llm_provider": "anthropic",
"litellm_metadata": metadata,
}
captured: Final[list[tuple[LiteLLMMessagesRequest, tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response("anthropic/claude-sonnet-4-5")
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback
pytest.fail("Required Rust dispatch must not call Python")
def native(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse:
captured.append((request, args, kwargs))
return expected
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=messages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
assert result is expected
request, call_args, call_kwargs = captured[0]
assert request.model == "anthropic/claude-sonnet-4-5"
assert request.messages is MESSAGES
assert request.max_tokens == 16
assert request.stream is True
assert request.api_key == "sk-test"
assert request.api_base == "https://example.invalid"
assert request.custom_llm_provider == "anthropic"
assert request.kwargs == {"litellm_metadata": metadata}
assert request.kwargs["litellm_metadata"] is metadata
assert call_args == args
assert call_args[1] is MESSAGES
assert call_kwargs == kwargs
assert call_kwargs["litellm_metadata"] is metadata
def test_internal_async_marker_bypasses_native() -> None:
args: Final[tuple[object, ...]] = (16, MESSAGES, "claude-sonnet-4-5")
kwargs: Final[Mapping[str, object]] = {"is_async": True}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape
captured.append((call_args, call_kwargs))
return expected
def native(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse:
pytest.fail("The async handler's inner sync call must stay on Python")
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=messages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
assert result is expected
assert captured == [(args, kwargs)]
@pytest.mark.parametrize(
("args", "kwargs"),
(
((16, MESSAGES, "claude-sonnet-4-5"), {"model": "duplicate"}),
((), {}),
),
)
def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Mapping[str, object]) -> None:
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records invalid call
captured.append((call_args, call_kwargs))
return expected
def native(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse:
pytest.fail("Binding failures must be delegated to Python")
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=messages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
assert result is expected
assert captured == [(args, kwargs)]
def test_anthropic_create_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
captured: Final[list[LiteLLMMessagesRequest]] = []
expected: Final = response()
def native(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse:
captured.append(request)
return expected
NATIVE_MESSAGES.override(native)
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
public_create: Final = cast(Callable[..., AnthropicMessagesResponse], litellm.anthropic.create)
try:
result: Final = public_create(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5")
finally:
NATIVE_MESSAGES.reset()
assert result is expected
assert [request.model for request in captured] == ["claude-sonnet-4-5"]
@pytest.mark.asyncio
async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
captured: Final[list[LiteLLMMessagesRequest]] = []
expected: Final = response()
async def native(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse:
captured.append(request)
return expected
NATIVE_AMESSAGES.override(native)
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
public_acreate: Final = cast(Callable[..., Awaitable[AnthropicMessagesResponse]], litellm.anthropic.acreate)
try:
result: Final = await public_acreate(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5")
finally:
NATIVE_AMESSAGES.reset()
assert result is expected
assert [request.model for request in captured] == ["claude-sonnet-4-5"]

389
tests/test_litellm/ocr/test_dispatch.py Normal file
View file

@ -0,0 +1,389 @@
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect
import httpx
import pytest
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr.dispatch import (
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
)
from litellm.rust_bridge import catalog
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Route, Rule, Rules
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.ocr.entrypoints import (
NATIVE_AOCR,
NATIVE_OCR,
LiteLLMOcrRequest,
NativeAocr,
NativeOcr,
)
PYTHON_RULES: Final[Rules] = (Rule(Route.OCR, Rollout.PYTHON_ONLY),)
RUST_RULES: Final[Rules] = (Rule(Route.OCR, Rollout.RUST_REQUIRED),)
def ocr_binding(native: NativeOcr | None) -> NativeBinding[NativeOcr]:
binding: Final[NativeBinding[NativeOcr]] = NativeBinding("ocr", validate=lambda _: None)
binding.override(native)
return binding
def aocr_binding(native: NativeAocr | None) -> NativeBinding[NativeAocr]:
binding: Final[NativeBinding[NativeAocr]] = NativeBinding("aocr", validate=lambda _: None)
binding.override(native)
return binding
def response(model: str = "mistral/mistral-ocr-latest") -> OCRResponse:
return OCRResponse(pages=[], model=model)
def test_python_route_forwards_original_call_shape() -> None:
document: Final[Mapping[str, object]] = {
"type": "document_url",
"document_url": "https://example.invalid/document.pdf",
}
pages: Final = [0]
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document)
kwargs: Final[Mapping[str, object]] = {"pages": pages}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> OCRResponse: # kwargs-ok: records public call shape
captured.append((call_args, call_kwargs))
return expected
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
pytest.fail("Python-only dispatch must not call native")
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=ocr_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=PYTHON_RULES,
)
assert result is expected
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[1] is document
assert call_kwargs == kwargs
assert call_kwargs["pages"] is pages
assert kwargs == {"pages": pages}
@pytest.mark.asyncio
async def test_async_python_route_forwards_original_call_shape() -> None:
document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"}
pages: Final = [1]
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document)
kwargs: Final[Mapping[str, object]] = {"pages": pages}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
async def python(
*call_args: object,
**call_kwargs: object, # kwargs-ok: records public call shape
) -> OCRResponse:
captured.append((call_args, call_kwargs))
return expected
async def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
pytest.fail("Python-only dispatch must not call native")
result: Final = await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=aocr_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=PYTHON_RULES,
)
assert result is expected
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[1] is document
assert call_kwargs == kwargs
assert call_kwargs["pages"] is pages
assert kwargs == {"pages": pages}
def test_native_receives_normalized_positional_request_and_original_call_shape() -> None:
document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"}
timeout: Final = httpx.Timeout(30)
extra_headers: Final[dict[str, object]] = {"x-test": "1"}
pages: Final = [0, 2]
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document)
kwargs: Final[Mapping[str, object]] = {
"api_key": "test-key",
"api_base": "https://example.invalid",
"timeout": timeout,
"custom_llm_provider": "mistral",
"extra_headers": extra_headers,
"pages": pages,
}
captured: Final[list[tuple[LiteLLMOcrRequest, tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> OCRResponse: # kwargs-ok: rejected fallback
pytest.fail("Required Rust dispatch must not call Python")
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
captured.append((request, args, kwargs))
return expected
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=ocr_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
request, call_args, call_kwargs = captured[0]
assert result is expected
assert request.model == "mistral/mistral-ocr-latest"
assert request.document is document
assert request.api_key == "test-key"
assert request.api_base == "https://example.invalid"
assert request.timeout is timeout
assert request.custom_llm_provider == "mistral"
assert request.extra_headers is extra_headers
assert request.kwargs == {"pages": pages}
assert request.kwargs["pages"] is pages
assert call_args is args
assert call_kwargs is kwargs
def test_native_preserves_keyword_model_and_document_in_original_call_shape() -> None:
document: Final[Mapping[str, object]] = {
"type": "document_url",
"document_url": "https://example.invalid/document.pdf",
}
pages: Final = [1]
args: Final[tuple[object, ...]] = ()
kwargs: Final[Mapping[str, object]] = {
"model": "mistral/mistral-ocr-latest",
"document": document,
"pages": pages,
}
captured: Final[list[tuple[LiteLLMOcrRequest, tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> OCRResponse: # kwargs-ok: rejected fallback
pytest.fail("Required Rust dispatch must not call Python")
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
captured.append((request, args, kwargs))
return expected
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=ocr_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
request, call_args, call_kwargs = captured[0]
assert result is expected
assert request.model == "mistral/mistral-ocr-latest"
assert request.document is document
assert request.kwargs == {"pages": pages}
assert call_args is args
assert call_kwargs is kwargs
assert call_kwargs["model"] == "mistral/mistral-ocr-latest"
assert call_kwargs["document"] is document
def test_aocr_marker_bypasses_native() -> None:
document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"}
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document)
kwargs: Final[Mapping[str, object]] = {"aocr": True}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> OCRResponse: # kwargs-ok: records public call shape
captured.append((call_args, call_kwargs))
return expected
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
pytest.fail("aocr's inner ocr call must stay on Python")
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=ocr_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
assert result is expected
assert captured == [(args, kwargs)]
@pytest.mark.parametrize(
("args", "kwargs", "message"),
(
(
("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}),
{"model": "duplicate"},
r"ocr\(\) got multiple values for argument 'model'",
),
(
("mistral/mistral-ocr-latest",),
{},
r"ocr\(\) missing 1 required positional argument: 'document'",
),
),
)
def test_ocr_parser_errors_before_python_or_native(
args: tuple[object, ...], kwargs: Mapping[str, object], message: str
) -> None:
def python(*call_args: object, **call_kwargs: object) -> OCRResponse: # kwargs-ok: rejects parser failures
pytest.fail("OCR parser failures must not call Python")
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
pytest.fail("OCR parser failures must not call native")
with pytest.raises(TypeError, match=message):
_DISPATCH.run(
args,
kwargs,
python=python,
binding=ocr_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("args", "kwargs", "message"),
(
(
("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}),
{"model": "duplicate"},
r"aocr\(\) got multiple values for argument 'model'",
),
(
("mistral/mistral-ocr-latest",),
{},
r"aocr\(\) missing 1 required positional argument: 'document'",
),
),
)
async def test_aocr_parser_errors_before_python_or_native(
args: tuple[object, ...], kwargs: Mapping[str, object], message: str
) -> None:
async def python(
*call_args: object,
**call_kwargs: object, # kwargs-ok: rejects parser failures
) -> OCRResponse:
pytest.fail("OCR parser failures must not call Python")
async def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
pytest.fail("OCR parser failures must not call native")
with pytest.raises(TypeError, match=message):
await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=aocr_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
def test_public_ocr_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
document: Final[Mapping[str, object]] = {
"type": "document_url",
"document_url": "https://example.invalid/document.pdf",
}
captured: Final[list[LiteLLMOcrRequest]] = []
expected: Final = response()
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
captured.append(request)
return expected
NATIVE_OCR.override(native)
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
public_ocr: Final = cast(Callable[..., OCRResponse], litellm.ocr)
try:
result: Final = public_ocr(model="mistral/mistral-ocr-latest", document=document)
finally:
NATIVE_OCR.reset()
assert result is expected
assert [request.model for request in captured] == ["mistral/mistral-ocr-latest"]
@pytest.mark.asyncio
async def test_public_aocr_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
document: Final[Mapping[str, object]] = {
"type": "document_url",
"document_url": "https://example.invalid/document.pdf",
}
captured: Final[list[LiteLLMOcrRequest]] = []
expected: Final = response()
async def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> OCRResponse:
captured.append(request)
return expected
NATIVE_AOCR.override(native)
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
public_aocr: Final = cast(Callable[..., Awaitable[OCRResponse]], litellm.aocr)
try:
result: Final = await public_aocr(model="mistral/mistral-ocr-latest", document=document)
finally:
NATIVE_AOCR.reset()
assert result is expected
assert [request.model for request in captured] == ["mistral/mistral-ocr-latest"]

83
tests/test_litellm/ocr/test_legacy.py → tests/test_litellm/ocr/test_main.py
View file

@ -1,4 +1,3 @@
import importlib
from collections.abc import AsyncGenerator
from datetime import datetime
from io import BytesIO
@ -15,9 +14,9 @@ from litellm.litellm_core_utils.litellm_logging import Logging, use_custom_prici
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
from litellm.llms.custom_httpx import llm_http_handler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.ocr.legacy import _prepare_ocr_request
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE
from litellm.ocr.main import _prepare_ocr_request
from litellm.rust_bridge import bindings, configuration, runtime
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
@pytest.fixture
@ -45,7 +44,8 @@ async def provider(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[Mock]:
monkeypatch.setattr(llm_http_handler, "_get_httpx_client", lambda: sync_handler)
monkeypatch.setattr(llm_http_handler, "get_async_httpx_client", lambda llm_provider: async_handler)
yield handler
NATIVE_OCR_LIFECYCLE.reset()
NATIVE_OCR.reset()
NATIVE_AOCR.reset()
configuration.reset_rust_configuration()
@ -60,9 +60,9 @@ async def test_python_request_response_and_callbacks(
if dispatch != "disabled":
monkeypatch.setenv("LITELLM_RUST", "1")
NATIVE_OCR_LIFECYCLE.override(Mock(side_effect=Declined()) if dispatch == "declined" else None)
main: Final = importlib.import_module("litellm.ocr.main")
monkeypatch.setattr(main, "native_exception_types", lambda: (Declined, RuntimeError))
binding: Final = NATIVE_AOCR if mode == "async" else NATIVE_OCR
binding.override(Mock(side_effect=Declined()) if dispatch == "declined" else None)
monkeypatch.setattr(runtime, "native_exception_types", lambda: (Declined, RuntimeError))
logger: Final = Mock(spec=CustomLogger)
monkeypatch.setattr(litellm, "input_callback", [logger])
arguments: Final = {
@ -257,3 +257,70 @@ def test_direct_ocr_call_bills_request_level_per_page_pricing() -> None:
)
assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.05 * 3)
def _prepare(model: str, document: object, **kwargs: object) -> object:
return _prepare_ocr_request(
model=model,
document=document, # pyright: ignore[reportArgumentType] # exercises the runtime guard for untyped callers
api_key="test-key",
api_base=None,
timeout=None,
custom_llm_provider=None,
extra_headers=None,
kwargs={"litellm_logging_obj": Mock(), **kwargs},
)
@pytest.mark.parametrize(
("document", "match"),
(
("https://example.com/file.pdf", "document must be a dict"),
({"type": "video_url", "video_url": "https://example.com/clip.mp4"}, "Invalid document type: video_url"),
),
)
def test_prepare_ocr_request_rejects_malformed_documents(document: object, match: str) -> None:
with pytest.raises(ValueError, match=match):
_prepare("mistral/mistral-ocr-latest", document)
def test_prepare_ocr_request_rejects_provider_without_ocr_support() -> None:
with pytest.raises(ValueError, match="OCR is not supported for provider: openai"):
_prepare("openai/gpt-4o", dict(PRICING_DOCUMENT))
@pytest.mark.parametrize(
("request_format", "match"),
(("markdown", "Invalid `req_format`"), ("native", "`req_format='native'` is not supported")),
)
def test_prepare_ocr_request_rejects_unsupported_request_format(request_format: str, match: str) -> None:
with pytest.raises(litellm.UnsupportedParamsError, match=match):
_prepare("mistral/mistral-ocr-latest", dict(PRICING_DOCUMENT), req_format=request_format)
@pytest.mark.asyncio
async def test_python_none_provider_response_raises_public_error(
provider: Mock, monkeypatch: pytest.MonkeyPatch
) -> None:
from litellm.ocr import main
monkeypatch.setattr(main.base_llm_http_handler, "ocr", Mock(return_value=None))
with pytest.raises(litellm.APIConnectionError, match="unexpected None response") as error:
await litellm.aocr(model="mistral/mistral-ocr-latest", document=dict(PRICING_DOCUMENT), api_key="test-key")
assert error.value.llm_provider == "mistral"
assert provider.call_count == 0
@pytest.mark.parametrize(
("model", "expected_provider"),
(("mistral-ocr-latest", "mistral"), ("azure_ai/doc-intelligence/prebuilt-layout", "azure_ai")),
)
def test_preparation_errors_map_to_public_exception_for_inferred_provider(
provider: Mock, model: str, expected_provider: str
) -> None:
with pytest.raises(litellm.APIConnectionError) as error:
litellm.ocr(model=model, document="not-a-document") # pyright: ignore[reportArgumentType] # exercises the runtime guard
assert error.value.llm_provider == expected_provider
assert "document must be a dict" in str(error.value)
assert provider.call_count == 0

2
tests/test_litellm/ocr/test_ocr_file_input.py
View file

@ -21,7 +21,7 @@ import orjson
import pytest
from starlette.datastructures import FormData
from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type
class TestGetMimeType:

322
tests/test_litellm/responses/test_dispatch.py Normal file
View file

@ -0,0 +1,322 @@
import inspect
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect
import pytest
import litellm
from litellm.responses import dispatch as responses_dispatch
from litellm.responses import main as python_responses
from litellm.responses.dispatch import (
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
)
from litellm.rust_bridge import catalog
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Route, Rule
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.responses.entrypoints import (
NATIVE_ARESPONSES,
NATIVE_RESPONSES,
LiteLLMResponsesRequest,
NativeAresponses,
NativeResponses,
)
from litellm.types.llms.openai import ResponsesAPIResponse
INPUT: Final = [{"role": "user", "content": "hi"}]
PYTHON_RULES: Final = ()
RUST_RULES: Final = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED),)
def _response(model: str = "gpt-4o") -> ResponsesAPIResponse:
return ResponsesAPIResponse(
id="resp_test", object="response", created_at=0, model=model, output=[], status="completed"
)
def responses_binding(native: NativeResponses | None) -> NativeBinding[NativeResponses]:
binding: Final[NativeBinding[NativeResponses]] = NativeBinding("responses", validate=lambda _: None)
binding.override(native)
return binding
def aresponses_binding(native: NativeAresponses | None) -> NativeBinding[NativeAresponses]:
binding: Final[NativeBinding[NativeAresponses]] = NativeBinding("aresponses", validate=lambda _: None)
binding.override(native)
return binding
def test_public_signature_is_the_legacy_signature() -> None:
public_responses: Final = cast(Callable[..., object], litellm.responses)
legacy_responses: Final = cast(Callable[..., object], python_responses.responses)
public_aresponses: Final = cast(Callable[..., object], litellm.aresponses)
legacy_aresponses: Final = cast(Callable[..., object], python_responses.aresponses)
assert inspect.signature(public_responses) == inspect.signature(legacy_responses)
assert inspect.signature(public_aresponses) == inspect.signature(legacy_aresponses)
def test_python_route_forwards_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = (INPUT, "gpt-4o")
kwargs: Final[Mapping[str, object]] = {"temperature": 0.1, "litellm_metadata": metadata}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
response: Final = _response()
def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: records call shape
captured.append((call_args, call_kwargs))
return response
def native(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse:
pytest.fail("Python-only dispatch must not call native")
assert (
_DISPATCH.run(
args,
kwargs,
python=python,
binding=responses_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=PYTHON_RULES,
)
is response
)
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[0] is INPUT
assert call_kwargs == kwargs
assert call_kwargs["litellm_metadata"] is metadata
assert kwargs == {"temperature": 0.1, "litellm_metadata": metadata}
@pytest.mark.asyncio
async def test_async_python_route_forwards_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = (INPUT, "gpt-4o")
kwargs: Final[Mapping[str, object]] = {"temperature": 0.1, "litellm_metadata": metadata}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
response: Final = _response()
async def python(
*call_args: object, **call_kwargs: object # kwargs-ok: records call shape
) -> ResponsesAPIResponse:
captured.append((call_args, call_kwargs))
return response
async def native(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse:
pytest.fail("Python-only dispatch must not call native")
result: Final = await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=aresponses_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=PYTHON_RULES,
)
assert result is response
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[0] is INPUT
assert call_kwargs == kwargs
assert call_kwargs["litellm_metadata"] is metadata
assert kwargs == {"temperature": 0.1, "litellm_metadata": metadata}
def test_native_receives_normalized_request_and_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
extra_headers: Final = {"x-test": "1"}
args: Final[tuple[object, ...]] = (INPUT, "anthropic/claude-sonnet-4-5")
kwargs: Final[Mapping[str, object]] = {
"stream": True,
"api_key": "sk-test",
"base_url": "https://example.invalid",
"extra_headers": extra_headers,
"custom_llm_provider": "anthropic",
"litellm_metadata": metadata,
}
captured: Final[
list[tuple[LiteLLMResponsesRequest, tuple[object, ...], Mapping[str, object]]]
] = []
response: Final = _response("anthropic/claude-sonnet-4-5")
def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: rejected fallback
pytest.fail("Required Rust dispatch must not call Python")
def native(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse:
captured.append((request, args, kwargs))
return response
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=responses_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
request, call_args, call_kwargs = captured[0]
assert result is response
assert request.model == "anthropic/claude-sonnet-4-5"
assert request.input is INPUT
assert request.stream is True
assert request.api_key == "sk-test"
assert request.api_base == "https://example.invalid"
assert request.custom_llm_provider == "anthropic"
assert request.extra_headers is extra_headers
assert request.kwargs == {
"api_key": "sk-test",
"base_url": "https://example.invalid",
"litellm_metadata": metadata,
}
assert request.kwargs["litellm_metadata"] is metadata
assert call_args == args
assert call_args[0] is INPUT
assert call_kwargs == kwargs
assert call_kwargs["extra_headers"] is extra_headers
assert call_kwargs["litellm_metadata"] is metadata
def test_internal_async_marker_bypasses_native() -> None:
args: Final[tuple[object, ...]] = (INPUT, "gpt-4o")
kwargs: Final[Mapping[str, object]] = {"aresponses": True}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
response: Final = _response()
def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: records call shape
captured.append((call_args, call_kwargs))
return response
def native(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse:
pytest.fail("aresponses' inner responses call must stay on Python")
assert (
_DISPATCH.run(
args,
kwargs,
python=python,
binding=responses_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
is response
)
assert captured == [(args, kwargs)]
@pytest.mark.parametrize(
("args", "kwargs"),
(
((INPUT, "gpt-4o"), {"model": "duplicate"}),
((), {}),
),
)
def test_binding_errors_delegate_unchanged_to_python(
args: tuple[object, ...], kwargs: Mapping[str, object]
) -> None:
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
response: Final = _response()
def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: records invalid call
captured.append((call_args, call_kwargs))
return response
def native(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse:
pytest.fail("Binding failures must be delegated to Python")
assert (
_DISPATCH.run(
args,
kwargs,
python=python,
binding=responses_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
rules=RUST_RULES,
)
is response
)
assert captured == [(args, kwargs)]
def test_public_responses_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
captured: Final[list[LiteLLMResponsesRequest]] = []
expected: Final = _response()
def native(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse:
captured.append(request)
return expected
NATIVE_RESPONSES.override(native)
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
public_responses: Final = cast(Callable[..., ResponsesAPIResponse], litellm.responses)
try:
result: Final = public_responses(input=INPUT, model="gpt-4o")
finally:
NATIVE_RESPONSES.reset()
assert result is expected
assert [request.model for request in captured] == ["gpt-4o"]
@pytest.mark.asyncio
async def test_public_aresponses_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
captured: Final[list[LiteLLMResponsesRequest]] = []
expected: Final = _response()
async def native(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse:
captured.append(request)
return expected
NATIVE_ARESPONSES.override(native)
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
public_aresponses: Final = cast(Callable[..., Awaitable[ResponsesAPIResponse]], litellm.aresponses)
try:
result: Final = await public_aresponses(input=INPUT, model="gpt-4o")
finally:
NATIVE_ARESPONSES.reset()
assert result is expected
assert [request.model for request in captured] == ["gpt-4o"]
def test_responses_with_retries_uses_the_dispatch_entrypoint(monkeypatch: pytest.MonkeyPatch) -> None:
calls: Final[list[Mapping[str, object]]] = []
expected: Final = _response()
def dispatch_responses(*args: object, **kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: records call shape
calls.append(kwargs)
return expected
monkeypatch.setattr(responses_dispatch, "responses", dispatch_responses)
retry: Final = cast(Callable[..., ResponsesAPIResponse], litellm.responses_with_retries)
result: Final = retry(input=INPUT, model="gpt-4o", num_retries=1)
assert result is expected
assert calls[0]["num_retries"] == 0
assert calls[0]["max_retries"] == 0

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

@ -2,8 +2,8 @@ from __future__ import annotations
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 import configuration
from litellm.rust_bridge.responses import websocket as responses_websocket
class _FakeNativeConnection:
@ -47,14 +47,6 @@ def reset_responses_websocket():
configuration.reset_rust_configuration()
def test_rust_websocket_bridge_uses_process_enablement() -> None:
configuration.rust(False)
assert not _rust_responses_websocket_enabled("openai")
configuration.rust(True)
assert _rust_responses_websocket_enabled("openai")
assert not _rust_responses_websocket_enabled("anthropic")
@pytest.mark.asyncio
async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None:
adapter = responses_websocket._ConnectionAdapter(_ClosedNativeConnection())

0
tests/test_litellm/rust_bridge/__init__.py Normal file
View file

0
tests/test_litellm/rust_bridge/chat_completions/__init__.py Normal file
View file

49
tests/test_litellm/rust_bridge/chat_completions/test_callbacks.py Normal file
View file

@ -0,0 +1,49 @@
from types import MappingProxyType
from typing import Final
from litellm.rust_bridge.chat_completions.callbacks import arguments, response
from litellm.rust_bridge.chat_completions.entrypoints import LiteLLMChatCompletionsRequest
from litellm.types.utils import ModelResponse
def test_response_builds_the_public_model_response() -> None:
built: Final = response(
MappingProxyType(
{
"id": "chatcmpl-native",
"object": "chat.completion",
"created": 1,
"model": "claude-sonnet-4-5",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {"role": "assistant", "content": "native"},
}
],
"usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
}
)
)
assert isinstance(built, ModelResponse)
assert built.id == "chatcmpl-native"
assert built.choices[0].message.content == "native"
assert built.usage is not None
assert built.usage.total_tokens == 5
def test_arguments_are_the_public_kwargs_view() -> None:
kwargs: Final = MappingProxyType({"metadata": {"user_id": "u"}})
request: Final = LiteLLMChatCompletionsRequest(
model="anthropic/claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
stream=None,
api_key=None,
api_base=None,
custom_llm_provider="anthropic",
extra_headers=None,
kwargs=kwargs,
)
assert arguments(request) is kwargs

0
tests/test_litellm/rust_bridge/messages/__init__.py Normal file
View file

42
tests/test_litellm/rust_bridge/messages/test_callbacks.py Normal file
View file

@ -0,0 +1,42 @@
from types import MappingProxyType
from typing import Final
from litellm.rust_bridge.messages.callbacks import arguments, response
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
def test_response_is_a_detached_public_messages_dict() -> None:
native: Final = MappingProxyType(
{
"id": "msg_native",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [{"type": "text", "text": "native"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 2, "output_tokens": 3},
}
)
built: Final = response(native)
assert built == dict(native)
assert isinstance(built, dict)
built["_hidden_params"] = {"annotated": True}
assert "_hidden_params" not in native
def test_arguments_are_the_public_kwargs_view() -> None:
kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}})
request: Final = LiteLLMMessagesRequest(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
max_tokens=16,
stream=None,
api_key=None,
api_base=None,
custom_llm_provider="anthropic",
kwargs=kwargs,
)
assert arguments(request) is kwargs

78
tests/test_litellm/rust_bridge/native_route_wheel_test.py
View file

@ -73,32 +73,12 @@ def assert_native_request(
headers: HTTPMessage,
body: object,
) -> None:
if route not in {"ocr", "azure_ocr", "azure_di", "transcription", "messages", "chat_completions"}:
if route not in {"transcription", "messages", "chat_completions"}:
raise AssertionError(f"unexpected route marker: {route!r}")
if outcome not in {"success", "429", "hang"}:
raise AssertionError(f"unexpected outcome marker: {outcome!r}")
if not isinstance(body, dict):
raise TypeError(f"{route} sent {type(body).__name__}, expected a JSON object")
if route == "ocr":
assert path == "/v1/ocr"
assert headers.get("authorization") == "Bearer sk-native"
assert body["model"] == "mistral-ocr-latest"
assert body["document"]["document_url"] == "https://example.com/document.pdf"
assert body["include_image_base64"] is True
return
if route == "azure_ocr":
assert path == "/providers/mistral/azure/ocr"
assert headers.get("authorization") == "Bearer prepared-azure-token"
assert body["model"] == "mistral-ocr-2505"
assert body["document"]["document_url"] == "data:application/pdf;base64,YWJj"
return
if route == "azure_di":
assert path.startswith("/documentintelligence/documentModels/prebuilt-read:analyze?")
assert "api-version=2024-11-30" in path
assert "pages=1%2C3" in path
assert headers.get("ocp-apim-subscription-key") == "di-key"
assert body == {"base64Source": "YWJj"}
return
if route == "transcription":
assert path == "/model/mistral.voxtral-mini-3b-2507/converse"
assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ")
@ -120,10 +100,6 @@ def assert_native_request(
def native_response(status: int, route: str | None) -> bytes:
if status == 429:
return b'{"error":"native-rate-limit"}'
if route in {"ocr", "azure_ocr"}:
return b'{"pages":[{"index":0,"markdown":"native-ocr"}]}'
if route == "azure_di":
return b'{"status":"succeeded","analyzeResult":{"pages":[]}}'
if route == "transcription":
return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}'
return ANTHROPIC_RESPONSE
@ -144,14 +120,6 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
"extra_headers": {"x-test-outcome": outcome, "x-test-route": route},
"timeout_seconds": 3.0,
}
if route == "ocr":
return common | {
"model": "mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
"api_key": "sk-native",
"custom_llm_provider": "mistral",
"optional_params": {"include_image_base64": True},
}
if route == "transcription":
return common | {
"model": "mistral.voxtral-mini-3b-2507",
@ -189,42 +157,12 @@ def assert_success(route: str, response: object) -> None:
if not isinstance(response, dict):
raise TypeError(f"{route} returned {type(response).__name__}, expected dict")
actual: Final = success_value(route, response)
expected: Final = (
"native-ocr" if route == "ocr" else "native-transcription" if route == "transcription" else "native-message"
)
expected: Final = "native-transcription" if route == "transcription" else "native-message"
if actual != expected:
raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}")
def azure_ocr_kwargs(api_base: str) -> dict[str, object]:
return {
"model": "mistral-ocr-2505",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"api_base": api_base,
"custom_llm_provider": "azure_ai",
"extra_headers": {
"x-test-outcome": "success",
"x-test-route": "azure_ocr",
},
"optional_params": {"azure_ad_token": "prepared-azure-token"},
}
def azure_di_kwargs(api_base: str) -> dict[str, object]:
return {
"model": "doc-intelligence/prebuilt-read",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"api_key": "di-key",
"api_base": api_base,
"custom_llm_provider": "azure_ai",
"extra_headers": {"x-test-outcome": "success", "x-test-route": "azure_di"},
"optional_params": {"req_format": "native", "pages": [0, 2]},
}
def success_value(route: str, response: dict[object, object]) -> object:
if route == "ocr":
return response["pages"][0]["markdown"]
if route == "transcription":
return response["text"]
if route == "messages":
@ -233,7 +171,7 @@ def success_value(route: str, response: dict[object, object]) -> object:
def assert_rate_limit(native: object, route: str, error: BaseException) -> None:
if route in {"ocr", "chat_completions"}:
if route == "chat_completions":
upstream_error: Final = native.RustUpstreamError
if not isinstance(error, upstream_error) or error.args[0] != 429:
raise AssertionError(f"{route} returned the wrong 429 error: {error!r}")
@ -243,7 +181,7 @@ def assert_rate_limit(native: object, route: str, error: BaseException) -> None:
def exercise_sync(native: object, api_base: str) -> None:
for route in ("ocr", "transcription", "messages", "chat_completions"):
for route in ("transcription", "messages", "chat_completions"):
function: Final = getattr(native, route)
assert_success(route, function(**route_kwargs(route, api_base, "success")))
try:
@ -252,13 +190,10 @@ def exercise_sync(native: object, api_base: str) -> None:
assert_rate_limit(native, route, error)
else:
raise AssertionError(f"{route} accepted a 429 response")
assert_success("ocr", native.ocr(**azure_ocr_kwargs(api_base)))
di_response: Final = native.ocr(**azure_di_kwargs(api_base))
assert di_response["provider_native_response"]["status"] == "succeeded"
async def exercise_async(native: object, api_base: str) -> None:
for route in ("ocr", "transcription", "messages", "chat_completions"):
for route in ("transcription", "messages", "chat_completions"):
function: Final = getattr(native, f"a{route}")
assert_success(route, await function(**route_kwargs(route, api_base, "success")))
try:
@ -267,9 +202,6 @@ async def exercise_async(native: object, api_base: str) -> None:
assert_rate_limit(native, route, error)
else:
raise AssertionError(f"a{route} accepted a 429 response")
assert_success("ocr", await native.aocr(**azure_ocr_kwargs(api_base)))
di_response: Final = await native.aocr(**azure_di_kwargs(api_base))
assert di_response["provider_native_response"]["status"] == "succeeded"
async def exercise_async_concurrency(native: object, api_base: str) -> None:

0
tests/test_litellm/rust_bridge/ocr/__init__.py Normal file
View file

8
tests/test_litellm/ocr/test_ocr_native_format.py → tests/test_litellm/rust_bridge/ocr/test_callbacks.py
View file

@ -1,13 +1,9 @@
"""
Tests for the OCR `req_format` option in the SDK request path.
"""
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.ocr.callbacks import response as build_ocr_response
def test_rust_ocr_response_retains_provider_native_response():
provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}}
response = rust_ocr_bridge._response(
response = build_ocr_response(
{
"pages": [],
"model": "prebuilt-layout",

0
tests/test_litellm/rust_bridge/responses/__init__.py Normal file
View file

57
tests/test_litellm/rust_bridge/responses/test_callbacks.py Normal file
View file

@ -0,0 +1,57 @@
from types import MappingProxyType
from typing import Final
import pytest
from pydantic import ValidationError
from litellm.rust_bridge.responses.callbacks import arguments, response
from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest
from litellm.types.llms.openai import ResponsesAPIResponse
def test_response_validates_into_the_public_responses_model() -> None:
built: Final = response(
MappingProxyType(
{
"id": "resp_native",
"object": "response",
"created_at": 1,
"model": "gpt-4o",
"status": "completed",
"output": [
{
"type": "message",
"id": "msg_native",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "native", "annotations": []}],
}
],
}
)
)
assert isinstance(built, ResponsesAPIResponse)
assert built.id == "resp_native"
assert built.output[0].content[0].text == "native"
def test_response_rejects_a_payload_missing_required_fields() -> None:
with pytest.raises(ValidationError):
response(MappingProxyType({"object": "response"}))
def test_arguments_are_the_public_kwargs_view() -> None:
kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}})
request: Final = LiteLLMResponsesRequest(
model="gpt-4o",
input="hi",
stream=None,
api_key=None,
api_base=None,
custom_llm_provider="openai",
extra_headers=None,
kwargs=kwargs,
)
assert arguments(request) is kwargs

38
tests/test_litellm/rust_bridge/test_bindings.py
View file

@ -4,6 +4,11 @@ from typing import Final
import pytest
from litellm.rust_bridge import bindings
from litellm.rust_bridge.chat_completions import entrypoints as chat_completions
from litellm.rust_bridge.messages import entrypoints as messages
from litellm.rust_bridge.ocr import entrypoints as ocr
from litellm.rust_bridge.responses import entrypoints as responses
from litellm.rust_bridge.transcription import native as transcription
def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None:
@ -33,3 +38,36 @@ def test_binding_validates_native_attribute(
binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None)
assert binding.load() == expected
ROUTE_BINDINGS: Final = (
("completion", chat_completions.NATIVE_COMPLETION),
("acompletion", chat_completions.NATIVE_ACOMPLETION),
("anthropic_messages_handler", messages.NATIVE_MESSAGES),
("anthropic_messages", messages.NATIVE_AMESSAGES),
("responses", responses.NATIVE_RESPONSES),
("aresponses", responses.NATIVE_ARESPONSES),
("ocr", ocr.NATIVE_OCR),
("aocr", ocr.NATIVE_AOCR),
("transcription", transcription.NATIVE_TRANSCRIPTION),
("atranscription", transcription.NATIVE_ATRANSCRIPTION),
)
@pytest.mark.parametrize(
("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS]
)
def test_route_bindings_only_accept_callable_native_attributes(
monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object]
) -> None:
def native_route() -> None:
pass
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: "not callable"}))
route_binding.reset()
assert route_binding.load() is None
monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route}))
route_binding.reset()
assert route_binding.load() is native_route
route_binding.reset()

83
tests/test_litellm/rust_bridge/test_catalog.py Normal file
View file

@ -0,0 +1,83 @@
from __future__ import annotations
from collections.abc import Generator
from typing import Final
import pytest
from litellm.rust_bridge import catalog, configuration
from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule
from litellm.rust_bridge.configuration import Decision, Rollout
@pytest.fixture(autouse=True)
def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield
configuration.reset_rust_configuration()
@pytest.mark.parametrize("route", tuple(Route))
@pytest.mark.parametrize("provider", (None, "bedrock", "mistral", "anthropic", "openai", "azure_ai", "unknown"))
@pytest.mark.parametrize("delivery", tuple(Delivery))
@pytest.mark.parametrize("process", (None, False, True))
@pytest.mark.parametrize("environment", (None, "0", "1"))
def test_shipped_decisions(
monkeypatch: pytest.MonkeyPatch,
route: Route,
provider: str | None,
delivery: Delivery,
process: bool | None,
environment: str | None,
) -> None:
configuration.rust(process)
if environment is not None:
monkeypatch.setenv("LITELLM_RUST", environment)
context: Final = Context(route, provider=provider, model="test-model", delivery=delivery)
if route is Route.OCR:
enabled: Final = environment == "1" if environment is not None else process is not False
assert catalog.rollout(context) is Rollout.RUST_OPT_OUT
assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON)
elif route is Route.TRANSCRIPTION and provider == "bedrock":
assert catalog.rollout(context) is Rollout.RUST_REQUIRED
assert catalog.decision(context) is Decision.RUST_REQUIRED
else:
assert catalog.rollout(context) is Rollout.PYTHON_ONLY
assert catalog.decision(context) is Decision.PYTHON
@pytest.mark.parametrize("route", tuple(Route))
def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pytest.MonkeyPatch, route: Route) -> None:
configuration.rust(True)
monkeypatch.setenv("LITELLM_RUST", "1")
assert catalog.rollout(Context(route), rules=()) is Rollout.PYTHON_ONLY
assert catalog.decision(Context(route), rules=()) is Decision.PYTHON
@pytest.mark.parametrize(
("context", "expected"),
(
(Context(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), Decision.RUST_REQUIRED),
(Context(Route.RESPONSES, provider="openai", model="m"), Decision.PYTHON),
(Context(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.STREAMING), Decision.PYTHON),
(Context(Route.RESPONSES, provider="openai", model="other", delivery=Delivery.WEBSOCKET), Decision.PYTHON),
(Context(Route.RESPONSES, provider="anthropic", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON),
(Context(Route.MESSAGES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON),
),
)
def test_first_matching_rule_respects_every_constraint(context: Context, expected: Decision) -> None:
rules: Final = (
Rule(
Route.RESPONSES,
Rollout.RUST_REQUIRED,
providers=frozenset({"openai"}),
models=frozenset({"m"}),
deliveries=frozenset({Delivery.WEBSOCKET}),
),
Rule(Route.RESPONSES, Rollout.PYTHON_ONLY),
)
assert catalog.decision(context, rules) is expected

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

@ -1,395 +0,0 @@
"""Tests for the Rust chat completions bridge.
The native callables are dependency-injected through
``set_rust_chat_completions`` rather than patched, so these run without the
compiled extension present.
"""
from __future__ import annotations
import pytest
import litellm
from litellm.rust_bridge import configuration
from litellm.rust_bridge import chat_completions as bridge
from litellm.types.utils import ModelResponse
RUST_RESPONSE = {
"created": 1_700_000_000,
"model": "claude-sonnet-4-5-20260101",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hello from rust"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 11,
"completion_tokens": 4,
"total_tokens": 15,
"prompt_tokens_details": {
"cached_tokens": 0,
"cache_creation_tokens": 0,
"text_tokens": 11,
},
},
}
MESSAGES = [{"role": "user", "content": "hi"}]
class _FakeDeclined(Exception):
"""Stands in for the native `RustBridgeDeclined`."""
class _FakeUpstream(Exception):
"""Stands in for the native `RustUpstreamError`; args are (status, message)."""
class _FakeNative:
RustBridgeDeclined = _FakeDeclined
RustUpstreamError = _FakeUpstream
def _fake_native_bridge(monkeypatch):
"""Expose the bridge's exception classes without the compiled extension."""
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
def _hide_native_bridge(monkeypatch):
"""Simulate a wheel built without the compiled extension.
There is no injection seam for "the .so is absent", so the loader itself is
replaced; every other case here uses `set_rust_chat_completions`.
"""
monkeypatch.setattr(bridge, "get_native_bridge", lambda: None)
@pytest.fixture(autouse=True)
def reset_bridge(monkeypatch):
"""Every test starts with no injected callables, and leaves none behind."""
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
configuration.reset_rust_configuration()
monkeypatch.setenv("LITELLM_RUST", "1")
yield
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
configuration.reset_rust_configuration()
class _RecordingDecline:
"""A stand-in for the native gate that records what it was asked."""
def __init__(self, reason: str | None = None):
self.reason = reason
self.calls: list[dict] = []
def __call__(self, **kwargs):
self.calls.append(kwargs)
return self.reason
class _RecordingCall:
def __init__(self, result=None, error: Exception | None = None):
self.result = result if result is not None else dict(RUST_RESPONSE)
self.error = error
self.calls: list[dict] = []
def __call__(self, **kwargs):
self.calls.append(kwargs)
if self.error is not None:
raise self.error
return self.result
class _RecordingAsyncCall(_RecordingCall):
async def __call__(self, **kwargs):
return _RecordingCall.__call__(self, **kwargs)
def _accepts(**overrides) -> bool:
kwargs = {
"model": "claude-sonnet-4-5",
"messages": MESSAGES,
"optional_params": {"max_tokens": 16},
"custom_llm_provider": "anthropic",
"litellm_params": {},
"stream": None,
}
kwargs.update(overrides)
return bridge.rust_chat_completions_accepts(**kwargs)
class TestGate:
def test_declines_when_the_deployment_did_not_opt_in(self, monkeypatch):
monkeypatch.delenv("LITELLM_RUST", raising=False)
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
assert _accepts(litellm_params={}) is False
assert _accepts(litellm_params=None) is False
assert gate.calls == [], "the gate must not be consulted before opt-in"
def test_accepts_when_the_deployment_opted_in_and_the_core_agrees(self, monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
assert _accepts() is True
assert gate.calls[0]["model"] == "claude-sonnet-4-5"
assert gate.calls[0]["custom_llm_provider"] == "anthropic"
def test_process_enable_applies_without_request_override(self):
bridge.set_rust_chat_completions(decline=_RecordingDecline())
configuration.rust(True)
assert _accepts(litellm_params={}) is True
def test_the_env_var_opts_in_without_a_per_model_flag(self, monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "true")
bridge.set_rust_chat_completions(decline=_RecordingDecline())
assert _accepts(litellm_params={}) is True
def test_declines_streaming_and_providers_off_the_path(self, monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
assert _accepts(stream=True) is False
assert _accepts(custom_llm_provider="openai") is False
assert _accepts(custom_llm_provider=None) is False
assert gate.calls == []
def test_declines_an_anthropic_request_carrying_a_litellm_metadata_user_id(self, monkeypatch):
"""`AnthropicConfig.transform_request` copies a valid `user_id` into the Messages body.
It does that inside the function the Rust route replaces, and the core is
handed `optional_params` only, so accepting here would send the request
to Anthropic with the abuse-detection attribution silently missing.
"""
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False
assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of"
# Bedrock's Converse transform reads no `user_id`, and an Anthropic request
# whose metadata carries none is one Python would not attribute either.
assert (
_accepts(
custom_llm_provider="bedrock",
model="bedrock/us-east-1/anthropic.claude-v2",
litellm_params={"metadata": {"user_id": "u-123"}},
)
is True
)
assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True
assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True
assert _accepts(litellm_params={"metadata": None}) is True
def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch):
"""`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the
Converse body from `litellm_params`, and owning that field also means
evicting a caller-supplied one. The core can do neither, so an operator
who armed `bedrock_request_metadata_fields` keeps the Python path.
"""
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
bedrock = {
"custom_llm_provider": "bedrock",
"model": "bedrock/us-east-1/anthropic.claude-v2",
}
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ["user_api_key_team_id"])
assert _accepts(**bedrock) is False
assert gate.calls == [], "the core must not be consulted for a field it cannot write"
assert _accepts() is True, "arming Bedrock attribution must not decline Anthropic"
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None)
assert _accepts(**bedrock) is True, "the decline follows the operator's opt-in alone"
def test_declines_when_the_core_declines(self, monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
bridge.set_rust_chat_completions(decline=_RecordingDecline("streaming"))
assert _accepts() is False
def test_declines_when_the_bridge_is_unavailable(self, monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
_hide_native_bridge(monkeypatch)
assert _accepts() is False
def test_declines_when_the_gate_itself_raises(self, monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
def exploding(**_kwargs):
raise RuntimeError("boom")
bridge.set_rust_chat_completions(decline=exploding)
assert _accepts() is False
def _call_kwargs(model_response: ModelResponse) -> dict:
return {
"model": "claude-sonnet-4-5",
"messages": MESSAGES,
"optional_params": {"max_tokens": 16},
"model_response": model_response,
"api_key": "sk-test",
"api_base": None,
"custom_llm_provider": "anthropic",
"extra_headers": {},
"timeout": 30.0,
"on_response": lambda _rust_response: None,
}
class TestSyncCall:
def test_builds_a_model_response_and_stamps_the_rust_header(self):
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
model_response = ModelResponse()
original_id = model_response.id
result = bridge.chat_completions(**_call_kwargs(model_response))
assert result is not None
assert result.choices[0].message.content == "hello from rust"
assert result.choices[0].finish_reason == "stop"
assert result.model == "claude-sonnet-4-5-20260101"
assert result.usage.prompt_tokens == 11
assert result.usage.completion_tokens == 4
assert result.usage.total_tokens == 15
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
assert result.id == original_id, "the rust path must keep the chatcmpl id litellm already minted"
def test_passes_the_timeout_through_as_seconds(self):
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert native.calls[0]["timeout_seconds"] == 30.0
def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
def test_falls_back_when_the_core_declines_before_calling_the_provider(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
class TestAsyncCall:
@pytest.mark.asyncio
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 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):
_hide_native_bridge(monkeypatch)
assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None
@pytest.mark.asyncio
async def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")))
assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None
class TestAsyncFallbackWrapper:
@pytest.mark.asyncio
async def test_returns_the_rust_response_without_running_the_fallback(self):
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall())
ran = []
async def fallback():
ran.append(True)
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result.choices[0].message.content == "hello from rust"
assert ran == []
@pytest.mark.asyncio
async def test_runs_the_fallback_when_the_core_declines(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")))
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
@pytest.mark.asyncio
async def test_runs_the_fallback_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
class TestFailureClassification:
"""A failure the provider already saw must not be retried on the Python
path: it would bill the customer for the same work twice."""
@pytest.fixture(autouse=True)
def _native_exceptions(self, monkeypatch):
_fake_native_bridge(monkeypatch)
def test_a_decline_falls_back_because_nothing_was_sent(self):
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming")))
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
def test_an_upstream_failure_is_surfaced_with_its_status(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited")))
with pytest.raises(APIError) as raised:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert raised.value.status_code == 429
assert "rate limited" in str(raised.value)
def test_a_transport_failure_with_no_response_surfaces_as_a_500(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset")))
with pytest.raises(APIError) as raised:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert raised.value.status_code == 500
def test_an_unrecognized_error_is_not_swallowed(self):
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=RuntimeError("something else")))
with pytest.raises(RuntimeError):
bridge.chat_completions(**_call_kwargs(ModelResponse()))
@pytest.mark.asyncio
async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom")))
ran = []
async def fallback():
ran.append(True)
return "python"
with pytest.raises(APIError):
await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert ran == [], "a request the provider already served must not be re-issued"
@pytest.mark.asyncio
async def test_the_async_wrapper_falls_back_on_a_decline(self):
bridge.set_rust_chat_completions(
achat_completions=_RecordingAsyncCall(error=_FakeDeclined("blank message text"))
)
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"

113
tests/test_litellm/rust_bridge/test_configuration.py
View file

@ -22,88 +22,101 @@ def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest
configuration.reset_rust_configuration()
Rollout: Final = configuration.Rollout
Decision: Final = configuration.Decision
@pytest.mark.parametrize(
("process", "environment", "release_default", "expected"),
("rollout", "process", "environment", "expected"),
(
(False, True, True, False),
(True, False, False, True),
(None, False, True, False),
(None, True, False, True),
(None, None, False, False),
(None, None, True, True),
(Rollout.PYTHON_ONLY, True, True, Decision.PYTHON),
(Rollout.RUST_REQUIRED, False, False, Decision.RUST_REQUIRED),
(Rollout.RUST_OPT_IN, None, None, Decision.PYTHON),
(Rollout.RUST_OPT_IN, None, True, Decision.RUST_WITH_FALLBACK),
(Rollout.RUST_OPT_IN, True, None, Decision.RUST_WITH_FALLBACK),
(Rollout.RUST_OPT_IN, True, False, Decision.PYTHON),
(Rollout.RUST_OPT_IN, False, True, Decision.RUST_WITH_FALLBACK),
(Rollout.RUST_OPT_OUT, None, None, Decision.RUST_WITH_FALLBACK),
(Rollout.RUST_OPT_OUT, None, False, Decision.PYTHON),
(Rollout.RUST_OPT_OUT, False, None, Decision.PYTHON),
(Rollout.RUST_OPT_OUT, False, True, Decision.RUST_WITH_FALLBACK),
(Rollout.RUST_OPT_OUT, True, False, Decision.PYTHON),
),
)
def test_resolution_precedence(
def test_decide_precedence(
rollout: configuration.Rollout,
process: bool | None,
environment: bool | None,
release_default: bool,
expected: bool,
expected: configuration.Decision,
) -> None:
assert (
configuration.resolve_rust_enabled(
process_override=process,
environment_override=environment,
release_default=release_default,
)
is expected
)
assert configuration.decide(rollout, process_override=process, environment_override=environment) is expected
def test_release_default_remains_disabled() -> None:
assert configuration.DEFAULT_RUST_ENABLED is False
def test_release_default_keeps_opt_in_routes_on_python() -> None:
assert configuration.decision(Rollout.RUST_OPT_IN) is Decision.PYTHON
assert configuration.decision(Rollout.RUST_OPT_OUT) is Decision.RUST_WITH_FALLBACK
assert configuration.rust_enabled() is False
assert configuration.rust_ocr_enabled() is True
@pytest.mark.parametrize("process", [None, False, True])
@pytest.mark.parametrize("environment", [None, "0", "1", "off"])
def test_ocr_configuration(monkeypatch: pytest.MonkeyPatch, process: bool | None, environment: str | None) -> None:
@pytest.mark.parametrize("process", (None, False, True))
@pytest.mark.parametrize("environment", (None, "0", "1", "off"))
def test_opt_out_route_configuration(
monkeypatch: pytest.MonkeyPatch, process: bool | None, environment: str | None
) -> None:
if environment is not None:
monkeypatch.setenv("LITELLM_RUST", environment)
if process is not None:
configuration.rust(process)
assert configuration.rust_ocr_enabled() is (environment not in {"0", "off"} and process is not False)
expected: Final = (
Decision.RUST_WITH_FALLBACK
if environment == "1" or (environment is None and process is not False)
else Decision.PYTHON
)
assert configuration.decision(Rollout.RUST_OPT_OUT) is expected
def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
@pytest.mark.parametrize(
("environment", "process", "expected"),
(
*((value, True, False) for value in ("0", "false", "False", "no", "off", "f", "n", " 0 ")),
*((value, False, True) for value in ("1", "true", "TRUE", "yes", "on", "t", "y", " 1 ")),
),
)
def test_environment_wins_over_process_override(
monkeypatch: pytest.MonkeyPatch, environment: str, process: bool, expected: bool
) -> None:
monkeypatch.setenv("LITELLM_RUST", environment)
configuration.rust(process)
assert configuration.rust_enabled() is expected
def test_process_override_applies_when_environment_is_unset() -> None:
configuration.rust(True)
assert configuration.rust_enabled() is True
def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_RUST", "off")
assert configuration.rust_enabled() is False
@pytest.mark.parametrize("value", ("", " ", "sometimes", "2"))
def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
def test_invalid_environment_value_is_ignored(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
monkeypatch.setenv("LITELLM_RUST", value)
assert configuration.rust_enabled() is False
def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_RUST", "1")
with ThreadPoolExecutor(max_workers=1) as executor:
assert executor.submit(configuration.rust_enabled).result() is True
configuration.rust(False)
assert executor.submit(configuration.rust_enabled).result() is False
configuration.reset_rust_configuration()
assert executor.submit(configuration.rust_enabled).result() is True
def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_RUST", "sometimes")
assert configuration.decision(Rollout.RUST_OPT_OUT) is Decision.RUST_WITH_FALLBACK
configuration.rust(True)
assert configuration.rust_enabled() is True
def test_process_override_and_reset_apply_to_existing_threads() -> None:
with ThreadPoolExecutor(max_workers=1) as executor:
assert executor.submit(configuration.rust_enabled).result() is False
configuration.rust(True)
assert executor.submit(configuration.rust_enabled).result() is True
configuration.reset_rust_configuration()
assert executor.submit(configuration.rust_enabled).result() is False
@pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False")))
def test_environment_controls_startup(value: str, expected: str) -> None:
environment: Final = {**os.environ, "LITELLM_RUST": value}

229
tests/test_litellm/rust_bridge/test_dispatch.py Normal file
View file

@ -0,0 +1,229 @@
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping
from dataclasses import dataclass
from typing import Final
import pytest
from litellm.rust_bridge import configuration
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule, Rules
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.dispatch import PublicDispatch
@dataclass(frozen=True, slots=True)
class Request:
model: str
def binding() -> NativeBinding[object]:
bound: Final[NativeBinding[object]] = NativeBinding("unused", validate=lambda value: value)
bound.override(None)
return bound
def test_route_without_rules_forwards_before_request_projection() -> None:
stream: Final[Iterator[int]] = iter((1, 2))
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
pytest.fail("Python-only routes must not project the request")
dispatch: Final = PublicDispatch(
route=Route.CHAT_COMPLETIONS, request=reject_request, context=lambda _: Context(Route.CHAT_COMPLETIONS)
)
result: Final = dispatch.run(
("model",),
{"stream": True},
python=lambda *args, **kwargs: stream,
binding=binding(),
native=lambda hook, request, args, kwargs: pytest.fail("Python-only routes must not call native"),
rules=(),
)
assert result is stream
def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None:
rules: Final[Rules] = (
Rule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY),
Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),
)
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
pytest.fail("First-match Python rule must prevent request projection")
dispatch: Final = PublicDispatch(
route=Route.CHAT_COMPLETIONS,
request=reject_request,
context=lambda _: Context(Route.CHAT_COMPLETIONS),
)
expected: Final = object()
result: Final = dispatch.run(
("model",),
{},
python=lambda *args, **kwargs: expected,
binding=binding(),
native=lambda hook, request, args, kwargs: pytest.fail("First-match Python rule must prevent native"),
rules=rules,
)
assert result is expected
def test_disabled_optional_rust_rule_forwards_before_projection() -> None:
rules: Final[Rules] = (Rule(Route.OCR, Rollout.RUST_OPT_OUT),)
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
pytest.fail("Disabled optional Rust must not project the request")
dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: Context(Route.OCR))
expected: Final = object()
configuration.rust(False)
try:
result: Final = dispatch.run(
("model",),
{},
python=lambda *args, **kwargs: expected,
binding=binding(),
native=lambda hook, request, args, kwargs: pytest.fail("Disabled optional Rust must not call native"),
rules=rules,
)
finally:
configuration.rust(None)
assert result is expected
def test_native_stream_result_is_not_consumed_or_wrapped() -> None:
request: Final = Request(model="streaming-model")
stream: Final[Iterator[int]] = iter((1, 2))
rules: Final[Rules] = (
Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.STREAMING})),
)
dispatch: Final = PublicDispatch(
route=Route.CHAT_COMPLETIONS,
request=lambda args, kwargs: request,
context=lambda value: Context(Route.CHAT_COMPLETIONS, model=value.model, delivery=Delivery.STREAMING),
)
def native(request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Iterator[int]:
return stream
native_binding: Final[
NativeBinding[Callable[[Request, tuple[object, ...], Mapping[str, object]], Iterator[int]]]
] = NativeBinding("stream", validate=lambda _: None)
native_binding.override(native)
result: Final = dispatch.run(
("streaming-model",),
{"stream": True},
python=lambda *args, **kwargs: pytest.fail("Required native stream dispatch must not call Python"),
binding=native_binding,
native=lambda hook, value, args, kwargs: hook(value, args, kwargs),
rules=rules,
)
assert result is stream
@pytest.mark.asyncio
async def test_async_route_without_rules_preserves_async_iterator_result() -> None:
async def chunks() -> AsyncGenerator[int, None]:
yield 1
stream: Final = chunks()
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
pytest.fail("Python-only routes must not project the request")
async def python(*args: object, **kwargs: object) -> AsyncGenerator[int, None]: # kwargs-ok: pass-through shape
return stream
dispatch: Final = PublicDispatch(
route=Route.RESPONSES, request=reject_request, context=lambda _: Context(Route.RESPONSES)
)
result: Final = await dispatch.arun(
("model",),
{"stream": True},
python=python,
binding=binding(),
native=lambda hook, request, args, kwargs: pytest.fail("Python-only routes must not call native"),
rules=(),
)
assert result is stream
await stream.aclose()
@pytest.mark.asyncio
async def test_async_dispatch_accepts_websocket_style_none_result() -> None:
request: Final = Request(model="realtime-model")
rules: Final[Rules] = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.WEBSOCKET})),)
dispatch: Final = PublicDispatch(
route=Route.RESPONSES,
request=lambda args, kwargs: request,
context=lambda value: Context(Route.RESPONSES, model=value.model, delivery=Delivery.WEBSOCKET),
)
async def python(*args: object, **kwargs: object) -> None: # kwargs-ok: public pass-through shape
pytest.fail("Required native WebSocket dispatch must not call Python")
async def native(request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> None:
return None
native_binding: Final[
NativeBinding[Callable[[Request, tuple[object, ...], Mapping[str, object]], Awaitable[None]]]
] = NativeBinding("websocket", validate=lambda _: None)
native_binding.override(native)
result: Final = await dispatch.arun(
("realtime-model",),
{},
python=python,
binding=native_binding,
native=lambda hook, value, args, kwargs: hook(value, args, kwargs),
rules=rules,
)
assert result is None
def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() -> None:
rules: Final[Rules] = (
Rule(Route.MESSAGES, Rollout.RUST_REQUIRED),
Rule(Route.OCR, Rollout.PYTHON_ONLY, providers=frozenset({"mistral"})),
)
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
pytest.fail("Rules that cannot select Rust must not project the request")
dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: Context(Route.OCR))
expected: Final = object()
result: Final = dispatch.run(
("model",),
{},
python=lambda *args, **kwargs: expected,
binding=binding(),
native=lambda hook, request, args, kwargs: pytest.fail("Rules that cannot select Rust must not call native"),
rules=rules,
)
assert result is expected
@pytest.mark.asyncio
async def test_async_bypass_forwards_to_python_without_native() -> None:
request: Final = Request(model="bypassed-model")
rules: Final[Rules] = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED),)
dispatch: Final = PublicDispatch(
route=Route.RESPONSES,
request=lambda args, kwargs: request,
context=lambda value: Context(Route.RESPONSES, model=value.model),
bypass=lambda value: value.model == "bypassed-model",
)
expected: Final = object()
async def python(*args: object, **kwargs: object) -> object: # kwargs-ok: public pass-through shape
return expected
result: Final = await dispatch.arun(
("bypassed-model",),
{},
python=python,
binding=binding(),
native=lambda hook, value, args, kwargs: pytest.fail("Bypassed requests must not call native"),
rules=rules,
)
assert result is expected

54
tests/test_litellm/rust_bridge/test_failures.py Normal file
View file

@ -0,0 +1,54 @@
from types import MappingProxyType
from typing import Final
import pytest
import litellm
from litellm.rust_bridge import failures
class UpstreamRateLimited(Exception):
status_code = 429
message = "rate limited"
def test_upstream_status_maps_onto_the_public_exception_contract() -> None:
upstream: Final = UpstreamRateLimited("rate limited")
mapped: Final = failures.map_failure(upstream, "anthropic/claude-sonnet-4-5", "anthropic", MappingProxyType({}))
assert isinstance(mapped, litellm.RateLimitError)
assert mapped.llm_provider == "anthropic"
assert mapped.model == "claude-sonnet-4-5"
def test_mapper_failure_keeps_the_native_error_as_context(monkeypatch: pytest.MonkeyPatch) -> None:
def explode(**_kwargs: object) -> Exception:
raise ValueError("mapper broke")
monkeypatch.setattr(litellm, "exception_type", explode)
native_error: Final = RuntimeError("native")
mapped: Final = failures.map_failure(native_error, "mistral/mistral-ocr-latest", "mistral", MappingProxyType({}))
assert isinstance(mapped, ValueError)
assert mapped.__context__ is native_error
def test_kwargs_are_handed_to_the_mapper_as_owned_copies(monkeypatch: pytest.MonkeyPatch) -> None:
seen: Final[list[dict[str, object]]] = []
def record(**kwargs: object) -> Exception:
seen.append(dict(kwargs))
return RuntimeError("mapped")
monkeypatch.setattr(litellm, "exception_type", record)
request_kwargs: Final = MappingProxyType({"metadata": {"user_id": "u"}})
failures.map_failure(RuntimeError("native"), "gpt-4o", "openai", request_kwargs)
assert seen[0]["completion_kwargs"] == {"metadata": {"user_id": "u"}}
assert seen[0]["extra_kwargs"] == {"metadata": {"user_id": "u"}}
assert seen[0]["completion_kwargs"] is not request_kwargs
assert seen[0]["model"] == "gpt-4o"
assert seen[0]["custom_llm_provider"] == "openai"

230
tests/test_litellm/rust_bridge/test_ocr_lifecycle.py
View file

@ -1,230 +0,0 @@
from collections.abc import Generator, Mapping
from typing import Final
from unittest.mock import AsyncMock, Mock
import pytest
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import legacy
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge.ocr import LiteLLMOcrRequest
from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE
@pytest.fixture(autouse=True)
def isolated_ocr_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield
NATIVE_OCR_LIFECYCLE.reset()
configuration.reset_rust_configuration()
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_unavailable_native_uses_legacy(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None:
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback)
NATIVE_OCR_LIFECYCLE.override(None)
document: Final = {"type": "document_url", "document_url": "https://example.com"}
result: Final = (
await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[0])
if asynchronous
else litellm.ocr("mistral/mistral-ocr-latest", document, pages=[0])
)
assert result is response
fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[0])
def test_admitted_failure_is_returned_without_replay() -> None:
failure: Final = RuntimeError("admitted")
native: Final = Mock(side_effect=failure)
litellm.rust(True)
NATIVE_OCR_LIFECYCLE.override(native)
try:
with pytest.raises(RuntimeError) as caught:
litellm.ocr("mistral/mistral-ocr-latest", {"type": "document_url", "document_url": "https://example.com"})
assert caught.value is failure
finally:
NATIVE_OCR_LIFECYCLE.reset()
litellm.rust(None)
assert native.call_count == 1
def test_public_binding_keeps_positional_fields_and_defaults_out_of_native_hook_kwargs() -> None:
document: Final = {"type": "document_url", "document_url": "https://example.com"}
captured: Final = []
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
asynchronous: bool,
) -> OCRResponse:
captured.append((request, args, kwargs, asynchronous))
return OCRResponse(pages=[], model=request.model)
litellm.rust(True)
NATIVE_OCR_LIFECYCLE.override(native)
try:
response: Final = litellm.ocr("mistral/mistral-ocr-latest", document)
finally:
NATIVE_OCR_LIFECYCLE.reset()
litellm.rust(None)
request, call_args, hook_kwargs, asynchronous = captured[0]
assert response.model == "mistral/mistral-ocr-latest"
assert request.model == "mistral/mistral-ocr-latest"
assert request.document is document
assert call_args == ("mistral/mistral-ocr-latest", document)
assert hook_kwargs == {}
assert asynchronous is False
def test_public_binding_keeps_keyword_model_and_document_in_native_hook_kwargs() -> None:
document: Final = {"type": "document_url", "document_url": "https://example.com"}
captured: Final = []
def native(
request: LiteLLMOcrRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
asynchronous: bool,
) -> OCRResponse:
assert args == ()
captured.append(kwargs)
return OCRResponse(pages=[], model=request.model)
litellm.rust(True)
NATIVE_OCR_LIFECYCLE.override(native)
try:
litellm.ocr(model="mistral/mistral-ocr-latest", document=document)
finally:
NATIVE_OCR_LIFECYCLE.reset()
litellm.rust(None)
assert captured[0]["model"] == "mistral/mistral-ocr-latest"
assert captured[0]["document"] is document
assert "timeout" not in captured[0]
@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"])
def test_public_duplicate_argument_error_does_not_depend_on_native_selection(enabled: bool) -> None:
native: Final = Mock(side_effect=AssertionError("binding errors precede admission"))
document: Final = {"type": "document_url", "document_url": "https://example.com"}
litellm.rust(enabled)
NATIVE_OCR_LIFECYCLE.override(native)
try:
with pytest.raises(TypeError, match=r"ocr\(\) got multiple values for argument 'model'"):
litellm.ocr("mistral/mistral-ocr-latest", document, model="duplicate")
finally:
NATIVE_OCR_LIFECYCLE.reset()
litellm.rust(None)
assert native.call_count == 0
@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"])
def test_public_missing_required_argument_error_does_not_depend_on_native_selection(enabled: bool) -> None:
native: Final = Mock(side_effect=AssertionError("binding errors precede admission"))
litellm.rust(enabled)
NATIVE_OCR_LIFECYCLE.override(native)
try:
with pytest.raises(TypeError, match=r"ocr\(\) missing 1 required positional argument: 'document'"):
litellm.ocr("mistral/mistral-ocr-latest")
finally:
NATIVE_OCR_LIFECYCLE.reset()
litellm.rust(None)
assert native.call_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("enabled", [False, True, None])
async def test_environment_opt_out_never_loads_native(
monkeypatch: pytest.MonkeyPatch, asynchronous: bool, enabled: bool | None
) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback)
load: Final = Mock(side_effect=AssertionError("native must not be loaded"))
monkeypatch.setattr(bindings, "get_native_bridge", load)
litellm.rust(enabled)
document: Final = {"type": "file", "file": b"pdf"}
result: Final = (
await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[1])
if asynchronous
else litellm.ocr("mistral/mistral-ocr-latest", document, pages=[1])
)
assert result is response
fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[1])
load.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("environment", [None, "1"])
async def test_native_is_enabled_by_default(
monkeypatch: pytest.MonkeyPatch, asynchronous: bool, environment: str | None
) -> None:
if environment is not None:
monkeypatch.setenv("LITELLM_RUST", environment)
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
native: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
NATIVE_OCR_LIFECYCLE.override(native)
fallback: Final = Mock(side_effect=AssertionError("legacy must not run"))
monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback)
result: Final = (
await litellm.aocr("mistral/mistral-ocr-latest", {})
if asynchronous
else litellm.ocr("mistral/mistral-ocr-latest", {})
)
assert result is response
assert native.call_count == 1
fallback.assert_not_called()
class Declined(Exception):
pass
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("declined", [False, True])
async def test_only_native_declines_replay_on_legacy(
monkeypatch: pytest.MonkeyPatch, asynchronous: bool, declined: bool
) -> None:
failure: Final = Declined("unsupported") if declined else RuntimeError("provider already called")
native: Final = AsyncMock(side_effect=failure) if asynchronous else Mock(side_effect=failure)
NATIVE_OCR_LIFECYCLE.override(native)
import importlib
main: Final = importlib.import_module("litellm.ocr.main")
monkeypatch.setattr(main, "native_exception_types", lambda: (Declined, RuntimeError))
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback)
document: Final = {"type": "file", "file": b"pdf"}
async def call() -> object:
if asynchronous:
return await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[0])
return litellm.ocr("mistral/mistral-ocr-latest", document, pages=[0])
if declined:
assert await call() is response
fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[0])
else:
with pytest.raises(RuntimeError) as caught:
await call()
assert caught.value is failure
fallback.assert_not_called()
assert native.call_count == 1

377
tests/test_litellm/rust_bridge/test_runtime.py
View file

@ -1,11 +1,17 @@
from __future__ import annotations
from collections.abc import Callable, Generator
from types import SimpleNamespace
from typing import Final, Protocol
import pytest
from litellm.exceptions import APIError
from litellm.rust_bridge import bindings, runtime
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.rust_bridge import bindings, configuration, runtime
from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule
from litellm.rust_bridge.configuration import Rollout
class RustBridgeDeclined(Exception):
@ -17,79 +23,352 @@ class RustUpstreamError(Exception):
@pytest.fixture(autouse=True)
def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
native = SimpleNamespace(
def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
native: Final = SimpleNamespace(
RustBridgeDeclined=RustBridgeDeclined,
RustUpstreamError=RustUpstreamError,
)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield
configuration.reset_rust_configuration()
def context() -> runtime.BridgeErrorContext:
return runtime.BridgeErrorContext(route="messages", provider="anthropic", model="model")
class NativeFn(Protocol):
def __call__(self) -> str: ...
def test_invoke_tags_native_decline_before_running_fallback() -> None:
calls: list[str] = []
CONTEXT: Final = Context(Route.MESSAGES, provider="anthropic", model="model")
RUST: Final = "rust"
PYTHON: Final = "python"
def decline() -> object:
calls.append("rust")
raise RustBridgeDeclined("unsupported")
value = runtime.invoke(
native_call=decline,
fallback=lambda: calls.append("python") or "fallback",
adapt=str,
mode=runtime.FallbackMode.PYTHON,
context=context(),
def binding(native: NativeFn | None) -> bindings.NativeBinding[NativeFn]:
bound: Final[bindings.NativeBinding[NativeFn]] = bindings.NativeBinding("_messages", validate=lambda _: None)
bound.override(native)
return bound
def rules(rollout: Rollout) -> tuple[Rule, ...]:
return (Rule(Route.MESSAGES, rollout, providers=frozenset({"anthropic"})),)
class Recorder:
def __init__(self, native_effect: BaseException | None = None) -> None:
self._native_effect: Final = native_effect
self.calls: tuple[str, ...] = ()
def rust(self) -> str:
self.calls = (*self.calls, RUST)
if self._native_effect is not None:
raise self._native_effect
return RUST
def python(self) -> str:
self.calls = (*self.calls, PYTHON)
return PYTHON
def recorder(native_effect: BaseException | None = None) -> Recorder:
return Recorder(native_effect)
def run(rollout: Rollout, calls: Recorder, *, native_missing: bool = False, context: Context = CONTEXT) -> str:
return runtime.run(
context,
binding=binding(None if native_missing else calls.rust),
native=lambda fn: fn(),
python=calls.python,
rules=rules(rollout),
)
assert value == "fallback"
assert calls == ["rust", "python"]
@pytest.mark.parametrize(
("rollout", "switch", "expected"),
(
(Rollout.PYTHON_ONLY, None, (PYTHON,)),
(Rollout.PYTHON_ONLY, True, (PYTHON,)),
(Rollout.RUST_OPT_IN, None, (PYTHON,)),
(Rollout.RUST_OPT_IN, True, (RUST,)),
(Rollout.RUST_OPT_OUT, None, (RUST,)),
(Rollout.RUST_OPT_OUT, False, (PYTHON,)),
(Rollout.RUST_REQUIRED, None, (RUST,)),
(Rollout.RUST_REQUIRED, False, (RUST,)),
),
)
def test_rollout_and_switch_select_native_or_python(
rollout: Rollout, switch: bool | None, expected: tuple[str, ...]
) -> None:
calls: Final = recorder()
if switch is not None:
configuration.rust(switch)
assert run(rollout, calls) == expected[-1]
assert calls.calls == expected
def test_invoke_translates_upstream_without_fallback() -> None:
def fail() -> object:
raise RustUpstreamError(429, "rate limited")
def test_environment_switch_enables_opt_in_route(monkeypatch: pytest.MonkeyPatch) -> None:
calls: Final = recorder()
monkeypatch.setenv("LITELLM_RUST", "1")
with pytest.raises(APIError, match="rate limited") as caught:
runtime.invoke(
native_call=fail,
fallback=lambda: pytest.fail("fallback must not run"),
adapt=str,
mode=runtime.FallbackMode.PYTHON,
context=context(),
)
assert run(Rollout.RUST_OPT_IN, calls) == "rust"
assert calls.calls == (RUST,)
assert caught.value.status_code == 429
@pytest.mark.parametrize(
("rollout", "environment", "switch", "expected"),
(
(Rollout.RUST_OPT_IN, "0", True, (PYTHON,)),
(Rollout.RUST_OPT_OUT, "0", True, (PYTHON,)),
(Rollout.RUST_OPT_IN, "1", False, (RUST,)),
(Rollout.RUST_OPT_OUT, "1", False, (RUST,)),
(Rollout.RUST_REQUIRED, "0", False, (RUST,)),
(Rollout.PYTHON_ONLY, "1", True, (PYTHON,)),
),
)
def test_environment_switch_wins_over_process_switch(
monkeypatch: pytest.MonkeyPatch,
rollout: Rollout,
environment: str,
switch: bool,
expected: tuple[str, ...],
) -> None:
calls: Final = recorder()
monkeypatch.setenv("LITELLM_RUST", environment)
configuration.rust(switch)
assert run(rollout, calls) == expected[-1]
assert calls.calls == expected
def test_context_outside_rule_stays_on_python() -> None:
calls: Final = recorder()
configuration.rust(True)
assert run(Rollout.RUST_REQUIRED, calls, context=Context(Route.MESSAGES, provider="openai")) == "python"
assert run(Rollout.RUST_REQUIRED, calls, context=Context(Route.RESPONSES, provider="anthropic")) == "python"
assert calls.calls == (PYTHON, PYTHON)
@pytest.mark.asyncio
async def test_ainvoke_handles_native_success() -> None:
async def native() -> int:
return 3
@pytest.mark.parametrize(
"context",
(
Context(Route.CHAT_COMPLETIONS, provider="anthropic"),
Context(Route.CHAT_COMPLETIONS, provider="bedrock"),
Context(Route.MESSAGES, provider="anthropic"),
Context(Route.RESPONSES, provider="openai"),
Context(Route.TRANSCRIPTION, provider="openai"),
),
)
@pytest.mark.parametrize("delivery", tuple(Delivery))
async def test_shipped_python_routes_never_load_native(
monkeypatch: pytest.MonkeyPatch, context: Context, delivery: Delivery
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1")
configuration.rust(True)
calls: Final = recorder()
request: Final = Context(context.route, provider=context.provider, delivery=delivery)
async def fallback() -> str:
pytest.fail("fallback must not run")
def reject_load(value: object) -> NativeFn | None:
pytest.fail("Python-only dispatch must not load a native binding")
bound: Final = bindings.NativeBinding("_messages", validate=reject_load)
async def native(fn: NativeFn) -> str:
return fn()
async def python() -> str:
return calls.python()
assert runtime.run(request, binding=bound, native=lambda fn: fn(), python=calls.python) == PYTHON
assert await runtime.arun(request, binding=bound, native=native, python=python) == PYTHON
assert calls.calls == (PYTHON, PYTHON)
def test_native_decline_falls_back_to_python_once() -> None:
calls: Final = recorder(RustBridgeDeclined("unsupported"))
assert run(Rollout.RUST_OPT_OUT, calls) == "python"
assert calls.calls == (RUST, PYTHON)
def test_unavailable_native_falls_back_to_python() -> None:
calls: Final = recorder()
assert run(Rollout.RUST_OPT_OUT, calls, native_missing=True) == "python"
assert calls.calls == (PYTHON,)
@pytest.mark.asyncio
@pytest.mark.parametrize("missing", (False, True))
async def test_python_fallback_does_not_claim_rust_execution(missing: bool) -> None:
calls: Final = recorder(RustBridgeDeclined("unsupported"))
bound: Final = binding(None if missing else calls.rust)
expected: Final = OCRResponse(pages=[], model="python")
def native(fn: NativeFn) -> OCRResponse:
fn()
pytest.fail("native must decline before constructing a response")
async def anative(fn: NativeFn) -> OCRResponse:
return native(fn)
async def python() -> OCRResponse:
return expected
assert (
await runtime.ainvoke(
native_call=native,
fallback=fallback,
adapt=str,
mode=runtime.FallbackMode.PYTHON,
context=context(),
runtime.run(CONTEXT, binding=bound, native=native, python=lambda: expected, rules=rules(Rollout.RUST_OPT_OUT))
is expected
)
assert (
await runtime.arun(CONTEXT, binding=bound, native=anative, python=python, rules=rules(Rollout.RUST_OPT_OUT))
is expected
)
assert get_hidden_params_dict(expected) == {}
@pytest.mark.asyncio
@pytest.mark.parametrize("shape", ("model", "dict"))
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_native_response_marker_reaches_caller_with_existing_metadata(shape: str, asynchronous: bool) -> None:
hidden: Final = {"additional_headers": {"x-request-id": "upstream"}, "response_cost": 0.01}
response: Final[OCRResponse | dict[str, object]] = (
OCRResponse(pages=[], model="native") if shape == "model" else {"content": "native", "_hidden_params": hidden}
)
if isinstance(response, OCRResponse):
response._hidden_params = hidden # pyright: ignore[reportPrivateUsage] # seed SDK metadata to verify it survives native marking
bound: Final[bindings.NativeBinding[Callable[[], object]]] = bindings.NativeBinding("ocr", validate=lambda _: None)
bound.override(lambda: response)
def python() -> object:
pytest.fail("native success must not fall back")
async def anative(fn: Callable[[], object]) -> object:
return fn()
async def apython() -> object:
return python()
result: Final = (
await runtime.arun(CONTEXT, binding=bound, native=anative, python=apython, rules=rules(Rollout.RUST_REQUIRED))
if asynchronous
else runtime.run(
CONTEXT, binding=bound, native=lambda fn: fn(), python=python, rules=rules(Rollout.RUST_REQUIRED)
)
== "3"
)
assert result is response
assert get_hidden_params_dict(result) == {
"response_cost": 0.01,
"additional_headers": {"x-request-id": "upstream", "x-litellm-rust": "true"},
}
def test_upstream_error_maps_to_api_error_without_fallback() -> None:
calls: Final = recorder(RustUpstreamError(429, "rate limited"))
with pytest.raises(APIError, match="rate limited") as caught:
run(Rollout.RUST_OPT_OUT, calls)
assert caught.value.status_code == 429
assert calls.calls == (RUST,)
def test_other_native_errors_propagate_without_fallback() -> None:
failure: Final = ValueError("admitted")
calls: Final = recorder(failure)
with pytest.raises(ValueError, match="admitted") as caught:
run(Rollout.RUST_OPT_OUT, calls)
assert caught.value is failure
assert calls.calls == (RUST,)
def test_required_route_rejects_unavailable_bridge() -> None:
calls: Final = recorder()
with pytest.raises(RuntimeError, match="Rust messages bridge is unavailable"):
run(Rollout.RUST_REQUIRED, calls, native_missing=True)
assert PYTHON not in calls.calls
def test_required_route_rejects_native_decline() -> None:
calls: Final = recorder(RustBridgeDeclined("unsupported"))
with pytest.raises(RuntimeError, match="declined the request: unsupported"):
run(Rollout.RUST_REQUIRED, calls)
assert PYTHON not in calls.calls
@pytest.mark.asyncio
@pytest.mark.parametrize(
("native_effect", "native_missing", "expected"),
(
(None, False, (RUST,)),
(RustBridgeDeclined("unsupported"), False, (RUST, PYTHON)),
(None, True, (PYTHON,)),
),
)
async def test_arun_mirrors_sync_fallback(
native_effect: BaseException | None, native_missing: bool, expected: tuple[str, ...]
) -> None:
calls: Final = recorder(native_effect)
async def native(fn: NativeFn) -> str:
return fn()
async def python() -> str:
return calls.python()
result: Final = await runtime.arun(
CONTEXT,
binding=binding(None if native_missing else calls.rust),
native=native,
python=python,
rules=rules(Rollout.RUST_OPT_OUT),
)
assert result == expected[-1]
assert calls.calls == expected
@pytest.mark.asyncio
async def test_arun_required_route_rejects_unavailable_bridge() -> None:
async def python() -> str:
pytest.fail("fallback must not run")
def test_required_mode_rejects_unavailable_bridge() -> None:
with pytest.raises(RuntimeError, match="is unavailable"):
runtime.invoke(
native_call=None,
fallback=lambda: pytest.fail("fallback must not run"),
adapt=str,
mode=runtime.FallbackMode.RUST_REQUIRED,
context=context(),
await runtime.arun(
CONTEXT,
binding=binding(None),
native=lambda fn: python(),
python=python,
rules=rules(Rollout.RUST_REQUIRED),
)
@pytest.mark.asyncio
async def test_arun_upstream_error_maps_to_api_error_without_fallback() -> None:
calls: Final = recorder(RustUpstreamError(503, "upstream unavailable"))
async def native(fn: NativeFn) -> str:
return fn()
async def python() -> str:
return calls.python()
with pytest.raises(APIError, match="upstream unavailable") as caught:
await runtime.arun(
CONTEXT,
binding=binding(calls.rust),
native=native,
python=python,
rules=rules(Rollout.RUST_OPT_OUT),
)
assert caught.value.status_code == 503
assert calls.calls == (RUST,)

204
tests/test_litellm/test_audio_transcription_rust_bridge.py
View file

@ -1,16 +1,44 @@
import importlib
from __future__ import annotations
from collections.abc import Generator
from types import SimpleNamespace
from typing import Final
import pytest
import litellm
from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge.transcription.native import NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION
rust_bridge = importlib.import_module("litellm.rust_bridge.transcription")
MODEL: Final = "bedrock/mistral.voxtral-mini-3b-2507"
AUDIO_FILE: Final = ("audio.wav", b"audio", "audio/wav")
class RustBridgeDeclined(Exception):
pass
class RustUpstreamError(Exception):
pass
@pytest.fixture(autouse=True)
def isolated_bridge(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
native: Final = SimpleNamespace(RustBridgeDeclined=RustBridgeDeclined, RustUpstreamError=RustUpstreamError)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield
NATIVE_TRANSCRIPTION.reset()
NATIVE_ATRANSCRIPTION.reset()
configuration.reset_rust_configuration()
class SyncBridge:
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
def __init__(self, effect: BaseException | None = None) -> None:
self._effect: Final = effect
self.calls: tuple[dict[str, object], ...] = ()
def __call__(
self,
@ -23,11 +51,19 @@ class SyncBridge:
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]:
self.calls.append({"model": model, "audio": audio, "optional_params": optional_params})
return {"text": "hello"}
self.calls = (
*self.calls,
{"model": model, "audio": audio, "provider": custom_llm_provider, "timeout": timeout_seconds},
)
if self._effect is not None:
raise self._effect
return {"text": "rust"}
class AsyncBridge:
def __init__(self) -> None:
self.calls: tuple[str, ...] = ()
async def __call__(
self,
model: str,
@ -39,113 +75,109 @@ class AsyncBridge:
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]:
return {"text": "async"}
self.calls = (*self.calls, model)
return {"text": "async rust"}
def test_enabled_sync_bridge_receives_audio() -> None:
bridge = SyncBridge()
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"},
def dispatch_sync() -> litellm.TranscriptionResponse:
return BedrockAudioTranscriptionRustDispatch().audio_transcriptions(
model=MODEL,
audio_file=AUDIO_FILE,
api_key=None,
api_base=None,
custom_llm_provider="bedrock",
extra_headers=None,
optional_params={"temperature": 0},
timeout=5.0,
timeout=5,
)
assert result == {"text": "hello"}
assert bridge.calls[0]["audio"] == {"data": "AQI=", "format": "wav", "filename": "audio.wav"}
@pytest.mark.asyncio
async def test_enabled_async_bridge() -> None:
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"},
api_key=None,
api_base=None,
custom_llm_provider="bedrock",
extra_headers=None,
optional_params={},
timeout=None,
def test_dispatch_marshals_audio_into_rust_call() -> None:
bridge: Final = SyncBridge()
NATIVE_TRANSCRIPTION.override(bridge)
response: Final = dispatch_sync()
assert response.text == "rust"
assert bridge.calls == (
{
"model": MODEL,
"audio": {"data": "YXVkaW8=", "format": "wav", "filename": "audio.wav"},
"provider": "bedrock",
"timeout": 5.0,
},
)
assert result == {"text": "async"}
def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None:
rust_bridge.configure_rust_transcription(transcription=None, atranscription=None)
monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", lambda: None)
assert rust_bridge.load_rust_transcription() is None
assert rust_bridge.load_rust_atranscription() is None
@pytest.mark.parametrize("disable", ("process", "environment"))
def test_bedrock_transcription_ignores_optional_rust_switches(disable: str, monkeypatch: pytest.MonkeyPatch) -> None:
if disable == "process":
litellm.rust(False)
else:
monkeypatch.setenv("LITELLM_RUST", "0")
bridge: Final = SyncBridge()
NATIVE_TRANSCRIPTION.override(bridge)
assert dispatch_sync().text == "rust"
assert len(bridge.calls) == 1
def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None)
def test_missing_native_binding_raises_without_python_fallback() -> None:
NATIVE_TRANSCRIPTION.override(None)
with pytest.raises(RuntimeError, match="bridge is unavailable"):
BedrockAudioTranscriptionRustDispatch().audio_transcriptions(
model="bedrock/mistral.voxtral-mini-3b-2507",
audio_file=("audio.wav", b"audio", "audio/wav"),
api_key=None,
api_base=None,
custom_llm_provider="bedrock",
extra_headers=None,
optional_params={},
timeout=5,
)
dispatch_sync()
def test_admission_decline_raises_for_required_route() -> None:
NATIVE_TRANSCRIPTION.override(SyncBridge(RustBridgeDeclined("unsupported format")))
with pytest.raises(RuntimeError, match="declined the request: unsupported format"):
dispatch_sync()
def test_upstream_error_maps_to_api_error() -> None:
NATIVE_TRANSCRIPTION.override(SyncBridge(RustUpstreamError(503, "bedrock down")))
with pytest.raises(litellm.APIError, match="bedrock down") as raised:
dispatch_sync()
assert raised.value.status_code == 503
def test_bedrock_transcription_dispatches_to_rust_from_sdk_entrypoint() -> None:
bridge: Final = SyncBridge()
NATIVE_TRANSCRIPTION.override(bridge)
response: Final = litellm.transcription(model=MODEL, file=AUDIO_FILE)
assert isinstance(response, litellm.TranscriptionResponse)
assert response.text == "rust"
assert bridge.calls[0]["model"] == MODEL.removeprefix("bedrock/")
@pytest.mark.asyncio
async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None:
async def unavailable(**_: object) -> None:
return None
async def test_bedrock_atranscription_dispatches_to_rust_from_sdk_entrypoint() -> None:
bridge: Final = AsyncBridge()
NATIVE_ATRANSCRIPTION.override(bridge)
monkeypatch.setattr(rust_bridge, "atranscription", unavailable)
response: Final = await litellm.atranscription(model=MODEL, file=AUDIO_FILE)
assert response.text == "async rust"
assert bridge.calls == (MODEL.removeprefix("bedrock/"),)
@pytest.mark.asyncio
async def test_async_missing_native_binding_raises_without_python_fallback() -> None:
NATIVE_ATRANSCRIPTION.override(None)
with pytest.raises(RuntimeError, match="bridge is unavailable"):
await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions(
model="bedrock/mistral.voxtral-mini-3b-2507",
audio_file=("audio.wav", b"audio", "audio/wav"),
model=MODEL,
audio_file=AUDIO_FILE,
api_key=None,
api_base=None,
custom_llm_provider="bedrock",
extra_headers=None,
optional_params={},
timeout=5,
timeout=None,
)
def test_bedrock_transcription_uses_rust_only_path() -> None:
rust_bridge.configure_rust_transcription(
transcription=lambda **_: {"text": "rust"},
atranscription=None,
)
try:
response = litellm.transcription(
model="bedrock/mistral.voxtral-mini-3b-2507",
file=("audio.wav", b"audio", "audio/wav"),
)
finally:
rust_bridge.configure_rust_transcription(transcription=None, atranscription=None)
assert response.text == "rust"
@pytest.mark.asyncio
async def test_bedrock_atranscription_uses_rust_only_path() -> None:
async def rust_response(**_: object) -> dict[str, object]:
return {"text": "rust"}
rust_bridge.configure_rust_transcription(transcription=None, atranscription=rust_response)
try:
response = await litellm.atranscription(
model="bedrock/mistral.voxtral-mini-3b-2507",
file=("audio.wav", b"audio", "audio/wav"),
)
finally:
rust_bridge.configure_rust_transcription(transcription=None, atranscription=None)
assert response.text == "rust"

4
tests/test_litellm_rust/ocr/test_lifecycle.py
View file

@ -580,7 +580,7 @@ async def test_retained_argument_aliases_and_body_roots_survive_envelope_replace
def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_server: RecordingServer) -> None:
from litellm.ocr.main import _public_request
from litellm.ocr.dispatch import _public_request
from litellm.rust_bridge import _native
ocr_server.expected_requests = 0
@ -594,7 +594,7 @@ def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_serv
def create():
file: Final = File()
kwargs: Final = {"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": file}}
coroutine: Final = _native._ocr_lifecycle(_public_request("aocr", (), kwargs), (), kwargs, True)
coroutine: Final = _native.aocr(_public_request("aocr", (), kwargs), (), kwargs)
file.owner = coroutine
coroutine.close()
return weakref.ref(file)

46
tests/test_litellm_rust/test_ocr.py
View file

@ -8,7 +8,6 @@ from typing import Final
import pytest
import litellm
from litellm.rust_bridge import ocr as rust_ocr_bridge
pytestmark = pytest.mark.requires_rust_extension
@ -71,35 +70,6 @@ def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[dict[str, object]]
thread.join()
def test_native_ocr_with_compiled_rust_extension(
ocr_server: tuple[ThreadingHTTPServer, list[dict[str, object]]],
) -> None:
server, requests = ocr_server
address: Final = server.server_address
host: Final = str(address[0])
port: Final = int(address[1])
response: Final = rust_ocr_bridge.ocr(
model="mistral-ocr-latest",
document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
api_key="test-key",
api_base=f"http://{host}:{port}",
custom_llm_provider="mistral",
extra_headers=None,
optional_params={},
timeout=None,
)
assert response is not None
assert response["pages"][0]["markdown"] == "native OCR response"
assert len(requests) == 1
assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx")
assert requests[0]["body"] == {
"model": "mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
}
@pytest.mark.parametrize(
"file_input,mime_type,expected_type,expected_field,expected_uri",
[
@ -219,22 +189,6 @@ async def test_native_ocr_failures_do_not_retry_on_python(ocr_server, asynchrono
assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx")
@pytest.mark.parametrize("custom_provider", ["mistral", "not-a-provider"])
def test_native_ocr_rejects_invalid_input_before_network(ocr_server, custom_provider):
from litellm.rust_bridge import _native
server, requests = ocr_server
with pytest.raises(ValueError, match="Document URL is required"):
_native.ocr(
model="mistral-ocr-latest",
custom_llm_provider=custom_provider,
document={"type": "document_url"},
api_key="test-key",
api_base=f"http://127.0.0.1:{server.server_port}",
)
assert requests == []
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_native_ocr_enforces_request_deadline_without_fallback(ocr_server, asynchronous):