feat(rust): discover providers at shared SDK dispatch boundaries

This commit is contained in:
Yujong Lee 2026-09-05 14:01:30 -07:00 committed by yujonglee
parent a278ee0ef1
commit f177246277
39 changed files with 3296 additions and 3102 deletions

View file

@ -1 +1 @@
pub use crate::ocr::{OcrRequest, ocr};
pub use crate::ocr::{OcrRequest, ocr, ocr_provider_supported};

View file

@ -40,6 +40,14 @@ impl ResponsesWebSocketConnection {
options: &RequestOptions,
_context: &LiteLlmRequestContext,
) -> Result<Self, Error> {
if !litellm_core::responses::websocket::native_websocket_supported(
options
.custom_llm_provider
.as_deref()
.unwrap_or("openai"),
) {
return Err(Error::Unsupported("unsupported native WebSocket provider"));
}
let headers = string_headers("Responses WebSocket", options.extra_headers.clone())?;
let mut request = input
.url

View file

@ -35,6 +35,10 @@ pub async fn ocr(
.await
}
pub fn ocr_provider_supported(model: &str, provider: &str) -> bool {
common_utils::ocr_provider_config(provider, model).is_some()
}
#[cfg(test)]
mod tests {
use crate::integrations::types::RequestHooks;

View file

@ -26,5 +26,9 @@ pub async fn audio_transcription(
.await
}
pub fn transcription_provider_supported(provider: &str) -> bool {
prepare::provider_config(provider).is_some()
}
#[cfg(test)]
mod tests;

View file

@ -9,7 +9,9 @@ use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderCo
use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
pub(super) fn provider_config(
provider: &str,
) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
#[cfg(feature = "bedrock-auth")]
if provider == "bedrock" {
return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);

View file

@ -37,5 +37,9 @@ pub async fn messages_stream(
execute_messages_provider_stream(request, options.clone()).await
}
pub fn messages_provider_supported(provider: &str) -> bool {
common_utils::messages_provider_config(provider).is_some()
}
#[cfg(test)]
mod tests;

View file

@ -2,6 +2,14 @@ use crate::Error;
use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH};
use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult};
pub fn native_websocket_supported(provider: &str) -> bool {
match provider {
"openai" => crate::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG
.supports_native_websocket(),
_ => false,
}
}
pub trait ResponsesWebSocketProviderConfig: Sync {
fn supports_native_websocket(&self) -> bool {
false

View file

@ -35,6 +35,16 @@ impl ResponsesWebSocketConnection {
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<Bound<'py, PyAny>> {
if let Some(reason) = responses_websocket_decline(
"responses websocket",
options.provider("openai"),
false,
false,
false,
None,
) {
return Err(crate::errors::RustBridgeDeclined::new_err(reason));
}
let options: litellm_core::request_options::RequestOptions = options.into();
let context: litellm_core::request_context::LiteLlmRequestContext = context.into();
let request = ResponsesWebSocketRequest { url: request.url };
@ -68,6 +78,25 @@ impl ResponsesWebSocketConnection {
}
}
#[pyfunction]
#[pyo3(signature = (_model, custom_llm_provider, *, stream=false, has_agentic_hook=false, has_custom_client=false, request_format=None))]
fn responses_websocket_decline(
_model: &str,
custom_llm_provider: &str,
stream: bool,
has_agentic_hook: bool,
has_custom_client: bool,
request_format: Option<&str>,
) -> Option<String> {
routes::definition::request_decline(
litellm_core::responses::websocket::native_websocket_supported(custom_llm_provider),
stream,
has_agentic_hook,
has_custom_client,
request_format,
)
}
#[pymodule(gil_used = false)]
mod _native {
use pyo3::prelude::*;
@ -77,6 +106,10 @@ mod _native {
super::errors::register(module)?;
super::routes::register(module)?;
module.add_class::<super::ResponsesWebSocketConnection>()?;
module.add_function(wrap_pyfunction!(
super::responses_websocket_decline,
module
)?)?;
super::diagnostics::register(module)
}
}
@ -101,16 +134,20 @@ mod tests {
let expected = [
"RustBridgeDeclined",
"RustUpstreamError",
"ocr_decline",
"ocr",
"aocr",
"transcription_decline",
"transcription",
"atranscription",
"messages_decline",
"messages",
"amessages",
"chat_completions_decline",
"chat_completions",
"achat_completions",
"ResponsesWebSocketConnection",
"responses_websocket_decline",
"gil_stats",
];

View file

@ -84,6 +84,12 @@ pub(crate) struct NativeRequestOptions {
vertex: Option<NativeVertexOptions>,
}
impl NativeRequestOptions {
pub(crate) fn provider(&self, default: &'static str) -> &str {
self.custom_llm_provider.as_deref().unwrap_or(default)
}
}
impl From<NativeRequestOptions> for litellm_core::request_options::RequestOptions {
fn from(input: NativeRequestOptions) -> Self {
Self {

View file

@ -22,6 +22,23 @@ fn prepare_transcription(
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
if let Some(reason) = transcription_decline(
&input.model,
input.options.provider("bedrock"),
input
.optional_params
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(false),
false,
false,
input
.optional_params
.get("response_format")
.and_then(Value::as_str),
) {
return Err(crate::errors::RustBridgeDeclined::new_err(reason));
}
let context: LiteLlmRequestContext = context.into();
let audio = input.audio;
Ok(async move {
@ -38,10 +55,30 @@ fn prepare_transcription(
})
}
#[pyfunction]
#[pyo3(signature = (_model, custom_llm_provider, *, stream=false, has_agentic_hook=false, has_custom_client=false, request_format=None))]
fn transcription_decline(
_model: &str,
custom_llm_provider: &str,
stream: bool,
has_agentic_hook: bool,
has_custom_client: bool,
request_format: Option<&str>,
) -> Option<String> {
super::definition::request_decline(
litellm_core::audio_transcription::transcription_provider_supported(custom_llm_provider),
stream,
has_agentic_hook,
has_custom_client,
request_format,
)
}
bridge_route! {
sync = transcription,
asynchronous = atranscription,
request = AudioTranscriptionInputs,
prepare = prepare_transcription,
errors = core_error_to_pyerr,
extra = [transcription_decline],
}

View file

@ -5,7 +5,6 @@ use litellm_core::chat_completions::chat_completions as run_route;
use litellm_core::chat_completions::chat_completions_decline_reason;
use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
use litellm_core::request_context::LiteLlmRequestContext;
use litellm_core::request_options::RequestOptions;
use pyo3::prelude::*;
use serde_json::{Map, Value};
use std::future::Future;
@ -40,26 +39,52 @@ fn prepare_chat_completions(
})
}
fn preflight_context(context: &Bound<'_, PyAny>) -> PyResult<LiteLlmRequestContext> {
let metadata = context.getattr("metadata")?;
let user_id = if metadata.is_none() {
None
} else {
match metadata.get_item("user_id") {
Ok(value) => Some(if value.is_none() {
Value::Null
} else {
Value::Bool(true)
}),
Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(context.py()) => {
None
}
Err(error) => return Err(error),
}
};
Ok(LiteLlmRequestContext {
metadata: user_id.map(|value| Map::from_iter([("user_id".into(), value)])),
request_metadata_fields: context.getattr("request_metadata_fields")?.extract()?,
..Default::default()
})
}
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None, *, options, context))]
#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None, *, context, stream=false, has_custom_client=false, has_agentic_hook=false))]
#[allow(
clippy::too_many_arguments,
reason = "PyO3 preserves chat preflight arguments alongside request features"
)]
fn chat_completions_decline(
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)] messages: Value,
#[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option<Value>,
custom_llm_provider: Option<String>,
options: NativeRequestOptions,
context: NativeRequestContext,
context: &Bound<'_, PyAny>,
stream: bool,
has_custom_client: bool,
has_agentic_hook: bool,
) -> PyResult<Option<String>> {
if !matches!(
custom_llm_provider.as_deref(),
Some("anthropic" | "bedrock")
) {
return Ok(Some(
"provider is not on the rust chat completions path".into(),
));
if let Some(reason) =
super::definition::request_decline(true, stream, has_agentic_hook, has_custom_client, None)
{
return Ok(Some(reason));
}
let context: LiteLlmRequestContext = context.into();
let options: RequestOptions = options.into();
let context = preflight_context(context)?;
let optional_params = match optional_params {
None | Some(Value::Null) => Map::new(),
Some(Value::Object(params)) => params,
@ -74,7 +99,6 @@ fn chat_completions_decline(
custom_llm_provider.as_deref(),
messages,
&optional_params,
&options,
&context,
)
.map(str::to_string))

View file

@ -98,6 +98,29 @@ pub(super) fn add_function(
module.add_function(function)
}
pub(crate) fn request_decline(
provider_supported: bool,
stream: bool,
has_agentic_hook: bool,
has_custom_client: bool,
request_format: Option<&str>,
) -> Option<String> {
let reason = if !provider_supported {
Some("unsupported native provider")
} else if stream {
Some("native streaming is unavailable")
} else if has_agentic_hook {
Some("native agentic hooks are unavailable")
} else if has_custom_client {
Some("native custom clients are unavailable")
} else if request_format == Some("native") {
Some("native OCR response format is unavailable")
} else {
None
};
reason.map(str::to_string)
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
@ -262,6 +285,64 @@ for field in ('litellm_call_id', 'trace_id', 'request_model'):
});
}
#[test]
fn acceptance_and_execution_decline_unsupported_requests_before_io() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "routes").expect("module should be created");
crate::routes::register(&module).expect("routes should register");
module
.add_class::<crate::ResponsesWebSocketConnection>()
.unwrap();
module
.add_function(
wrap_pyfunction!(crate::responses_websocket_decline, &module).unwrap(),
)
.unwrap();
let locals = crate::marshal::request_fixtures(py);
locals.set_item("routes", module).unwrap();
py.run(
c"
for route, provider in (
('messages', 'anthropic'),
('transcription', 'bedrock'),
('ocr', 'mistral'),
('responses_websocket', 'openai'),
):
decline = getattr(routes, route + '_decline')
assert decline('model', provider) is None, route
for flag in ('stream', 'has_agentic_hook', 'has_custom_client'):
assert decline('model', provider, **{flag: True}) is not None, (route, flag)
reason = decline('model', 'unsupported-native-provider')
assert reason is not None, route
request = Request(
messages=[], body={}, audio={}, document={}, optional_params={},
url='invalid-url-must-not-be-used',
options=Options(custom_llm_provider='unsupported-native-provider'),
)
functions = (
(routes.ResponsesWebSocketConnection.connect,)
if route == 'responses_websocket'
else (getattr(routes, route), getattr(routes, 'a' + route))
)
for execute in functions:
try:
execute(request, context=context)
except Exception as error:
assert type(error).__name__ == 'RustBridgeDeclined', (route, error)
assert str(error) == reason, (route, reason, error)
else:
raise AssertionError('unsupported request reached provider execution')
assert routes.ocr_decline('model', 'mistral', request_format='native') is not None
assert routes.ocr_decline('model', 'mistral', request_format='litellm') is None
",
Some(&locals),
Some(&locals),
)
.expect("acceptance must match execution eligibility without I/O");
});
}
#[test]
fn generated_routes_execute_sync_and_async_contracts() {
Python::initialize();

View file

@ -20,6 +20,20 @@ fn prepare_messages(
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<impl Future<Output = Result<AnthropicMessagesResponse, Error>> + Send + 'static> {
if let Some(reason) = messages_decline(
&input.model,
input.options.provider("anthropic"),
input
.body
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(false),
false,
false,
input.body.get("response_format").and_then(Value::as_str),
) {
return Err(crate::errors::RustBridgeDeclined::new_err(reason));
}
let context: LiteLlmRequestContext = context.into();
let body = required_value("body", input.body, Value::is_object, "dict")?;
Ok(async move {
@ -35,10 +49,30 @@ fn prepare_messages(
})
}
#[pyfunction]
#[pyo3(signature = (_model, custom_llm_provider, *, stream=false, has_agentic_hook=false, has_custom_client=false, request_format=None))]
fn messages_decline(
_model: &str,
custom_llm_provider: &str,
stream: bool,
has_agentic_hook: bool,
has_custom_client: bool,
request_format: Option<&str>,
) -> Option<String> {
super::definition::request_decline(
litellm_core::messages::messages_provider_supported(custom_llm_provider),
stream,
has_agentic_hook,
has_custom_client,
request_format,
)
}
bridge_route! {
sync = messages,
asynchronous = amessages,
request = MessagesInputs,
prepare = prepare_messages,
errors = core_error_to_pyerr,
extra = [messages_decline],
}

View file

@ -1,7 +1,7 @@
use pyo3::prelude::*;
#[macro_use]
mod definition;
pub(crate) mod definition;
#[cfg(feature = "trace-parity")]
mod gateway_messages;

View file

@ -23,6 +23,23 @@ fn prepare_ocr(
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
if let Some(reason) = ocr_decline(
&input.model,
input.options.provider("mistral"),
input
.optional_params
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(false),
false,
false,
input
.optional_params
.get("req_format")
.and_then(Value::as_str),
) {
return Err(crate::errors::RustBridgeDeclined::new_err(reason));
}
let context: LiteLlmRequestContext = context.into();
let document = input.document;
Ok(async move {
@ -43,10 +60,30 @@ fn prepare_ocr(
})
}
#[pyfunction]
#[pyo3(signature = (model, custom_llm_provider, *, stream=false, has_agentic_hook=false, has_custom_client=false, request_format=None))]
fn ocr_decline(
model: &str,
custom_llm_provider: &str,
stream: bool,
has_agentic_hook: bool,
has_custom_client: bool,
request_format: Option<&str>,
) -> Option<String> {
super::definition::request_decline(
litellm_ai_gateway::io::ocr::ocr_provider_supported(model, custom_llm_provider),
stream,
has_agentic_hook,
has_custom_client,
request_format,
)
}
bridge_route! {
sync = ocr,
asynchronous = aocr,
request = OcrInputs,
prepare = prepare_ocr,
errors = ocr_error_to_pyerr,
extra = [ocr_decline],
}

View file

@ -10,7 +10,8 @@ import subprocess
import sys
import time
import traceback
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Generator, Iterator, Mapping, Sequence
from contextlib import contextmanager
from datetime import datetime as dt_object
from functools import lru_cache
from types import MappingProxyType, TracebackType
@ -459,6 +460,7 @@ class Logging(LiteLLMLoggingBaseClass):
log_raw_request_response: bool = False,
supports_correlation_logging: bool = True,
):
self._suppress_next_pre_call: bool = False
_input: Final[str | None] = messages # save original value of messages
if messages is not None:
if isinstance(messages, str):
@ -1188,7 +1190,19 @@ class Logging(LiteLLMLoggingBaseClass):
additional_args.get("api_base", "")
)
@contextmanager
def suppress_next_pre_call(self) -> Generator[None]:
previous: Final = self._suppress_next_pre_call
self._suppress_next_pre_call = True
try:
yield
finally:
self._suppress_next_pre_call = previous
def pre_call(self, input, api_key, model=None, additional_args={}):
if self._suppress_next_pre_call:
self._suppress_next_pre_call = False
return
# Log the exact input to the LLM API
try:
self._pre_call(

View file

@ -25,11 +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.rust_bridge.dispatch import anative_first, native_first, provider_errors
from litellm.rust_bridge.request import anthropic_options, request_context
from litellm.rust_bridge.runtime import DispatchResult
from litellm.types.llms.anthropic import (
ContentBlockDelta,
ContentBlockStart,
@ -371,7 +366,12 @@ class AnthropicChatCompletion(BaseLLM):
if config is None:
raise ValueError(f"Provider config not found for model: {model} and provider: {custom_llm_provider}")
def prepare_python() -> tuple[dict[str, str], dict[str, object]]: # mutable-ok: stream mutates data
def build_request() -> tuple[dict, dict]: # mutable-ok: rewritten in place downstream
"""Translate the request the Python way, returning `(headers, data)`.
The pair stays mutable because the streaming path rewrites it in
place (`data["stream"] = True`) before sending.
"""
request_data: Final = config.transform_request(
model=model,
messages=messages,
@ -379,119 +379,32 @@ class AnthropicChatCompletion(BaseLLM):
litellm_params=litellm_params,
headers=headers,
)
python_headers, data = update_request_with_filtered_beta(
return update_request_with_filtered_beta(
headers=headers,
request_data=request_data,
provider=custom_llm_provider,
)
## LOGGING
# Reaching here with `serves_via_rust` set means the Rust attempt
# declined at call time, before the provider was called, and already
# logged this request. That is the same attempt continuing.
if not serves_via_rust:
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": python_headers,
},
)
print_verbose(f"_is_function_call: {_is_function_call}")
return python_headers, data
headers, data = build_request()
# 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,
)
rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict
"complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent
"model": model,
"messages": messages,
**rust_optional_params,
},
"api_base": api_base,
"headers": headers,
}
if serves_via_rust:
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,
## LOGGING
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args=rust_logging_args,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": headers,
},
)
rust_context: Final = request_context(
logging_obj=logging_obj,
request_model=logging_obj.model,
litellm_params=litellm_params,
)
def native_completion() -> DispatchResult[ModelResponse]:
return 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,
anthropic=anthropic_options(litellm_params),
stream=bool(stream),
has_custom_client=client is not None,
eligible=serves_via_rust,
context=rust_context,
)
async def native_acompletion() -> DispatchResult[ModelResponse]:
return await rust_chat_completions_bridge.achat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
anthropic=anthropic_options(litellm_params),
stream=bool(stream),
has_custom_client=client is not None,
eligible=serves_via_rust,
context=rust_context,
)
@anative_first(
native=native_acompletion,
route="chat_completions",
errors=lambda: provider_errors(custom_llm_provider or "", model),
)
async def execute_async() -> ModelResponse | CustomStreamWrapper:
headers, data = prepare_python()
print_verbose(f"_is_function_call: {_is_function_call}")
if acompletion is True:
if (
stream is True
): # if function call - fake the streaming (need complete blocks for output parsing in openai format)
print_verbose("makes async anthropic streaming POST request")
data["stream"] = stream
return await self.acompletion_stream_function(
return self.acompletion_stream_function(
model=model,
messages=messages,
data=data,
@ -513,7 +426,7 @@ class AnthropicChatCompletion(BaseLLM):
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
)
else:
return await self.acompletion_function(
return self.acompletion_function(
model=model,
messages=messages,
data=data,
@ -535,14 +448,7 @@ class AnthropicChatCompletion(BaseLLM):
json_mode=json_mode,
timeout=timeout,
)
@native_first(
native=native_completion,
route="chat_completions",
errors=lambda: provider_errors(custom_llm_provider or "", model),
)
def execute_sync() -> ModelResponse | CustomStreamWrapper:
headers, data = prepare_python()
else:
## COMPLETION CALL
if (
stream is True
@ -574,12 +480,13 @@ class AnthropicChatCompletion(BaseLLM):
)
else:
python_client: Final = (
client if isinstance(client, HTTPHandler) else _get_httpx_client(params={"timeout": timeout})
)
if client is None or not isinstance(client, HTTPHandler):
client = _get_httpx_client(params={"timeout": timeout})
else:
client = client
try:
response: Final = python_client.post(
response: Final = client.post(
api_base,
headers=headers,
data=json.dumps(data),
@ -600,21 +507,20 @@ class AnthropicChatCompletion(BaseLLM):
status_code=status_code,
headers=error_headers,
)
return config.transform_response(
model=model,
raw_response=response,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
request_data=data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
json_mode=json_mode,
)
return execute_async() if acompletion else execute_sync()
return config.transform_response(
model=model,
raw_response=response,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
request_data=data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
json_mode=json_mode,
)
def embedding(self):
# logic for parsing in - calling - parsing out model embedding calls

View file

@ -315,6 +315,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
def _resolved_provider(self) -> str:
return self.custom_llm_provider or "anthropic"
@classmethod
def get_config_for_model(cls, model: str) -> dict[str, object]:
from pydantic import TypeAdapter
return TypeAdapter(dict[str, object]).validate_python(cls.get_config(model=model))
@classmethod
def get_config(cls, *, model: str | None = None):
config: Final = super().get_config()

View file

@ -559,27 +559,54 @@ def anthropic_messages_handler(
**kwargs,
)
anthropic_messages_provider_config: BaseAnthropicMessagesConfig | None = None
def python_fallback():
anthropic_messages_provider_config: BaseAnthropicMessagesConfig | None = None
if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]:
anthropic_messages_provider_config = ProviderConfigManager.get_provider_anthropic_messages_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
if anthropic_messages_provider_config is None and _deployment_passes_through_anthropic_messages(
kwargs.get("model_info")
):
from litellm.llms.openai_like.messages.transformation import (
OpenAILikeAnthropicMessagesConfig,
)
if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]:
anthropic_messages_provider_config = ProviderConfigManager.get_provider_anthropic_messages_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
if anthropic_messages_provider_config is None and _deployment_passes_through_anthropic_messages(
kwargs.get("model_info")
):
from litellm.llms.openai_like.messages.transformation import (
OpenAILikeAnthropicMessagesConfig,
)
anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig(
cache_control_ttl=_deployment_supports_cache_control_ttl(kwargs.get("model_info")),
)
if anthropic_messages_provider_config is None:
# Route to Responses API for OpenAI / Azure, chat/completions for everything else.
if _should_route_to_responses_api(custom_llm_provider, original_model, model):
return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(
anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig(
cache_control_ttl=_deployment_supports_cache_control_ttl(kwargs.get("model_info")),
)
if anthropic_messages_provider_config is None:
# Route to Responses API for OpenAI / Azure, chat/completions for everything else.
if _should_route_to_responses_api(custom_llm_provider, original_model, model):
return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(
max_tokens=max_tokens,
messages=messages,
model=original_model,
metadata=metadata,
stop_sequences=stop_sequences,
stream=stream,
system=system,
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
tools=tools,
top_k=top_k,
top_p=top_p,
_is_async=is_async,
api_key=api_key,
api_base=api_base,
client=client,
custom_llm_provider=custom_llm_provider,
**kwargs,
)
# The in-gateway context_management polyfill runs inside
# ``async_anthropic_messages_handler`` so it can ``await`` the
# summarization model for ``compact_20260112``. ``context_management``
# is passed through as a regular kwarg.
return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler(
max_tokens=max_tokens,
messages=messages,
model=original_model,
@ -601,66 +628,87 @@ def anthropic_messages_handler(
**kwargs,
)
# The in-gateway context_management polyfill runs inside
# ``async_anthropic_messages_handler`` so it can ``await`` the
# summarization model for ``compact_20260112``. ``context_management``
# is passed through as a regular kwarg.
return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler(
max_tokens=max_tokens,
messages=messages,
model=original_model,
metadata=metadata,
stop_sequences=stop_sequences,
stream=stream,
system=system,
temperature=temperature,
thinking=thinking,
tool_choice=tool_choice,
tools=tools,
top_k=top_k,
top_p=top_p,
if custom_llm_provider is None:
raise ValueError(
f"custom_llm_provider is required for Anthropic messages, passed in model={model}, custom_llm_provider={custom_llm_provider}"
)
local_vars.update(kwargs)
anthropic_messages_optional_request_params: Final = (
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params=local_vars,
model=model,
drop_params=litellm_params.get("drop_params") is True,
custom_llm_provider=custom_llm_provider,
)
)
if is_reasoning_auto_summary_enabled():
thinking_param: Final = anthropic_messages_optional_request_params.get("thinking")
if isinstance(thinking_param, dict) and thinking_param.get("type") != "disabled":
anthropic_messages_optional_request_params["thinking"] = {
**thinking_param,
"display": "summarized",
}
return base_llm_http_handler.anthropic_messages_handler(
model=model,
messages=strip_provider_specific_fields_from_anthropic_messages(messages),
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=dict(anthropic_messages_optional_request_params),
_is_async=is_async,
api_key=api_key,
api_base=api_base,
client=client,
custom_llm_provider=custom_llm_provider,
**kwargs,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
api_key=api_key,
api_base=api_base,
stream=stream,
kwargs=kwargs,
)
if custom_llm_provider is None:
raise ValueError(
f"custom_llm_provider is required for Anthropic messages, passed in model={model}, custom_llm_provider={custom_llm_provider}"
)
from pydantic import TypeAdapter
local_vars.update(kwargs)
anthropic_messages_optional_request_params: Final = (
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params=local_vars,
from litellm.rust_bridge.messages import dispatch_messages
return dispatch_messages(
model=model,
provider=custom_llm_provider,
messages=TypeAdapter(list[dict[str, object]]).validate_python(
strip_provider_specific_fields_from_anthropic_messages(messages)
),
body=lambda: _native_messages_body(
params={**local_vars, **kwargs},
model=model,
drop_params=litellm_params.get("drop_params") is True,
custom_llm_provider=custom_llm_provider,
provider=custom_llm_provider,
),
params=litellm_params,
logging=litellm_logging_obj,
api_key=api_key or dynamic_api_key,
api_base=api_base or dynamic_api_base,
stream=bool(stream),
asynchronous=bool(is_async),
has_custom_client=client is not None,
fallback=python_fallback,
)
def _native_messages_body(params: dict[str, object], model: str, drop_params: bool, provider: str) -> dict[str, object]:
from pydantic import TypeAdapter
requested: Final = TypeAdapter(dict[str, object]).validate_python(
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params=params,
model=model,
drop_params=drop_params,
custom_llm_provider=provider,
)
)
if is_reasoning_auto_summary_enabled():
thinking_param: Final = anthropic_messages_optional_request_params.get("thinking")
if isinstance(thinking_param, dict) and thinking_param.get("type") != "disabled":
anthropic_messages_optional_request_params["thinking"] = {
**thinking_param,
"display": "summarized",
}
return base_llm_http_handler.anthropic_messages_handler(
model=model,
messages=strip_provider_specific_fields_from_anthropic_messages(messages),
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=dict(anthropic_messages_optional_request_params),
_is_async=is_async,
client=client,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
api_key=api_key,
api_base=api_base,
stream=stream,
kwargs=kwargs,
)
thinking_param: Final = requested.get("thinking")
if (
is_reasoning_auto_summary_enabled()
and isinstance(thinking_param, dict)
and thinking_param.get("type") != "disabled"
):
return {**requested, "thinking": {**thinking_param, "display": "summarized"}}
return requested

View file

@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Iterator
from typing import TYPE_CHECKING, Any, Final, Union
import httpx
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from litellm.constants import DEFAULT_MAX_TOKENS, RESPONSE_FORMAT_TOOL_NAME
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@ -70,6 +70,10 @@ class BaseConfig(ABC):
def __init__(self):
pass
@classmethod
def get_config_for_model(cls, model: str) -> dict[str, object]:
return TypeAdapter(dict[str, object]).validate_python(cls.get_config())
@classmethod
def get_config(cls):
return {

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,11 +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.rust_bridge.dispatch import anative_first, native_first, provider_errors
from litellm.rust_bridge.request import bedrock_options, request_context
from litellm.rust_bridge.runtime import DispatchResult
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
@ -29,22 +22,6 @@ from ..common_utils import BedrockError, _get_all_bedrock_regions
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,
@ -215,9 +192,6 @@ class BedrockConverseLLM(BaseAWSLLM):
)
## LOGGING
# The Rust path already logged this request's pre_call before handing
# it here, and it only declines before the provider is called, so this
# is the same attempt continuing rather than a second one.
if not skip_pre_call_logging:
logging_obj.pre_call(
input=messages,
@ -390,94 +364,12 @@ 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,
)
rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict
"complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent
"messages": messages,
**optional_params,
},
"api_base": proxy_endpoint_url,
"headers": headers,
}
if serves_via_rust:
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,
)
rust_context: Final = request_context(
logging_obj=logging_obj,
request_model=logging_obj.model,
litellm_params=litellm_params,
)
def native_completion() -> DispatchResult[ModelResponse]:
return 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,
bedrock=bedrock_options(rust_optional_params),
stream=bool(stream),
has_custom_client=client is not None,
eligible=serves_via_rust,
context=rust_context,
)
async def native_acompletion() -> DispatchResult[ModelResponse]:
return await rust_chat_completions_bridge.achat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=proxy_endpoint_url,
custom_llm_provider="bedrock",
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
bedrock=bedrock_options(rust_optional_params),
stream=bool(stream),
has_custom_client=client is not None,
eligible=serves_via_rust,
context=rust_context,
)
@anative_first(
native=native_acompletion,
route="chat_completions",
errors=lambda: provider_errors("bedrock", model),
)
async def execute_async() -> ModelResponse | CustomStreamWrapper:
python_client: Final = None if isinstance(client, HTTPHandler) else client
### ROUTING (ASYNC, STREAMING, SYNC)
if acompletion:
if isinstance(client, HTTPHandler):
client = None
if stream is True:
return await self.async_streaming(
return self.async_streaming(
model=model,
messages=messages,
api_base=proxy_endpoint_url,
@ -490,7 +382,7 @@ class BedrockConverseLLM(BaseAWSLLM):
logger_fn=logger_fn,
headers=headers,
timeout=timeout,
client=python_client,
client=client,
json_mode=json_mode,
fake_stream=fake_stream,
credentials=credentials,
@ -498,7 +390,7 @@ class BedrockConverseLLM(BaseAWSLLM):
stream_chunk_size=stream_chunk_size,
)
### ASYNC COMPLETION
return await self.async_completion(
return self.async_completion(
model=model,
messages=messages,
api_base=proxy_endpoint_url,
@ -511,112 +403,102 @@ class BedrockConverseLLM(BaseAWSLLM):
logger_fn=logger_fn,
headers=headers,
timeout=timeout,
client=python_client,
client=client,
credentials=credentials,
api_key=api_key,
skip_pre_call_logging=serves_via_rust,
)
@native_first(
native=native_completion,
route="chat_completions",
errors=lambda: provider_errors("bedrock", model),
## TRANSFORMATION ##
_data: Final = litellm.AmazonConverseConfig()._transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=extra_headers,
)
def execute_sync() -> ModelResponse | CustomStreamWrapper:
## TRANSFORMATION ##
data: Final = json.dumps(_data)
_data: Final = litellm.AmazonConverseConfig()._transform_request(
prepped: Final = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=proxy_endpoint_url,
data=data,
headers=headers,
api_key=api_key,
)
## LOGGING
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:
if isinstance(timeout, float) or isinstance(timeout, int):
timeout = httpx.Timeout(timeout)
_params["timeout"] = timeout
client = _get_httpx_client(_params)
else:
client = client
if stream is not None and stream is True:
completion_stream, response_headers = make_sync_call(
client=(client if client is not None and isinstance(client, HTTPHandler) else None),
api_base=proxy_endpoint_url,
headers=prepped.headers,
data=data,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=extra_headers,
)
data: Final = json.dumps(_data)
prepped: Final = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=proxy_endpoint_url,
data=data,
headers=headers,
api_key=api_key,
)
## LOGGING
# Reaching here with `serves_via_rust` set means the synchronous Rust
# attempt declined at call time, before the provider was called, and
# already logged this request. That is the same attempt continuing.
if not serves_via_rust:
logging_obj.pre_call(
input=messages,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": proxy_endpoint_url,
"headers": prepped.headers,
},
)
resolved_timeout: Final = httpx.Timeout(timeout) if isinstance(timeout, (float, int)) else timeout
python_client: Final = (
_get_httpx_client({"timeout": resolved_timeout} if resolved_timeout is not None else None)
if client is None or isinstance(client, AsyncHTTPHandler)
else client
)
if stream is not None and stream is True:
completion_stream, response_headers = make_sync_call(
client=python_client,
api_base=proxy_endpoint_url,
headers=prepped.headers,
data=data,
model=model,
messages=messages,
logging_obj=logging_obj,
json_mode=json_mode,
fake_stream=fake_stream,
stream_chunk_size=stream_chunk_size,
)
streaming_response: Final = CustomStreamWrapper(
completion_stream=completion_stream,
model=model,
custom_llm_provider="bedrock",
logging_obj=logging_obj,
_response_headers=response_headers,
)
return streaming_response
### COMPLETION
try:
response: Final = python_client.post(
url=proxy_endpoint_url,
headers=prepped.headers,
data=data,
logging_obj=logging_obj,
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
sync_transformed_response: Final = litellm.AmazonConverseConfig()._transform_response(
model=model,
response=response,
model_response=model_response,
stream=stream if isinstance(stream, bool) else False,
logging_obj=logging_obj,
api_key="",
data=data,
messages=messages,
optional_params=optional_params,
encoding=encoding,
json_mode=json_mode,
fake_stream=fake_stream,
stream_chunk_size=stream_chunk_size,
)
streaming_response: Final = CustomStreamWrapper(
completion_stream=completion_stream,
model=model,
custom_llm_provider="bedrock",
logging_obj=logging_obj,
_response_headers=response_headers,
)
sync_transformed_response.set_provider_response_headers(response.headers)
return sync_transformed_response
return execute_async() if acompletion else execute_sync()
return streaming_response
### COMPLETION
try:
response: Final = client.post(
url=proxy_endpoint_url,
headers=prepped.headers,
data=data,
logging_obj=logging_obj,
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
sync_transformed_response: Final = litellm.AmazonConverseConfig()._transform_response(
model=model,
response=response,
model_response=model_response,
stream=stream if isinstance(stream, bool) else False,
logging_obj=logging_obj,
api_key="",
data=data,
messages=messages,
optional_params=optional_params,
encoding=encoding,
)
sync_transformed_response.set_provider_response_headers(response.headers)
return sync_transformed_response

View file

@ -1,8 +1,7 @@
import asyncio
import json
import ssl
from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from functools import lru_cache
from types import MappingProxyType, ModuleType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints
@ -92,8 +91,6 @@ from litellm.responses.streaming_iterator import (
ResponsesWebSocketStreaming,
SyncResponsesAPIStreamingIterator,
)
from litellm.rust_bridge.dispatch import anative_context, anative_first, provider_errors
from litellm.rust_bridge.runtime import DispatchResult, NativeSkipped, NativeSkipReason, adapt_result
from litellm.types.containers.main import (
ContainerFileListResponse,
ContainerListResponse,
@ -162,15 +159,6 @@ from litellm.utils import (
async_pre_call_deployment_hook,
)
def _rust_responses_websocket_enabled(
custom_llm_provider: str | None,
) -> bool:
from litellm.rust_bridge.configuration import rust_enabled
return custom_llm_provider == "openai" and rust_enabled()
from .http_handler import get_shared_realtime_ssl_context
if TYPE_CHECKING:
@ -183,9 +171,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,
@ -245,7 +230,7 @@ def _responses_api_optional_request_param_names() -> frozenset[str]:
return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys())
def _custom_logger_callbacks(logging_obj: LiteLLMLoggingObj) -> list["CustomLogger"]:
def _custom_logger_callbacks(logging_obj: LiteLLMLoggingObj | None) -> list["CustomLogger"]:
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import (
get_custom_logger_compatible_class,
@ -2055,7 +2040,7 @@ class BaseLLMHTTPHandler:
raise RuntimeError("unreachable: anthropic messages HTTP retry loop exited without return")
@staticmethod
def _resolve_anthropic_messages_timeout(
def resolve_anthropic_messages_timeout(
litellm_params: GenericLiteLLMParams,
stream: bool,
custom_llm_provider: str,
@ -2227,116 +2212,86 @@ class BaseLLMHTTPHandler:
},
)
async def native_messages() -> DispatchResult[AnthropicMessagesResponse | AsyncIterator[object]]:
result: Final = await self._attempt_rust_anthropic_messages(
custom_llm_provider=custom_llm_provider,
response: Final = await self._async_post_anthropic_messages_with_http_error_retry(
async_httpx_client=async_httpx_client,
request_url=request_url,
headers=headers,
signed_json_body=(signed_json_body if signed_json_body is not None else request_body_json),
request_body=request_body,
stream=stream or False,
logging_obj=logging_obj,
provider_config=anthropic_messages_provider_config,
litellm_params=litellm_params,
api_key=api_key,
model=model,
timeout=self.resolve_anthropic_messages_timeout(
litellm_params=litellm_params,
has_agentic_hook=self._has_agentic_completion_hook(logging_obj),
stream=bool(stream),
has_custom_client=client is not None,
model=model,
api_key=api_key,
api_base=api_base,
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,
),
logging_obj=logging_obj,
)
return adapt_result(result, self._rust_anthropic_messages_fake_stream) if stream else result
@anative_first(
native=native_messages, route="messages", errors=lambda: provider_errors(custom_llm_provider, model)
)
async def execute_messages() -> AnthropicMessagesResponse | AsyncIterator[object]:
response: Final = await self._async_post_anthropic_messages_with_http_error_retry(
async_httpx_client=async_httpx_client,
request_url=request_url,
headers=headers,
signed_json_body=(signed_json_body if signed_json_body is not None else request_body_json),
request_body=request_body,
stream=stream or False,
logging_obj=logging_obj,
provider_config=anthropic_messages_provider_config,
litellm_params=litellm_params,
api_key=api_key,
model=model,
timeout=self._resolve_anthropic_messages_timeout(
litellm_params=litellm_params,
stream=stream or False,
custom_llm_provider=custom_llm_provider,
),
custom_llm_provider=custom_llm_provider,
),
)
# used for logging + cost tracking
logging_obj.model_call_details["httpx_response"] = response
initial_response: AsyncIterator | AnthropicMessagesResponse
if stream:
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
AnthropicMessagesStreamingResponse,
anthropic_messages_stream_hidden_params,
)
# used for logging + cost tracking
logging_obj.model_call_details["httpx_response"] = response
completion_stream: Final = anthropic_messages_provider_config.get_async_streaming_response_iterator(
model=model,
httpx_response=response,
request_body=request_body,
litellm_logging_obj=logging_obj,
)
stream_hidden_params: Final = anthropic_messages_stream_hidden_params(response.headers)
initial_response: AsyncIterator | AnthropicMessagesResponse
if stream:
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
AnthropicMessagesStreamingResponse,
anthropic_messages_stream_hidden_params,
)
completion_stream: Final = anthropic_messages_provider_config.get_async_streaming_response_iterator(
model=model,
httpx_response=response,
request_body=request_body,
litellm_logging_obj=logging_obj,
)
stream_hidden_params: Final = anthropic_messages_stream_hidden_params(response.headers)
if not self._has_agentic_completion_hook(logging_obj):
# No callback overrides async_should_run_agentic_loop, so the
# agentic wrapper's only effect would be buffering every chunk
# and rebuilding the response from SSE at end-of-stream to call
# hooks that all return (False, {}). Stream through directly and
# skip that per-chunk + end-of-stream overhead.
return AnthropicMessagesStreamingResponse(
completion_stream=completion_stream,
hidden_params=stream_hidden_params,
)
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
AgenticAnthropicStreamingIterator,
)
held_back_tool_names: Final = self._server_fulfilled_tools_in_request(
logging_obj=logging_obj,
tools=anthropic_messages_optional_request_params.get("tools"),
)
initial_response = AgenticAnthropicStreamingIterator(
completion_stream=completion_stream,
http_handler=self,
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,
kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
hold_back=bool(held_back_tool_names),
server_fulfilled_tool_names=held_back_tool_names,
)
if not self.has_agentic_completion_hook(logging_obj):
# No callback overrides async_should_run_agentic_loop, so the
# agentic wrapper's only effect would be buffering every chunk
# and rebuilding the response from SSE at end-of-stream to call
# hooks that all return (False, {}). Stream through directly and
# skip that per-chunk + end-of-stream overhead.
return AnthropicMessagesStreamingResponse(
completion_stream=initial_response,
completion_stream=completion_stream,
hidden_params=stream_hidden_params,
)
else:
initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
)
return initial_response
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
AgenticAnthropicStreamingIterator,
)
held_back_tool_names: Final = self._server_fulfilled_tools_in_request(
logging_obj=logging_obj,
tools=anthropic_messages_optional_request_params.get("tools"),
)
initial_response = AgenticAnthropicStreamingIterator(
completion_stream=completion_stream,
http_handler=self,
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,
kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
hold_back=bool(held_back_tool_names),
server_fulfilled_tool_names=held_back_tool_names,
)
return AnthropicMessagesStreamingResponse(
completion_stream=initial_response,
hidden_params=stream_hidden_params,
)
else:
initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
)
initial_response: Final = await execute_messages()
if stream:
return initial_response
return await self._finalize_anthropic_messages_response(
initial_response=initial_response,
model=model,
@ -2385,80 +2340,6 @@ class BaseLLMHTTPHandler:
"anthropic_messages",
)
@staticmethod
async def _attempt_rust_anthropic_messages(
*,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
has_agentic_hook: bool,
stream: bool,
has_custom_client: bool,
model: str,
api_key: str | None,
api_base: str | None,
headers: dict,
request_body: dict,
timeout: float | httpx.Timeout | None,
logging_obj: LiteLLMLoggingObj | None = None,
) -> DispatchResult[AnthropicMessagesResponse]:
if custom_llm_provider not in ("azure_ai", "anthropic"):
return NativeSkipped(NativeSkipReason.INELIGIBLE)
from litellm.rust_bridge.configuration import rust_enabled
if not rust_enabled():
return NativeSkipped(NativeSkipReason.DISABLED)
if has_agentic_hook:
return NativeSkipped(NativeSkipReason.INELIGIBLE)
from litellm.rust_bridge import messages as rust_messages_bridge
from litellm.rust_bridge.request import request_context
upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"}
result: 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,
stream=stream,
has_custom_client=has_custom_client,
has_agentic_hook=has_agentic_hook,
context=request_context(
logging_obj=logging_obj,
request_model=logging_obj.model if logging_obj is not None else model,
litellm_params=litellm_params.model_dump(),
),
)
def adapt(rust_response: dict[str, object]) -> AnthropicMessagesResponse:
return cast(
AnthropicMessagesResponse,
{**rust_response, "_hidden_params": {"additional_headers": {"x-litellm-rust": "true"}}},
)
return adapt_result(result, adapt)
@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,
@ -2767,7 +2648,7 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
)
if self._has_agentic_completion_hook(logging_obj):
if self.has_agentic_completion_hook(logging_obj):
agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place
final_response: Final = run_async_function(
self._call_agentic_completion_hooks,
@ -5158,7 +5039,7 @@ class BaseLLMHTTPHandler:
return depth, max_loops, fingerprints
@staticmethod
def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool:
def has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj | None) -> bool:
"""
True if any registered callback actually overrides
``async_should_run_agentic_loop`` (the gate every agentic hook goes
@ -6516,41 +6397,21 @@ class BaseLLMHTTPHandler:
},
)
from litellm.rust_bridge import responses_websocket as rust_responses_websocket
from litellm.rust_bridge.request import request_context
from litellm.rust_bridge.responses_websocket import open_connection
async def attempt_connection() -> DispatchResult[
AbstractAsyncContextManager[rust_responses_websocket.ConnectionAdapter]
]:
if not _rust_responses_websocket_enabled(custom_llm_provider):
return NativeSkipped(NativeSkipReason.INELIGIBLE)
return await rust_responses_websocket.managed_connect(
url=ws_url,
headers={str(key): str(value) for key, value in headers.items()},
timeout=timeout,
context=request_context(
logging_obj=logging_obj,
request_model=logging_obj.model,
litellm_params=litellm_params.model_dump(),
),
)
@anative_context(
native=attempt_connection,
route="responses_websocket",
errors=lambda: provider_errors("openai", "responses websocket"),
)
@asynccontextmanager
async def _backend_connection() -> AsyncGenerator[ClientConnection, None]:
async with websockets.connect(
async with open_connection(
url=ws_url,
headers={str(key): str(value) for key, value in headers.items()},
timeout=timeout,
model=model,
provider=custom_llm_provider,
fallback=lambda: websockets.connect(
ws_url,
additional_headers=headers,
max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
ssl=ssl_context,
) as backend:
yield backend
async with _backend_connection() as backend_ws:
),
) as backend_ws:
_request_data: Final[dict[str, object]] = {}
if litellm_metadata:
_request_data["litellm_metadata"] = litellm_metadata

View file

@ -19,7 +19,7 @@ import random
import sys
import time
import traceback
from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping, Sequence
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence
from concurrent import futures
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from copy import deepcopy
@ -4958,6 +4958,183 @@ def _complete_langflow(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe
)
_PYTHON_PRIMARY_COMPLETION_HANDLERS: Final[
Mapping[str, Callable[[_CompletionDispatchContext], _CompletionDispatchResult]]
] = MappingProxyType(
{
"azure": _complete_azure,
"azure_text": _complete_azure_text,
"deepseek": _complete_deepseek,
"azure_ai": _complete_azure_ai,
}
)
_PYTHON_COMPATIBLE_COMPLETION_HANDLERS: Final[
Mapping[str, Callable[[_CompletionDispatchContext], _CompletionDispatchResult]]
] = MappingProxyType(
{
"fireworks_ai": _complete_fireworks_ai,
"together_ai": _complete_together_ai,
"heroku": _complete_heroku,
"ragflow": _complete_ragflow,
"xai": _complete_xai,
"groq": _complete_groq,
"bedrock_mantle": _complete_bedrock_mantle,
"a2a": _complete_a2a,
"gigachat": _complete_gigachat,
"sap": _complete_sap,
"aiohttp_openai": _complete_aiohttp_openai,
"cometapi": _complete_cometapi,
"minimax": _complete_minimax,
"hosted_vllm": _complete_hosted_vllm,
}
)
_PYTHON_LEGACY_COMPLETION_HANDLERS: Final[
Mapping[str, Callable[[_CompletionDispatchContext], _CompletionDispatchResult]]
] = MappingProxyType(
{
"anthropic_text": _complete_anthropic_text,
"anthropic": _complete_anthropic,
"nlp_cloud": _complete_nlp_cloud,
"aleph_alpha": _complete_aleph_alpha,
"cohere_chat": _complete_cohere_chat,
"cohere": _complete_cohere_chat,
"maritalk": _complete_maritalk,
"amazon_nova": _complete_amazon_nova,
"huggingface": _complete_huggingface,
"oci": _complete_oci,
"compactifai": _complete_compactifai,
"oobabooga": _complete_oobabooga,
"databricks": _complete_databricks,
"datarobot": _complete_datarobot,
"openrouter": _complete_openrouter,
"vercel_ai_gateway": _complete_vercel_ai_gateway,
}
)
_PYTHON_EXTENDED_COMPLETION_HANDLERS: Final[
Mapping[str, Callable[[_CompletionDispatchContext], _CompletionDispatchResult]]
] = MappingProxyType(
{
"vertex_ai_beta": _complete_vertex_ai_beta,
"gemini": _complete_vertex_ai_beta,
"vertex_ai": _complete_vertex_ai,
"predibase": _complete_predibase,
"text-completion-codestral": _complete_text_completion_codestral,
"text-completion-inception": _complete_text_completion_inception,
"sagemaker_chat": _complete_sagemaker_chat,
"sagemaker_nova": _complete_sagemaker_chat,
"sagemaker": _complete_sagemaker,
"bedrock": _complete_bedrock,
"watsonx": _complete_watsonx,
"watsonx_text": _complete_watsonx_text,
"vllm": _complete_vllm,
"ollama": _complete_ollama,
"ollama_chat": _complete_ollama_chat,
"triton": _complete_triton,
"cloudflare": _complete_cloudflare,
}
)
_PYTHON_ADDITIONAL_COMPLETION_HANDLERS: Final[
Mapping[str, Callable[[_CompletionDispatchContext], _CompletionDispatchResult]]
] = MappingProxyType(
{
"gradient_ai": _complete_gradient_ai,
"gdc": _complete_gdc,
"bytez": _complete_bytez,
"lemonade": _complete_lemonade,
}
)
_PYTHON_CUSTOM_COMPLETION_HANDLERS: Final[
Mapping[str, Callable[[_CompletionDispatchContext], _CompletionDispatchResult]]
] = MappingProxyType(
{
"langgraph": _complete_langgraph,
"langflow": _complete_langflow,
}
)
def _complete_python(_dispatch_ctx: _CompletionDispatchContext) -> _CompletionDispatchResult:
custom_llm_provider: Final = _dispatch_ctx.custom_llm_provider
model: Final = _dispatch_ctx.model
kwargs: Final = _dispatch_ctx.kwargs
if (primary_handler := _PYTHON_PRIMARY_COMPLETION_HANDLERS.get(custom_llm_provider)) is not None:
return primary_handler(_dispatch_ctx)
elif (
custom_llm_provider == "text-completion-openai"
or "ft:babbage-002" in model
or "ft:davinci-002" in model
or (
custom_llm_provider in litellm.openai_text_completion_compatible_providers
and kwargs.get("text_completion") is True
)
):
return _complete_text_completion_openai(_dispatch_ctx)
elif (compatible_handler := _PYTHON_COMPATIBLE_COMPLETION_HANDLERS.get(custom_llm_provider)) is not None:
return compatible_handler(_dispatch_ctx)
elif (
model in litellm.open_ai_chat_completion_models
and custom_llm_provider in (None, "openai")
or custom_llm_provider == "custom_openai"
or custom_llm_provider == "deepinfra"
or (custom_llm_provider == "perplexity")
or (custom_llm_provider == "nvidia_nim")
or (custom_llm_provider == "cerebras")
or (custom_llm_provider == "baseten")
or (custom_llm_provider == "sambanova")
or (custom_llm_provider == "volcengine")
or (custom_llm_provider == "anyscale")
or (custom_llm_provider == "openai")
or (custom_llm_provider == "nebius")
or (custom_llm_provider == "wandb")
or (custom_llm_provider == "clarifai")
or (custom_llm_provider in litellm.openai_compatible_providers)
or JSONProviderRegistry.exists(custom_llm_provider)
or ("ft:gpt-3.5-turbo" in model)
):
return _complete_custom_openai(_dispatch_ctx)
return _complete_python_provider(_dispatch_ctx)
def _complete_python_provider(_dispatch_ctx: _CompletionDispatchContext) -> _CompletionDispatchResult:
custom_llm_provider: Final = _dispatch_ctx.custom_llm_provider
model: Final = _dispatch_ctx.model
if custom_llm_provider == "mistral":
return _complete_mistral(_dispatch_ctx)
elif "replicate" in model or custom_llm_provider == "replicate" or model in litellm.replicate_models:
return _complete_replicate(_dispatch_ctx)
elif "clarifai" in model or custom_llm_provider == "clarifai" or model in litellm.clarifai_models:
return _complete_custom_openai(_dispatch_ctx)
elif (legacy_handler := _PYTHON_LEGACY_COMPLETION_HANDLERS.get(custom_llm_provider)) is not None:
return legacy_handler(_dispatch_ctx)
elif custom_llm_provider == "palm":
raise ValueError(
"Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en"
)
elif (extended_handler := _PYTHON_EXTENDED_COMPLETION_HANDLERS.get(custom_llm_provider)) is not None:
return extended_handler(_dispatch_ctx)
elif custom_llm_provider == "petals" or model in litellm.petals_models:
return _complete_petals(_dispatch_ctx)
elif custom_llm_provider == "snowflake" or model in litellm.snowflake_models:
return _complete_snowflake(_dispatch_ctx)
elif (additional_handler := _PYTHON_ADDITIONAL_COMPLETION_HANDLERS.get(custom_llm_provider)) is not None:
return additional_handler(_dispatch_ctx)
elif custom_llm_provider == "ovhcloud" or model in litellm.ovhcloud_models:
return _complete_ovhcloud(_dispatch_ctx)
elif custom_llm_provider == "custom":
return _complete_custom(_dispatch_ctx)
elif custom_llm_provider in litellm._custom_providers:
return _complete_custom_providers(_dispatch_ctx)
elif (custom_handler := _PYTHON_CUSTOM_COMPLETION_HANDLERS.get(custom_llm_provider)) is not None:
return custom_handler(_dispatch_ctx)
else:
raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider)
@tracer.wrap()
@client
def completion(
@ -5640,209 +5817,9 @@ def completion(
timeout=timeout,
top_p=top_p,
)
if custom_llm_provider == "azure":
# azure configs
## check dynamic params ##
response = _complete_azure(_dispatch_ctx)
elif custom_llm_provider == "azure_text":
# azure configs
response = _complete_azure_text(_dispatch_ctx)
elif custom_llm_provider == "deepseek":
## COMPLETION CALL
from litellm.rust_bridge.chat_completions import dispatch_completion
response = _complete_deepseek(_dispatch_ctx)
elif custom_llm_provider == "azure_ai":
response = _complete_azure_ai(_dispatch_ctx)
elif (
custom_llm_provider == "text-completion-openai"
or "ft:babbage-002" in model
or "ft:davinci-002" in model # support for finetuned completion models
or custom_llm_provider in litellm.openai_text_completion_compatible_providers
and kwargs.get("text_completion") is True
):
response = _complete_text_completion_openai(_dispatch_ctx)
elif custom_llm_provider == "fireworks_ai":
## COMPLETION CALL
response = _complete_fireworks_ai(_dispatch_ctx)
elif custom_llm_provider == "together_ai":
response = _complete_together_ai(_dispatch_ctx)
elif custom_llm_provider == "heroku":
response = _complete_heroku(_dispatch_ctx)
elif custom_llm_provider == "ragflow":
## COMPLETION CALL - RAGFlow uses HTTP handler to support custom URL paths
response = _complete_ragflow(_dispatch_ctx)
elif custom_llm_provider == "xai":
## COMPLETION CALL
response = _complete_xai(_dispatch_ctx)
elif custom_llm_provider == "groq":
response = _complete_groq(_dispatch_ctx)
elif custom_llm_provider == "bedrock_mantle":
response = _complete_bedrock_mantle(_dispatch_ctx)
elif custom_llm_provider == "a2a":
# A2A (Agent-to-Agent) Protocol
# Resolve agent configuration from registry if model format is "a2a/<agent-name>"
response = _complete_a2a(_dispatch_ctx)
elif custom_llm_provider == "gigachat":
# GigaChat - Sber AI's LLM (Russia)
response = _complete_gigachat(_dispatch_ctx)
elif custom_llm_provider == "sap":
response = _complete_sap(_dispatch_ctx)
elif custom_llm_provider == "aiohttp_openai":
# NEW aiohttp provider for 10-100x higher RPS
response = _complete_aiohttp_openai(_dispatch_ctx)
elif custom_llm_provider == "cometapi":
response = _complete_cometapi(_dispatch_ctx)
elif custom_llm_provider == "minimax":
response = _complete_minimax(_dispatch_ctx)
elif custom_llm_provider == "hosted_vllm":
response = _complete_hosted_vllm(_dispatch_ctx)
elif (
# A known OpenAI model name only decides the route when nothing else
# resolved a provider. get_llm_provider() already maps these names to
# "openai", so a different value here was asked for explicitly (or came
# from a register_model entry), and the provider config built for it
# would be handed to the OpenAI handler.
(model in litellm.open_ai_chat_completion_models and custom_llm_provider in (None, "openai"))
or custom_llm_provider == "custom_openai"
or custom_llm_provider == "deepinfra"
or custom_llm_provider == "perplexity"
or custom_llm_provider == "nvidia_nim"
or custom_llm_provider == "cerebras"
or custom_llm_provider == "baseten"
or custom_llm_provider == "sambanova"
or custom_llm_provider == "volcengine"
or custom_llm_provider == "anyscale"
or custom_llm_provider == "openai"
or custom_llm_provider == "nebius"
or custom_llm_provider == "wandb"
or custom_llm_provider == "clarifai"
or custom_llm_provider in litellm.openai_compatible_providers
or JSONProviderRegistry.exists(custom_llm_provider) # JSON-configured providers
or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo
): # allow user to make an openai call with a custom base
# note: if a user sets a custom base - we should ensure this works
# allow for the setting of dynamic and stateful api-bases
response = _complete_custom_openai(_dispatch_ctx)
elif custom_llm_provider == "mistral":
response = _complete_mistral(_dispatch_ctx)
elif "replicate" in model or custom_llm_provider == "replicate" or model in litellm.replicate_models:
# Setting the relevant API KEY for replicate, replicate defaults to using os.environ.get("REPLICATE_API_TOKEN")
response = _complete_replicate(_dispatch_ctx)
elif "clarifai" in model or custom_llm_provider == "clarifai" or model in litellm.clarifai_models:
pass # Deprecated - handled in the openai compatible provider section above
elif custom_llm_provider == "anthropic_text":
response = _complete_anthropic_text(_dispatch_ctx)
elif custom_llm_provider == "anthropic":
response = _complete_anthropic(_dispatch_ctx)
elif custom_llm_provider == "nlp_cloud":
response = _complete_nlp_cloud(_dispatch_ctx)
elif custom_llm_provider == "aleph_alpha":
response = _complete_aleph_alpha(_dispatch_ctx)
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
response = _complete_cohere_chat(_dispatch_ctx)
elif custom_llm_provider == "maritalk":
response = _complete_maritalk(_dispatch_ctx)
elif custom_llm_provider == "amazon_nova":
response = _complete_amazon_nova(_dispatch_ctx)
elif custom_llm_provider == "huggingface":
response = _complete_huggingface(_dispatch_ctx)
elif custom_llm_provider == "oci":
response = _complete_oci(_dispatch_ctx)
elif custom_llm_provider == "compactifai":
response = _complete_compactifai(_dispatch_ctx)
elif custom_llm_provider == "oobabooga":
response = _complete_oobabooga(_dispatch_ctx)
elif custom_llm_provider == "databricks":
response = _complete_databricks(_dispatch_ctx)
elif custom_llm_provider == "datarobot":
response = _complete_datarobot(_dispatch_ctx)
elif custom_llm_provider == "openrouter":
response = _complete_openrouter(_dispatch_ctx)
elif custom_llm_provider == "vercel_ai_gateway":
response = _complete_vercel_ai_gateway(_dispatch_ctx)
elif custom_llm_provider == "palm":
raise ValueError(
"Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en"
)
elif custom_llm_provider == "vertex_ai_beta" or custom_llm_provider == "gemini":
response = _complete_vertex_ai_beta(_dispatch_ctx)
elif custom_llm_provider == "vertex_ai":
response = _complete_vertex_ai(_dispatch_ctx)
elif custom_llm_provider == "predibase":
response = _complete_predibase(_dispatch_ctx)
elif custom_llm_provider == "text-completion-codestral":
response = _complete_text_completion_codestral(_dispatch_ctx)
elif custom_llm_provider == "text-completion-inception":
response = _complete_text_completion_inception(_dispatch_ctx)
elif custom_llm_provider in ("sagemaker_chat", "sagemaker_nova"):
# boto3 reads keys from .env
# sagemaker_chat: HF Messages API endpoints
# sagemaker_nova: Nova models on SageMaker (OpenAI-compatible)
response = _complete_sagemaker_chat(_dispatch_ctx)
elif custom_llm_provider == "sagemaker":
# boto3 reads keys from .env
response = _complete_sagemaker(_dispatch_ctx)
elif custom_llm_provider == "bedrock":
# boto3 reads keys from .env
response = _complete_bedrock(_dispatch_ctx)
elif custom_llm_provider == "watsonx":
response = _complete_watsonx(_dispatch_ctx)
elif custom_llm_provider == "watsonx_text":
response = _complete_watsonx_text(_dispatch_ctx)
elif custom_llm_provider == "vllm":
response = _complete_vllm(_dispatch_ctx)
elif custom_llm_provider == "ollama":
response = _complete_ollama(_dispatch_ctx)
elif custom_llm_provider == "ollama_chat":
response = _complete_ollama_chat(_dispatch_ctx)
elif custom_llm_provider == "triton":
response = _complete_triton(_dispatch_ctx)
elif custom_llm_provider == "cloudflare":
response = _complete_cloudflare(_dispatch_ctx)
elif custom_llm_provider == "petals" or model in litellm.petals_models:
response = _complete_petals(_dispatch_ctx)
elif custom_llm_provider == "snowflake" or model in litellm.snowflake_models:
response = _complete_snowflake(_dispatch_ctx)
elif custom_llm_provider == "gradient_ai":
response = _complete_gradient_ai(_dispatch_ctx)
elif custom_llm_provider == "gdc":
response = _complete_gdc(_dispatch_ctx)
elif custom_llm_provider == "bytez":
response = _complete_bytez(_dispatch_ctx)
elif custom_llm_provider == "lemonade":
response = _complete_lemonade(_dispatch_ctx)
elif custom_llm_provider == "ovhcloud" or model in litellm.ovhcloud_models:
response = _complete_ovhcloud(_dispatch_ctx)
elif custom_llm_provider == "custom":
response = _complete_custom(_dispatch_ctx)
elif custom_llm_provider in litellm._custom_providers: # Assume custom LLM provider
# Get the Custom Handler
response = _complete_custom_providers(_dispatch_ctx)
elif custom_llm_provider == "langgraph":
# LangGraph - Agent Runtime Provider
response = _complete_langgraph(_dispatch_ctx)
elif custom_llm_provider == "langflow":
# LangFlow - Visual AI Agent Platform
response = _complete_langflow(_dispatch_ctx)
else:
raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider)
return response
return dispatch_completion(_dispatch_ctx, lambda: _complete_python(_dispatch_ctx))
except Exception as e:
## Map to OpenAI Exception
raise exception_type(
@ -7771,179 +7748,176 @@ def transcription(
custom_llm_provider=custom_llm_provider,
)
response: TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] | None = None
def python_fallback(
file: FileTypes,
api_key: str | None = api_key,
api_base: str | None = api_base,
api_version: str | None = api_version,
) -> TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]:
response: TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] | None = None
provider_config: Final = ProviderConfigManager.get_provider_audio_transcription_config(
provider_config: Final = ProviderConfigManager.get_provider_audio_transcription_config(
model=model,
provider=LlmProviders(custom_llm_provider),
)
if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None:
# azure configs
api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
azure_ad_token: Final = kwargs.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN")
api_key = api_key or litellm.api_key or litellm.azure_key or get_secret_str("AZURE_API_KEY")
optional_params["extra_headers"] = extra_headers
response = azure_audio_transcriptions.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
model_response=model_response,
atranscription=atranscription,
client=client,
timeout=timeout,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
api_version=api_version,
azure_ad_token=azure_ad_token,
max_retries=max_retries,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers):
api_base = (
api_base
or litellm.api_base
or get_secret("OPENAI_BASE_URL")
or get_secret("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
openai.organization = (
litellm.organization
or get_secret("OPENAI_ORGANIZATION")
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY")
response = openai_audio_transcriptions.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
model_response=model_response,
atranscription=atranscription,
client=client,
timeout=timeout,
logging_obj=litellm_logging_obj,
max_retries=max_retries,
api_base=api_base,
api_key=api_key,
provider_config=provider_config,
litellm_params=litellm_params_dict,
shared_session=shared_session,
)
elif custom_llm_provider == "nvidia_riva":
# NVIDIA Riva is gRPC-based, not HTTP. It has its own dedicated handler
# rather than going through base_llm_http_handler.
response = nvidia_riva_audio_transcriptions.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
litellm_params=litellm_params_dict,
model_response=model_response,
atranscription=atranscription,
timeout=timeout,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
provider_config=(
provider_config if isinstance(provider_config, NvidiaRivaAudioTranscriptionConfig) else None
),
)
elif custom_llm_provider == "soniox":
from litellm.llms.soniox.audio_transcription.handler import (
SonioxAudioTranscriptionHandler,
)
response = SonioxAudioTranscriptionHandler().audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
litellm_params=litellm_params_dict,
model_response=model_response,
atranscription=atranscription,
client=(
client
if client is not None and (isinstance(client, HTTPHandler) or isinstance(client, AsyncHTTPHandler))
else None
),
timeout=timeout,
max_retries=max_retries,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
headers=extra_headers,
provider_config=provider_config,
)
elif provider_config is not None:
response = base_llm_http_handler.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
litellm_params=litellm_params_dict,
model_response=model_response,
atranscription=atranscription,
client=(
client
if client is not None and (isinstance(client, HTTPHandler) or isinstance(client, AsyncHTTPHandler))
else None
),
timeout=timeout,
max_retries=max_retries,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
custom_llm_provider=custom_llm_provider,
headers={},
provider_config=provider_config,
shared_session=shared_session,
)
if response is None:
raise ValueError("Unmapped provider passed in. Unable to get the response.")
return response
from litellm.rust_bridge.transcription import dispatch_transcription
response: Final = dispatch_transcription(
model=model,
provider=LlmProviders(custom_llm_provider),
provider=custom_llm_provider,
file=file,
api_key=api_key,
api_base=api_base,
headers=extra_headers,
optional_params=optional_params,
timeout=timeout,
logging=litellm_logging_obj,
asynchronous=bool(atranscription),
has_custom_client=client is not None or shared_session is not None,
fallback=python_fallback,
)
if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None:
# azure configs
api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
azure_ad_token: Final = kwargs.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN")
api_key = api_key or litellm.api_key or litellm.azure_key or get_secret_str("AZURE_API_KEY")
optional_params["extra_headers"] = extra_headers
response = azure_audio_transcriptions.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
model_response=model_response,
atranscription=atranscription,
client=client,
timeout=timeout,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
api_version=api_version,
azure_ad_token=azure_ad_token,
max_retries=max_retries,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers):
api_base = (
api_base
or litellm.api_base
or get_secret("OPENAI_BASE_URL")
or get_secret("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
openai.organization = (
litellm.organization
or get_secret("OPENAI_ORGANIZATION")
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY")
response = openai_audio_transcriptions.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
model_response=model_response,
atranscription=atranscription,
client=client,
timeout=timeout,
logging_obj=litellm_logging_obj,
max_retries=max_retries,
api_base=api_base,
api_key=api_key,
provider_config=provider_config,
litellm_params=litellm_params_dict,
shared_session=shared_session,
)
elif custom_llm_provider == "nvidia_riva":
# NVIDIA Riva is gRPC-based, not HTTP. It has its own dedicated handler
# rather than going through base_llm_http_handler.
response = nvidia_riva_audio_transcriptions.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
litellm_params=litellm_params_dict,
model_response=model_response,
atranscription=atranscription,
timeout=timeout,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
provider_config=(
provider_config if isinstance(provider_config, NvidiaRivaAudioTranscriptionConfig) else None
),
)
elif custom_llm_provider == "soniox":
from litellm.llms.soniox.audio_transcription.handler import (
SonioxAudioTranscriptionHandler,
)
response = SonioxAudioTranscriptionHandler().audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
litellm_params=litellm_params_dict,
model_response=model_response,
atranscription=atranscription,
client=(
client
if client is not None and (isinstance(client, HTTPHandler) or isinstance(client, AsyncHTTPHandler))
else None
),
timeout=timeout,
max_retries=max_retries,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
headers=extra_headers,
provider_config=provider_config,
)
elif custom_llm_provider == "bedrock":
from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch
dispatch: Final = BedrockAudioTranscriptionRustDispatch()
if atranscription:
response = dispatch.async_audio_transcriptions(
model=model,
audio_file=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,
logging_obj=litellm_logging_obj,
)
else:
response = dispatch.audio_transcriptions(
model=model,
audio_file=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,
logging_obj=litellm_logging_obj,
)
elif provider_config is not None:
response = base_llm_http_handler.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
litellm_params=litellm_params_dict,
model_response=model_response,
atranscription=atranscription,
client=(
client
if client is not None and (isinstance(client, HTTPHandler) or isinstance(client, AsyncHTTPHandler))
else None
),
timeout=timeout,
max_retries=max_retries,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
custom_llm_provider=custom_llm_provider,
headers={},
provider_config=provider_config,
shared_session=shared_session,
)
# Store duration in _hidden_params for cost calculation without
# exposing it in the response body (see sync path comment above).
if response is not None and not isinstance(response, Coroutine):
if not isinstance(response, Coroutine):
existing_duration: Final = getattr(response, "duration", None)
if existing_duration is None:
calculated_duration: Final = calculate_request_duration(file)
if calculated_duration is not None:
response._hidden_params["audio_transcription_duration"] = calculated_duration
if response is None:
raise ValueError("Unmapped provider passed in. Unable to get the response.")
return response

View file

@ -8,6 +8,8 @@ import mimetypes
import os
import re
from collections.abc import Callable, Coroutine, Mapping
from contextlib import nullcontext
from dataclasses import dataclass
from io import IOBase
from typing import Any, Final, cast
@ -17,16 +19,23 @@ import litellm
from litellm._logging import verbose_logger
from litellm.constants import request_timeout
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.azure_ai.ocr.common_utils import is_azure_document_intelligence_model
from litellm.llms.azure_ai.ocr.common_utils import (
is_azure_document_intelligence_model,
)
from litellm.llms.base_llm.ocr.transformation import (
OCR_REQUEST_FORMAT_PARAM,
BaseOCRConfig,
OCRResponse,
parse_ocr_request_format,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.dispatch import anative_first, native_first, provider_errors
from litellm.rust_bridge.runtime import DispatchResult
from litellm.rust_bridge.request import (
NativeRequestOptions,
PreparedNativeCall,
vertex_options,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -35,6 +44,21 @@ base_llm_http_handler = BaseLLMHTTPHandler()
#################################################
@dataclass
class _PreparedOCRRequest:
model: str
document: dict[str, Any]
api_key: str | None
api_base: str | None
custom_llm_provider: str
extra_headers: dict[str, object] | None
provider_config: BaseOCRConfig
optional_params: dict[str, object]
litellm_params: dict[str, object]
effective_timeout: float | httpx.Timeout
litellm_logging_obj: LiteLLMLoggingObj
def _prepare_ocr_request(
model: str,
document: Mapping[str, object],
@ -44,7 +68,7 @@ def _prepare_ocr_request(
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
kwargs: dict[str, object],
) -> rust_ocr_bridge.PreparedOCRRequest:
) -> _PreparedOCRRequest:
litellm_logging_obj: Final = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj"))
litellm_call_id: Final = cast(str | None, kwargs.get("litellm_call_id", None))
@ -141,7 +165,7 @@ def _prepare_ocr_request(
custom_llm_provider=custom_llm_provider,
)
return rust_ocr_bridge.PreparedOCRRequest(
return _PreparedOCRRequest(
model=model,
document=document,
api_key=api_key,
@ -156,71 +180,157 @@ def _prepare_ocr_request(
)
@anative_first(
native=rust_ocr_bridge.aattempt_ocr,
route="ocr",
errors=lambda prepared_request, resolve_api_key: provider_errors(
prepared_request.custom_llm_provider, prepared_request.model
),
)
async def _execute_aocr(
prepared_request: rust_ocr_bridge.PreparedOCRRequest,
def _rust_bridge_optional_params(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> dict[str, object]:
optional_params: Final = dict(prepared_request.optional_params)
if prepared_request.custom_llm_provider == "vertex_ai":
vertex_project: Final = (
prepared_request.litellm_params.get("vertex_project")
or prepared_request.litellm_params.get("vertex_ai_project")
or litellm.vertex_project
or resolve_secret("VERTEXAI_PROJECT")
)
vertex_location: Final = (
prepared_request.litellm_params.get("vertex_location")
or prepared_request.litellm_params.get("vertex_ai_location")
or litellm.vertex_location
or resolve_secret("VERTEXAI_LOCATION")
or resolve_secret("VERTEX_LOCATION")
)
if vertex_project is not None:
optional_params["vertex_project"] = vertex_project
if vertex_location is not None:
optional_params["vertex_location"] = vertex_location
return optional_params
def _rust_bridge_api_base(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> str | None:
if prepared_request.api_base is not None:
return prepared_request.api_base
if prepared_request.custom_llm_provider == "azure_ai":
if is_azure_document_intelligence_model(prepared_request.model):
return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
return resolve_secret("AZURE_AI_API_BASE")
return None
def _prepare_rust_ocr_call(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> OCRResponse:
pending: Final = base_llm_http_handler.ocr(
) -> PreparedNativeCall[rust_ocr_bridge.NativeOCRRequest]:
provider_config: Final = prepared_request.provider_config
api_key_env_var: Final = provider_config.get_api_key_env_var()
resolved_api_key: Final = prepared_request.api_key or (
resolve_api_key(api_key_env_var) if api_key_env_var is not None else None
)
resolved_headers: Final = provider_config.validate_environment(
headers=prepared_request.extra_headers or {},
model=prepared_request.model,
document=prepared_request.document,
optional_params=prepared_request.optional_params,
timeout=prepared_request.effective_timeout,
logging_obj=prepared_request.litellm_logging_obj,
api_key=prepared_request.api_key,
api_key=resolved_api_key,
api_base=prepared_request.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
aocr=True,
headers=prepared_request.extra_headers,
provider_config=prepared_request.provider_config,
litellm_params=prepared_request.litellm_params,
)
response: Final = await pending if asyncio.iscoroutine(pending) else pending
if response is None:
raise ValueError(f"Got an unexpected None response from the OCR API: {response}")
return response
def _attempt_ocr(
prepared_request: rust_ocr_bridge.PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
is_async: bool,
) -> DispatchResult[OCRResponse]:
return rust_ocr_bridge.attempt_ocr(prepared_request=prepared_request, resolve_api_key=resolve_api_key)
@native_first(
native=_attempt_ocr,
route="ocr",
errors=lambda prepared_request, resolve_api_key, is_async: provider_errors(
prepared_request.custom_llm_provider, prepared_request.model
),
)
def _execute_ocr(
prepared_request: rust_ocr_bridge.PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
is_async: bool,
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
return base_llm_http_handler.ocr(
model=prepared_request.model,
document=prepared_request.document,
optional_params=prepared_request.optional_params,
timeout=prepared_request.effective_timeout,
logging_obj=prepared_request.litellm_logging_obj,
api_key=prepared_request.api_key,
resolved_complete_url: Final = provider_config.get_complete_url(
api_base=prepared_request.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
aocr=is_async,
headers=prepared_request.extra_headers,
provider_config=prepared_request.provider_config,
model=prepared_request.model,
optional_params=prepared_request.optional_params,
litellm_params=prepared_request.litellm_params,
)
rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key)
rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key)
prepared_request.litellm_logging_obj.pre_call(
input="OCR document processing",
api_key=resolved_api_key,
additional_args={
"complete_input_dict": {
"model": prepared_request.model,
"document": prepared_request.document,
**rust_optional_params,
},
"api_base": resolved_complete_url,
"headers": resolved_headers,
},
)
return PreparedNativeCall(
request=rust_ocr_bridge.NativeOCRRequest(
model=prepared_request.model,
document=prepared_request.document,
optional_params=prepared_request.optional_params,
),
options=NativeRequestOptions(
vertex=vertex_options(rust_optional_params),
api_key=resolved_api_key,
api_base=rust_api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=cast( # cast-ok: provider header normalization returns string-object pairs
dict[str, object], resolved_headers
),
timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout),
),
)
@dataclass
class _OCROperation:
request: _PreparedOCRRequest
resolve_api_key: Callable[[str], str | None]
python: Callable[[], OCRResponse | Coroutine[object, object, OCRResponse]]
logged: bool = False
def prepare(self) -> PreparedNativeCall[rust_ocr_bridge.NativeOCRRequest]:
prepared: Final = _prepare_rust_ocr_call(self.request, self.resolve_api_key)
self.logged = True
return prepared
def fallback(self) -> OCRResponse | Coroutine[object, object, OCRResponse]:
with self.request.litellm_logging_obj.suppress_next_pre_call() if self.logged else nullcontext():
return self.python()
async def afallback(self) -> OCRResponse:
with self.request.litellm_logging_obj.suppress_next_pre_call() if self.logged else nullcontext():
result: Final = self.python()
return await result if isinstance(result, Coroutine) else result
def _run_rust_ocr(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
fallback: Callable[[], OCRResponse | Coroutine[object, object, OCRResponse]],
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
operation: Final = _OCROperation(prepared_request, resolve_api_key, fallback)
return rust_ocr_bridge.dispatch_ocr(
prepare=operation.prepare,
fallback=operation.fallback,
adapt=OCRResponse.model_validate,
model=prepared_request.model,
provider=prepared_request.custom_llm_provider,
request_format=(
"native" if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native" else None
),
)
async def _run_rust_aocr(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
fallback: Callable[[], Coroutine[object, object, OCRResponse]],
) -> OCRResponse:
operation: Final = _OCROperation(prepared_request, resolve_api_key, fallback)
return await rust_ocr_bridge.adispatch_ocr(
prepare=operation.prepare,
fallback=operation.afallback,
adapt=OCRResponse.model_validate,
model=prepared_request.model,
provider=prepared_request.custom_llm_provider,
request_format=(
"native" if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native" else None
),
)
@client
@ -319,7 +429,31 @@ async def aocr(
from litellm.secret_managers.main import get_secret_str
return await _execute_aocr(prepared_request=prepared, resolve_api_key=get_secret_str)
async def python_fallback() -> OCRResponse:
pending: Final = base_llm_http_handler.ocr(
model=prepared.model,
document=prepared.document,
optional_params=prepared.optional_params,
timeout=prepared.effective_timeout,
logging_obj=prepared.litellm_logging_obj,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared.custom_llm_provider,
aocr=True,
headers=prepared.extra_headers,
provider_config=prepared.provider_config,
litellm_params=prepared.litellm_params,
)
response: Final = await pending if asyncio.iscoroutine(pending) else pending
if response is None:
raise ValueError(f"Got an unexpected None response from the OCR API: {response}")
return response
return await _run_rust_aocr(
prepared_request=prepared,
resolve_api_key=get_secret_str,
fallback=python_fallback,
)
except Exception as e:
raise litellm.exception_type(
model=model,
@ -560,7 +694,27 @@ def ocr(
from litellm.secret_managers.main import get_secret_str
return _execute_ocr(prepared_request=prepared, resolve_api_key=get_secret_str, is_async=_is_async)
def python_fallback() -> OCRResponse | Coroutine[object, object, OCRResponse]:
return base_llm_http_handler.ocr(
model=prepared.model,
document=prepared.document,
optional_params=prepared.optional_params,
timeout=prepared.effective_timeout,
logging_obj=prepared.litellm_logging_obj,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared.custom_llm_provider,
aocr=_is_async,
headers=prepared.extra_headers,
provider_config=prepared.provider_config,
litellm_params=prepared.litellm_params,
)
return _run_rust_ocr(
prepared_request=prepared,
resolve_api_key=get_secret_str,
fallback=python_fallback,
)
except Exception as e:
raise litellm.exception_type(
model=model,

View file

@ -4,22 +4,31 @@ 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 Mapping, Sequence
from dataclasses import replace
from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence
from contextlib import nullcontext
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol
import httpx
from pydantic import TypeAdapter, ValidationError
import litellm
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 get_bedrock_request_metadata_fields
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.protocols import (
RustAchatCompletions,
@ -30,22 +39,32 @@ from litellm.rust_bridge.request import (
NativeAnthropicOptions,
NativeBedrockOptions,
NativeChatCompletionsRequest,
NativeRequestCapabilities,
NativePreCallDetails,
NativeRequestContext,
NativeRequestOptions,
PreparedNativeCall,
anthropic_options,
bedrock_options,
call_native,
with_capabilities,
)
from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointBinding,
EndpointDispatch,
PythonFallback,
async_none,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.secret_managers.main import get_secret_str
from litellm.types.completion import (
_CompletionDispatchContext, # pyright: ignore[reportPrivateUsage] # shared internal SDK dispatch context
_CompletionDispatchResult, # pyright: ignore[reportPrivateUsage] # shared internal SDK dispatch result
)
from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
RUST_RESPONSE_HEADER: Final = "x-litellm-rust"
@ -89,10 +108,16 @@ def response_logger(
return log
_CHAT: Final[NativeBinding[RustChatCompletions]] = NativeBinding(lambda native: native.chat_completions)
_ACHAT: Final[NativeBinding[RustAchatCompletions]] = NativeBinding(lambda native: native.achat_completions)
_CHAT_PREFLIGHT: Final[NativeBinding[RustChatCompletionsDecline]] = NativeBinding(
lambda native: native.chat_completions_decline
_CHAT: Final[EndpointDispatch[RustChatCompletions, RustAchatCompletions]] = EndpointDispatch.native(
route="chat_completions",
sync=lambda native: native.chat_completions,
asynchronous=lambda native: native.achat_completions,
enabled=rust_enabled,
)
_CHAT_PREFLIGHT: Final[EndpointBinding[RustChatCompletionsDecline]] = EndpointBinding.native(
route="chat_completions",
select=lambda native: native.chat_completions_decline,
enabled=rust_enabled,
)
@ -106,14 +131,14 @@ def set_rust_chat_completions(
patching module attributes."""
if not isinstance(chat_completions, Unchanged):
if chat_completions is None:
_CHAT.reset()
_CHAT.sync.reset()
else:
_CHAT.override(chat_completions)
_CHAT.sync.override(chat_completions)
if not isinstance(achat_completions, Unchanged):
if achat_completions is None:
_ACHAT.reset()
_CHAT.asynchronous.reset()
else:
_ACHAT.override(achat_completions)
_CHAT.asynchronous.override(achat_completions)
if not isinstance(decline, Unchanged):
if decline is None:
_CHAT_PREFLIGHT.reset()
@ -121,21 +146,16 @@ def set_rust_chat_completions(
_CHAT_PREFLIGHT.override(decline)
def _provider_eligibility_options(
provider: str | None,
litellm_params: Mapping[str, object] | None,
optional_params: Mapping[str, object],
) -> NativeRequestOptions:
bedrock: Final = (
replace(
bedrock_options(optional_params),
request_metadata_fields=get_bedrock_request_metadata_fields(),
)
if provider == "bedrock"
else None
def _preflight_context(litellm_params: Mapping[str, object] | None) -> NativeRequestContext:
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 NativeRequestContext(request_metadata_fields=get_bedrock_request_metadata_fields())
return NativeRequestContext(
metadata=entries,
request_metadata_fields=get_bedrock_request_metadata_fields(),
)
anthropic: Final = anthropic_options(litellm_params) if provider == "anthropic" else None
return NativeRequestOptions(custom_llm_provider=provider, bedrock=bedrock, anthropic=anthropic)
def rust_chat_completions_accepts(
@ -154,28 +174,16 @@ def rust_chat_completions_accepts(
capability gate answers the second half; it resolves no credentials and
performs no I/O.
"""
if not rust_enabled():
return False
decline: Final = _CHAT_PREFLIGHT.load()
if decline is None:
return False
try:
reason: Final = decline(
return _CHAT_PREFLIGHT.accepts(
check=lambda decline: decline(
model=model,
messages=messages,
optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
options=_provider_eligibility_options(custom_llm_provider, litellm_params, optional_params),
context=NativeRequestContext(
capabilities=NativeRequestCapabilities(stream=bool(stream))
),
)
except Exception as error: # noqa: BLE001 # capability checks perform no provider I/O
verbose_logger.debug("Native chat acceptance check failed: %s", error)
return False
if reason is not None:
verbose_logger.debug("Native chat request is ineligible: %s", reason)
return reason is None
context=_preflight_context(litellm_params),
stream=bool(stream),
),
)
def _build_model_response(
@ -206,26 +214,18 @@ def chat_completions(
on_response: ResponseObserver,
bedrock: NativeBedrockOptions | None = None,
anthropic: NativeAnthropicOptions | None = None,
stream: bool = False,
has_custom_client: bool = False,
eligible: bool = True,
context: NativeRequestContext | None = None,
) -> DispatchResult[ModelResponse]:
) -> ModelResponse | None:
def adapt(rust_response: Mapping[str, object]) -> ModelResponse:
on_response(rust_response)
return _build_model_response(rust_response, model_response)
def call(
native: RustChatCompletions, prepared: PreparedNativeCall[NativeChatCompletionsRequest]
) -> Mapping[str, object]:
return call_native(native, prepared)
return attempt(
load=_CHAT.load,
enabled=rust_enabled(),
eligible=eligible,
return _CHAT.invoke(
prepare=lambda: PreparedNativeCall(
request=NativeChatCompletionsRequest(model=model, messages=messages, optional_params=optional_params),
NativeChatCompletionsRequest(
model=model,
messages=messages,
optional_params=optional_params,
),
options=NativeRequestOptions(
api_key=api_key,
api_base=api_base,
@ -235,17 +235,12 @@ def chat_completions(
bedrock=bedrock,
anthropic=anthropic,
),
context=with_capabilities(
context or NativeRequestContext(),
NativeRequestCapabilities(
execution_mode="sync",
stream=stream,
has_custom_client=has_custom_client,
),
),
context=NativeRequestContext(),
),
call=call,
call=call_native,
fallback=lambda: None,
adapt=adapt,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -263,27 +258,18 @@ async def achat_completions(
on_response: ResponseObserver,
bedrock: NativeBedrockOptions | None = None,
anthropic: NativeAnthropicOptions | None = None,
stream: bool = False,
has_custom_client: bool = False,
eligible: bool = True,
context: NativeRequestContext | None = None,
) -> DispatchResult[ModelResponse]:
) -> ModelResponse | None:
def adapt(rust_response: Mapping[str, object]) -> ModelResponse:
on_response(rust_response)
return _build_model_response(rust_response, model_response)
async def call(
native: RustAchatCompletions,
prepared: PreparedNativeCall[NativeChatCompletionsRequest],
) -> Mapping[str, object]:
return await call_native(native, prepared)
return await aattempt(
load=_ACHAT.load,
enabled=rust_enabled(),
eligible=eligible,
return await _CHAT.ainvoke(
prepare=lambda: PreparedNativeCall(
request=NativeChatCompletionsRequest(model=model, messages=messages, optional_params=optional_params),
NativeChatCompletionsRequest(
model=model,
messages=messages,
optional_params=optional_params,
),
options=NativeRequestOptions(
api_key=api_key,
api_base=api_base,
@ -293,15 +279,197 @@ async def achat_completions(
bedrock=bedrock,
anthropic=anthropic,
),
context=with_capabilities(
context or NativeRequestContext(),
NativeRequestCapabilities(
execution_mode="async",
stream=stream,
has_custom_client=has_custom_client,
context=NativeRequestContext(),
),
call=call_native,
fallback=async_none,
adapt=adapt,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
async def achat_completions_or_fallback(
*,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
model_response: ModelResponse,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
python_fallback: Callable[[], Awaitable[object]],
bedrock: NativeBedrockOptions | None = None,
anthropic: NativeAnthropicOptions | None = None,
) -> object:
"""Await the Rust path, falling back to the caller's own Python path when
the bridge is unavailable or the call fails.
The caller supplies the fallback, so the bridge stays free of provider
dispatch. This exists because a caller that dispatches asynchronously has
already returned a coroutine by the time a Rust failure surfaces, and so
cannot fall back on its own.
"""
def adapt(rust_response: Mapping[str, object]) -> object:
on_response(rust_response)
return _build_model_response(rust_response, model_response)
return await _CHAT.ainvoke(
prepare=lambda: PreparedNativeCall(
NativeChatCompletionsRequest(
model=model,
messages=messages,
optional_params=optional_params,
),
options=NativeRequestOptions(
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
bedrock=bedrock,
anthropic=anthropic,
),
context=NativeRequestContext(),
),
call=call_native,
fallback=python_fallback,
adapt=adapt,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object])
_STR_ADAPTER: Final = TypeAdapter(str | None)
@dataclass
class _ChatOperation:
context: _CompletionDispatchContext
python: Callable[[], _CompletionDispatchResult]
pre_call_logged: bool = False
def assess(self) -> PythonFallback | None:
ctx: Final = self.context
return _CHAT_PREFLIGHT.assess(
check=lambda decline: decline(
model=ctx.model,
messages=ctx.messages,
optional_params=ctx.optional_params,
custom_llm_provider=ctx.custom_llm_provider,
context=_preflight_context(ctx.litellm_params),
stream=bool(ctx.stream),
has_custom_client=ctx.client is not None or ctx.shared_session is not None,
has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(ctx.logging),
),
)
def prepare(self) -> PreparedNativeCall[NativeChatCompletionsRequest]:
ctx: Final = self.context
config: Final = ctx.provider_config
defaults: Final = (
_PARAMS_ADAPTER.validate_python(config.get_config_for_model(ctx.model))
if config is not None
else MappingProxyType({})
)
params: Final = _PARAMS_ADAPTER.validate_python(MappingProxyType({**defaults, **ctx.optional_params}))
key: Final = (
ctx.api_key
or _STR_ADAPTER.validate_python(getattr(litellm, f"{ctx.custom_llm_provider}_key", None))
or litellm.api_key
or get_secret_str(f"{ctx.custom_llm_provider.upper()}_API_KEY")
)
base: Final = (
ctx.api_base
or litellm.api_base
or get_secret_str(f"{ctx.custom_llm_provider.upper()}_API_BASE")
or get_secret_str(f"{ctx.custom_llm_provider.upper()}_BASE_URL")
)
initial_headers: Final = _PARAMS_ADAPTER.validate_python(
MappingProxyType({**(ctx.headers or MappingProxyType({})), **(ctx.extra_headers or MappingProxyType({}))})
)
headers: Final = (
_PARAMS_ADAPTER.validate_python(
config.validate_environment(
api_key=key,
api_base=base,
headers=initial_headers,
model=ctx.model,
messages=ctx.messages,
optional_params=params,
litellm_params=ctx.litellm_params,
)
)
if config is not None
else initial_headers
)
log_details: Final[NativePreCallDetails] = {
"complete_input_dict": {"model": ctx.model, "messages": ctx.messages, **params},
"api_base": base or "",
"headers": headers,
}
ctx.logging.pre_call(input=ctx.messages, api_key=key, additional_args=log_details)
self.pre_call_logged = True
return PreparedNativeCall(
NativeChatCompletionsRequest(
model=ctx.model,
messages=ctx.messages,
optional_params=provider_request_params(params),
options=NativeRequestOptions(
api_key=key,
api_base=base,
custom_llm_provider=ctx.custom_llm_provider,
extra_headers=headers,
timeout_seconds=timeout_to_seconds(
float(ctx.timeout) if isinstance(ctx.timeout, str) else ctx.timeout
),
provider_connection=provider_connection_params(params),
),
),
),
call=call,
adapt=adapt,
context=_preflight_context(ctx.litellm_params),
)
def fallback(self) -> _CompletionDispatchResult:
with self.context.logging.suppress_next_pre_call() if self.pre_call_logged else nullcontext():
return self.python()
async def afallback(self) -> ModelResponse | litellm.CustomStreamWrapper:
with self.context.logging.suppress_next_pre_call() if self.pre_call_logged else nullcontext():
result: Final = self.python()
return await result if isinstance(result, Coroutine) else result
def adapt(self, response: Mapping[str, object]) -> ModelResponse:
self.context.logging.post_call(
input=self.context.messages,
api_key=self.context.api_key,
original_response=json.dumps(response),
)
return _build_model_response(response, self.context.model_response)
def dispatch_completion(
context: _CompletionDispatchContext,
fallback: Callable[[], _CompletionDispatchResult],
) -> _CompletionDispatchResult:
operation: Final = _ChatOperation(context, fallback)
error_context: Final = BridgeErrorContext(provider=context.custom_llm_provider, model=context.model)
if context.acompletion:
return _CHAT.ainvoke(
prepare=operation.prepare,
call=call_native,
fallback=operation.afallback,
adapt=operation.adapt,
error_context=error_context,
preflight=operation.assess,
)
return _CHAT.invoke(
prepare=operation.prepare,
call=call_native,
fallback=operation.fallback,
adapt=operation.adapt,
error_context=error_context,
preflight=operation.assess,
)

View file

@ -2,51 +2,94 @@
from __future__ import annotations
from typing import Final
import json
from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Sequence
from contextlib import nullcontext
from dataclasses import dataclass
from functools import reduce
from types import MappingProxyType
from typing import Final, Literal
import httpx
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged
from litellm.rust_bridge.protocols import RustAmessages, RustMessages
from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.protocols import RustAmessages, RustMessages, RustRouteDecline
from litellm.rust_bridge.request import (
NativeMessagesRequest,
NativeRequestCapabilities,
NativePreCallDetails,
NativeRequestContext,
NativeRequestOptions,
PreparedNativeCall,
call_native,
with_capabilities,
)
from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt, identity
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointBinding,
EndpointDispatch,
PythonFallback,
assess_route,
async_none,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
from litellm.types.llms.openai import ChatCompletionUserMessage
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders, ProviderSpecificHeader
from litellm.utils import ProviderConfigManager
_MESSAGES: Final[NativeBinding[RustMessages]] = NativeBinding(lambda native: native.messages)
_AMESSAGES: Final[NativeBinding[RustAmessages]] = NativeBinding(lambda native: native.amessages)
_MESSAGES: Final[EndpointDispatch[RustMessages, RustAmessages]] = EndpointDispatch.native(
route="messages",
sync=lambda native: native.messages,
asynchronous=lambda native: native.amessages,
enabled=rust_enabled,
)
_PREFLIGHT: Final[EndpointBinding[RustRouteDecline]] = EndpointBinding.native(
route="messages",
select=lambda native: native.messages_decline,
enabled=rust_enabled,
)
def set_rust_messages(
*,
messages: RustMessages | None | Unchanged = UNCHANGED,
amessages: RustAmessages | None | Unchanged = UNCHANGED,
decline: RustRouteDecline | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(decline, Unchanged):
if decline is None:
_PREFLIGHT.reset()
else:
_PREFLIGHT.override(decline)
if not isinstance(messages, Unchanged):
if messages is None:
_MESSAGES.reset()
_MESSAGES.sync.reset()
else:
_MESSAGES.override(messages)
_MESSAGES.sync.override(messages)
if not isinstance(amessages, Unchanged):
if amessages is None:
_AMESSAGES.reset()
_MESSAGES.asynchronous.reset()
else:
_AMESSAGES.override(amessages)
_MESSAGES.asynchronous.override(amessages)
def load_rust_messages() -> RustMessages | None:
return _MESSAGES.load()
return _MESSAGES.sync.load()
def load_rust_amessages() -> RustAmessages | None:
return _AMESSAGES.load()
return _MESSAGES.asynchronous.load()
def messages(
@ -58,17 +101,13 @@ def messages(
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
stream: bool = False,
has_custom_client: bool = False,
has_agentic_hook: bool = False,
context: NativeRequestContext | None = None,
) -> DispatchResult[dict[str, object]]:
return attempt(
load=_MESSAGES.load,
enabled=True,
eligible=True,
) -> dict[str, object] | None:
return _MESSAGES.invoke(
prepare=lambda: PreparedNativeCall(
request=NativeMessagesRequest(model=model, body=body),
NativeMessagesRequest(
model=model,
body=body,
),
options=NativeRequestOptions(
api_key=api_key,
api_base=api_base,
@ -76,18 +115,13 @@ def messages(
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
),
context=with_capabilities(
context or NativeRequestContext(),
NativeRequestCapabilities(
execution_mode="sync",
stream=stream,
has_custom_client=has_custom_client,
has_agentic_hook=has_agentic_hook,
),
),
context=NativeRequestContext(),
),
call=call_native,
preflight=lambda: assess_route(_PREFLIGHT, model, custom_llm_provider or ""),
fallback=lambda: None,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
@ -100,17 +134,13 @@ async def amessages(
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
stream: bool = False,
has_custom_client: bool = False,
has_agentic_hook: bool = False,
context: NativeRequestContext | None = None,
) -> DispatchResult[dict[str, object]]:
return await aattempt(
load=_AMESSAGES.load,
enabled=True,
eligible=True,
) -> dict[str, object] | None:
return await _MESSAGES.ainvoke(
prepare=lambda: PreparedNativeCall(
request=NativeMessagesRequest(model=model, body=body),
NativeMessagesRequest(
model=model,
body=body,
),
options=NativeRequestOptions(
api_key=api_key,
api_base=api_base,
@ -118,16 +148,197 @@ async def amessages(
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
),
context=with_capabilities(
context or NativeRequestContext(),
NativeRequestCapabilities(
execution_mode="async",
stream=stream,
has_custom_client=has_custom_client,
has_agentic_hook=has_agentic_hook,
),
),
context=NativeRequestContext(),
),
call=call_native,
preflight=lambda: assess_route(_PREFLIGHT, model, custom_llm_provider or ""),
fallback=async_none,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
MessagesResponse = AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]
MessagesResult = MessagesResponse | Coroutine[object, None, MessagesResponse]
_BODY_ADAPTER: Final = TypeAdapter(dict[str, object])
class _NativeMessagesResponse(BaseModel):
id: str
type: Literal["message"]
role: Literal["assistant"]
model: str
content: list[dict[str, object]]
usage: dict[str, object]
class _BridgedMessagesResponse(AnthropicMessagesResponse):
_hidden_params: ReadOnly[dict[str, dict[str, str]]]
_RESPONSE_ADAPTER: Final = TypeAdapter(AnthropicMessagesResponse)
@dataclass
class _MessagesOperation:
model: str
provider: str
messages: list[dict[str, object]]
body: Callable[[], dict[str, object]]
params: GenericLiteLLMParams
logging: Logging | None
api_key: str | None
api_base: str | None
python: Callable[[], MessagesResult]
logged: bool = False
def prepare(self) -> PreparedNativeCall[NativeMessagesRequest]:
requested: Final = self.body()
body: Final = _BODY_ADAPTER.validate_python(
reduce(
delete_nested_value,
TypeAdapter(tuple[str, ...]).validate_python(self.params.get("additional_drop_params") or ()),
requested,
)
)
provider_headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
TypeAdapter(ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None).validate_python(
self.params.get("provider_specific_header")
),
self.provider,
)
config: Final = ProviderConfigManager.get_provider_anthropic_messages_config(
model=self.model, provider=LlmProviders(self.provider)
)
initial_headers: Final = _BODY_ADAPTER.validate_python(
MappingProxyType(
{
**(self.params.get("headers") or MappingProxyType({})),
**(self.params.get("extra_headers") or MappingProxyType({})),
**provider_headers,
}
)
)
validated_headers, base = (
config.validate_anthropic_messages_environment(
headers=initial_headers,
model=self.model,
messages=self.messages,
optional_params=body,
litellm_params=self.params.model_dump(),
api_key=self.api_key,
api_base=self.api_base,
)
if config is not None
else (initial_headers, self.api_base)
)
headers: Final = (
update_headers_with_filtered_beta(headers=validated_headers, provider=self.provider)
if config is not None and config.should_filter_anthropic_beta_headers()
else validated_headers
)
request_body: Final = _BODY_ADAPTER.validate_python(
MappingProxyType({**body, "model": self.model, "messages": self.messages})
)
if self.logging is not None:
self.logging.update_from_kwargs(
kwargs=self.params.model_dump(),
model=self.model,
optional_params=body,
litellm_params=self.params.model_dump(),
custom_llm_provider=self.provider,
)
self.logging.model_call_details.update(request_body)
log_details: Final[NativePreCallDetails] = {
"complete_input_dict": request_body,
"api_base": base or "",
"headers": headers,
}
log_input: Final[ChatCompletionUserMessage] = {"role": "user", "content": json.dumps(request_body)}
self.logging.pre_call(
input=[log_input], # mutable-ok: logging callbacks expect a concrete message list
api_key=self.api_key,
additional_args=log_details,
)
self.logged = True
return PreparedNativeCall(
NativeMessagesRequest(
model=self.model,
body=request_body,
options=NativeRequestOptions(
api_key=self.api_key,
api_base=base,
custom_llm_provider=self.provider,
extra_headers=headers,
timeout_seconds=timeout_to_seconds(
BaseLLMHTTPHandler.resolve_anthropic_messages_timeout(self.params, False, self.provider)
),
),
)
)
def fallback(self) -> MessagesResult:
with self.logging.suppress_next_pre_call() if self.logged and self.logging is not None else nullcontext():
return self.python()
async def afallback(self) -> MessagesResponse:
with self.logging.suppress_next_pre_call() if self.logged and self.logging is not None else nullcontext():
response: Final = self.python()
return await response if isinstance(response, Coroutine) else response
def adapt(self, response: dict[str, object]) -> AnthropicMessagesResponse:
_NativeMessagesResponse.model_validate(response)
parsed: Final[_BridgedMessagesResponse] = {
**_RESPONSE_ADAPTER.validate_python(response),
"_hidden_params": {"additional_headers": {"x-litellm-rust": "true"}},
}
if self.logging is not None:
self.logging.post_call(input=self.messages, api_key=self.api_key, original_response=json.dumps(response))
return parsed
def dispatch_messages(
*,
model: str,
provider: str,
messages: list[dict[str, object]],
body: Callable[[], dict[str, object]],
params: GenericLiteLLMParams,
logging: Logging | None,
api_key: str | None,
api_base: str | None,
stream: bool,
asynchronous: bool,
has_custom_client: bool,
fallback: Callable[[], MessagesResult],
) -> MessagesResult:
operation: Final = _MessagesOperation(model, provider, messages, body, params, logging, api_key, api_base, fallback)
def preflight() -> PythonFallback | None:
return assess_route(
_PREFLIGHT,
model,
provider,
stream=stream,
has_custom_client=has_custom_client,
has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(logging),
)
error_context: Final = BridgeErrorContext(provider=provider, model=model)
if asynchronous:
return _MESSAGES.ainvoke(
prepare=operation.prepare,
call=call_native,
adapt=operation.adapt,
fallback=operation.afallback,
preflight=preflight,
error_context=error_context,
)
return _MESSAGES.invoke(
prepare=operation.prepare,
call=call_native,
adapt=operation.adapt,
fallback=operation.fallback,
preflight=preflight,
error_context=error_context,
)

View file

@ -1,277 +1,109 @@
"""Thin Python wrapper for the native Rust OCR bridge."""
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from typing import Final
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, TypeVar
import httpx
from pydantic import TypeAdapter
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.azure_ai.ocr.common_utils import is_azure_document_intelligence_model
from litellm.llms.base_llm.ocr.transformation import OCR_REQUEST_FORMAT_PARAM, BaseOCRConfig, OCRResponse
from litellm.rust_bridge import configuration as _configuration
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.protocols import RustAocr, RustOcr
from litellm.rust_bridge.request import (
NativeOCRRequest,
NativeRequestCapabilities,
NativeRequestOptions,
PreparedNativeCall,
call_native,
request_context,
vertex_options,
from . import configuration as _configuration
from .bindings import UNCHANGED, Unchanged
from .protocols import RustAocr, RustOcr, RustRouteDecline
from .request import NativeOCRRequest, PreparedNativeCall, call_native
from .runtime import (
BridgeErrorContext,
EndpointBinding,
EndpointDispatch,
assess_route,
)
from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt
from litellm.rust_bridge.timeouts import timeout_to_seconds
_OCR: Final[NativeBinding[RustOcr]] = NativeBinding(lambda native: native.ocr)
_AOCR: Final[NativeBinding[RustAocr]] = NativeBinding(lambda native: native.aocr)
_HEADERS: Final = TypeAdapter(dict[str, object])
rust_ocr_enabled = _configuration.rust_ocr_enabled
rust = _configuration.rust
ResultT = TypeVar("ResultT")
@dataclass(frozen=True, slots=True)
class PreparedOCRRequest:
model: str
document: dict[str, object]
api_key: str | None
api_base: str | None
custom_llm_provider: str
extra_headers: dict[str, object] | None
provider_config: BaseOCRConfig
optional_params: dict[str, object]
litellm_params: dict[str, object]
effective_timeout: float | httpx.Timeout
litellm_logging_obj: LiteLLMLoggingObj
@dataclass(frozen=True, slots=True)
class _PreparedRustOCRCall:
api_key: str | None
api_base: str | None
headers: dict[str, object]
optional_params: dict[str, object]
_RUST_OCR_PROVIDERS: Final = frozenset(
{
"mistral",
"azure_ai",
"vertex_ai",
}
_OCR: Final[EndpointDispatch[RustOcr, RustAocr]] = EndpointDispatch.native(
route="ocr",
sync=lambda native: native.ocr,
asynchronous=lambda native: native.aocr,
enabled=_configuration.rust_ocr_enabled,
)
_PREFLIGHT: Final[EndpointBinding[RustRouteDecline]] = EndpointBinding.native(
route="ocr",
select=lambda native: native.ocr_decline,
enabled=_configuration.rust_ocr_enabled,
)
def set_rust_ocr(
*,
ocr: RustOcr | None | Unchanged = UNCHANGED,
aocr: RustAocr | None | Unchanged = UNCHANGED,
decline: RustRouteDecline | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(decline, Unchanged):
if decline is None:
_PREFLIGHT.reset()
else:
_PREFLIGHT.override(decline)
if not isinstance(ocr, Unchanged):
if ocr is None:
_OCR.sync.reset()
else:
_OCR.sync.override(ocr)
if not isinstance(aocr, Unchanged):
if aocr is None:
_OCR.asynchronous.reset()
else:
_OCR.asynchronous.override(aocr)
def load_rust_ocr() -> RustOcr | None:
return _OCR.load()
return _OCR.sync.load()
def load_rust_aocr() -> RustAocr | None:
return _AOCR.load()
return _OCR.asynchronous.load()
def _rust_ocr_supported(prepared_request: PreparedOCRRequest) -> bool:
if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native":
return False
if not prepared_request.provider_config.supports_rust_bridge():
return False
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
def _ocr_input_source_kind(document: dict[str, object]) -> str:
if "document_url" in document:
return "document_url"
if "image_url" in document:
return "image_url"
if "file" in document:
return "file"
return "inline"
def _ocr_request_format(optional_params: dict[str, object]) -> str | None:
value = optional_params.get(OCR_REQUEST_FORMAT_PARAM)
return value if isinstance(value, str) else None
def _rust_bridge_optional_params(
prepared_request: PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> dict[str, object]:
if prepared_request.custom_llm_provider != "vertex_ai":
return prepared_request.optional_params
vertex_project: Final = (
prepared_request.litellm_params.get("vertex_project")
or prepared_request.litellm_params.get("vertex_ai_project")
or litellm.vertex_project
or resolve_secret("VERTEXAI_PROJECT")
)
vertex_location: Final = (
prepared_request.litellm_params.get("vertex_location")
or prepared_request.litellm_params.get("vertex_ai_location")
or litellm.vertex_location
or resolve_secret("VERTEXAI_LOCATION")
or resolve_secret("VERTEX_LOCATION")
)
return {
**prepared_request.optional_params,
**{
name: value
for name, value in (("vertex_project", vertex_project), ("vertex_location", vertex_location))
if value is not None
},
}
def _rust_bridge_api_base(
prepared_request: PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> str | None:
if prepared_request.api_base is not None:
return prepared_request.api_base
if prepared_request.custom_llm_provider == "azure_ai":
if is_azure_document_intelligence_model(prepared_request.model):
return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
return resolve_secret("AZURE_AI_API_BASE")
return None
def _prepare_rust_ocr_call(
prepared_request: PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> _PreparedRustOCRCall:
provider_config: Final = prepared_request.provider_config
api_key_env_var: Final = provider_config.get_api_key_env_var()
resolved_api_key: Final = prepared_request.api_key or (
resolve_api_key(api_key_env_var) if api_key_env_var is not None else None
)
resolved_headers: Final = _HEADERS.validate_python(
provider_config.validate_environment(
headers=prepared_request.extra_headers or {},
model=prepared_request.model,
api_key=resolved_api_key,
api_base=prepared_request.api_base,
litellm_params=prepared_request.litellm_params,
)
)
resolved_complete_url: Final = provider_config.get_complete_url(
api_base=prepared_request.api_base,
model=prepared_request.model,
optional_params=prepared_request.optional_params,
litellm_params=prepared_request.litellm_params,
)
rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key)
rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key)
prepared_request.litellm_logging_obj.pre_call(
input="OCR document processing",
api_key=resolved_api_key,
additional_args={
"complete_input_dict": {
"model": prepared_request.model,
"document": prepared_request.document,
**rust_optional_params,
},
"api_base": resolved_complete_url,
"headers": resolved_headers,
},
)
return _PreparedRustOCRCall(
api_key=resolved_api_key,
api_base=rust_api_base,
headers=resolved_headers,
optional_params=rust_optional_params,
def dispatch_ocr(
*,
prepare: Callable[[], PreparedNativeCall[NativeOCRRequest]],
fallback: Callable[[], ResultT],
adapt: Callable[[Mapping[str, object]], ResultT],
model: str,
provider: str,
eligible: bool = True,
request_format: str | None = None,
) -> ResultT:
return _OCR.invoke(
prepare=prepare,
call=call_native,
fallback=fallback,
adapt=adapt,
error_context=BridgeErrorContext(provider=provider, model=model),
eligible=eligible,
preflight=lambda: assess_route(_PREFLIGHT, model, provider, request_format=request_format),
)
def attempt_ocr(
prepared_request: PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> DispatchResult[OCRResponse]:
return attempt(
load=_OCR.load,
enabled=_configuration.rust_enabled(),
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
),
call=lambda native, prepared: call_native(
native,
PreparedNativeCall(
request=NativeOCRRequest(
model=prepared_request.model,
document=prepared_request.document,
optional_params=prepared.optional_params,
),
options=NativeRequestOptions(
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout),
vertex=vertex_options(prepared.optional_params),
),
context=request_context(
logging_obj=prepared_request.litellm_logging_obj,
request_model=prepared_request.model,
litellm_params=prepared_request.litellm_params,
capabilities=NativeRequestCapabilities(
execution_mode="sync",
input_source_kind=_ocr_input_source_kind(prepared_request.document),
request_format=_ocr_request_format(prepared_request.optional_params),
native_response_format=(
prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native"
),
),
),
),
),
adapt=OCRResponse.model_validate,
eligible=_rust_ocr_supported(prepared_request),
)
async def aattempt_ocr(
prepared_request: PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> DispatchResult[OCRResponse]:
return await aattempt(
load=_AOCR.load,
enabled=_configuration.rust_enabled(),
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
),
call=lambda native, prepared: call_native(
native,
PreparedNativeCall(
request=NativeOCRRequest(
model=prepared_request.model,
document=prepared_request.document,
optional_params=prepared.optional_params,
),
options=NativeRequestOptions(
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
timeout_seconds=timeout_to_seconds(prepared_request.effective_timeout),
vertex=vertex_options(prepared.optional_params),
),
context=request_context(
logging_obj=prepared_request.litellm_logging_obj,
request_model=prepared_request.model,
litellm_params=prepared_request.litellm_params,
capabilities=NativeRequestCapabilities(
execution_mode="async",
input_source_kind=_ocr_input_source_kind(prepared_request.document),
request_format=_ocr_request_format(prepared_request.optional_params),
native_response_format=(
prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native"
),
),
),
),
),
adapt=OCRResponse.model_validate,
eligible=_rust_ocr_supported(prepared_request),
async def adispatch_ocr(
*,
prepare: Callable[[], PreparedNativeCall[NativeOCRRequest]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[Mapping[str, object]], ResultT],
model: str,
provider: str,
eligible: bool = True,
request_format: str | None = None,
) -> ResultT:
return await _OCR.ainvoke(
prepare=prepare,
call=call_native,
fallback=fallback,
adapt=adapt,
error_context=BridgeErrorContext(provider=provider, model=model),
eligible=eligible,
preflight=lambda: assess_route(_PREFLIGHT, model, provider, request_format=request_format),
)

View file

@ -32,8 +32,10 @@ class RustChatCompletionsDecline(Protocol):
optional_params: Mapping[str, object] | None,
custom_llm_provider: str | None,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
stream: bool,
has_custom_client: bool = False,
has_agentic_hook: bool = False,
) -> str | None: ...
@ -56,6 +58,19 @@ class RustResponsesWebSocketConnection(Protocol):
) -> RustResponsesWebSocket: ...
class RustRouteDecline(Protocol):
def __call__(
self,
model: str,
custom_llm_provider: str,
*,
stream: bool = False,
has_agentic_hook: bool = False,
has_custom_client: bool = False,
request_format: str | None = None,
) -> str | None: ...
class NativeModule(Protocol):
@property
def chat_completions(self) -> RustChatCompletions: ...
@ -92,3 +107,15 @@ class NativeModule(Protocol):
@property
def atranscription(self) -> RustAtranscription: ...
@property
def ocr_decline(self) -> RustRouteDecline: ...
@property
def messages_decline(self) -> RustRouteDecline: ...
@property
def transcription_decline(self) -> RustRouteDecline: ...
@property
def responses_websocket_decline(self) -> RustRouteDecline: ...

View file

@ -68,6 +68,14 @@ def vertex_options(params: Mapping[str, object]) -> NativeVertexOptions:
location=location if isinstance(location, str) else None,
)
from typing_extensions import ReadOnly, TypedDict
class NativePreCallDetails(TypedDict):
complete_input_dict: ReadOnly[Mapping[str, object]]
api_base: ReadOnly[str]
headers: ReadOnly[Mapping[str, object] | None]
@dataclass(frozen=True, slots=True)
class NativeRequestOptions:

View file

@ -2,40 +2,59 @@
from __future__ import annotations
from collections.abc import AsyncGenerator
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from typing import Final
from collections.abc import AsyncGenerator, Awaitable, Callable
from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
from typing import Final, Protocol
import httpx
from websockets.exceptions import ConnectionClosedOK
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.protocols import (
RustResponsesWebSocket,
RustResponsesWebSocketConnection,
RustRouteDecline,
)
from litellm.rust_bridge.request import (
NativeRequestCapabilities,
NativeRequestContext,
NativeRequestOptions,
NativeResponsesWebSocketRequest,
PreparedNativeCall,
call_native,
with_capabilities,
)
from litellm.rust_bridge.runtime import DispatchResult, aattempt, adapt_result
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointBinding,
assess_route,
async_none,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
_RESPONSES_WEBSOCKET: Final[NativeBinding[RustResponsesWebSocketConnection]] = NativeBinding(
lambda native: native.ResponsesWebSocketConnection,
_RESPONSES_WEBSOCKET: Final[EndpointBinding[RustResponsesWebSocketConnection]] = EndpointBinding.native(
route="responses_websocket",
select=lambda native: native.ResponsesWebSocketConnection,
enabled=rust_enabled,
)
_PREFLIGHT: Final[EndpointBinding[RustRouteDecline]] = EndpointBinding.native(
route="responses_websocket",
select=lambda native: native.responses_websocket_decline,
enabled=rust_enabled,
)
def set_rust_responses_websocket(
*,
connection: RustResponsesWebSocketConnection | None | Unchanged = UNCHANGED,
decline: RustRouteDecline | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(decline, Unchanged):
if decline is None:
_PREFLIGHT.reset()
else:
_PREFLIGHT.override(decline)
if not isinstance(connection, Unchanged):
if connection is None:
_RESPONSES_WEBSOCKET.reset()
@ -43,7 +62,13 @@ def set_rust_responses_websocket(
_RESPONSES_WEBSOCKET.override(connection)
class ConnectionAdapter:
class Connection(Protocol):
async def send(self, text: str) -> None: ...
async def recv(self) -> str | bytes: ...
async def close(self) -> None: ...
class _ConnectionAdapter:
def __init__(self, connection: RustResponsesWebSocket):
self._connection: Final[RustResponsesWebSocket] = connection
@ -65,54 +90,55 @@ async def connect(
url: str,
headers: dict[str, str],
timeout: float | httpx.Timeout | None,
websocket_mode: str = "native",
requires_connection: bool = True,
context: NativeRequestContext | None = None,
) -> DispatchResult[ConnectionAdapter]:
return await aattempt(
load=_RESPONSES_WEBSOCKET.load,
enabled=rust_enabled(),
eligible=True,
model: str = "responses websocket",
provider: str = "openai",
fallback: Callable[[], Awaitable[Connection | None]] = async_none,
) -> Connection | None:
return await _RESPONSES_WEBSOCKET.ainvoke(
prepare=lambda: PreparedNativeCall(
request=NativeResponsesWebSocketRequest(url=url),
options=NativeRequestOptions(extra_headers=headers, timeout_seconds=timeout_to_seconds(timeout)),
context=with_capabilities(
context or NativeRequestContext(),
NativeRequestCapabilities(
execution_mode="async",
websocket_mode=websocket_mode,
requires_connection=requires_connection,
),
NativeResponsesWebSocketRequest(
url=url,
),
options=NativeRequestOptions(
extra_headers=headers,
timeout_seconds=timeout_to_seconds(timeout),
custom_llm_provider=provider,
),
context=NativeRequestContext(),
),
call=lambda connection_type, prepared: call_native(connection_type.connect, prepared),
adapt=ConnectionAdapter,
call=lambda connection_type, request: call_native(connection_type.connect, request),
preflight=lambda: assess_route(_PREFLIGHT, model, provider),
fallback=fallback,
adapt=_ConnectionAdapter,
error_context=BridgeErrorContext(provider=provider, model=model),
)
@asynccontextmanager
async def _connection_context(connection: ConnectionAdapter) -> AsyncGenerator[ConnectionAdapter, None]:
try:
yield connection
finally:
await connection.close()
async def managed_connect(
async def open_connection(
*,
url: str,
headers: dict[str, str],
timeout: float | httpx.Timeout | None,
websocket_mode: str = "managed",
requires_connection: bool = True,
context: NativeRequestContext | None = None,
) -> DispatchResult[AbstractAsyncContextManager[ConnectionAdapter]]:
result: Final = await connect(
url=url,
headers=headers,
timeout=timeout,
websocket_mode=websocket_mode,
requires_connection=requires_connection,
context=context,
)
return adapt_result(result, _connection_context)
model: str,
provider: str,
fallback: Callable[[], AbstractAsyncContextManager[Connection]],
) -> AsyncGenerator[Connection]:
async with AsyncExitStack() as stack:
async def python_connection() -> Connection:
return await stack.enter_async_context(fallback())
backend: Final = await connect(
url=url,
headers=headers,
timeout=timeout,
model=model,
provider=provider,
fallback=python_connection,
)
if backend is None:
raise RuntimeError("WebSocket connection returned no connection")
if isinstance(backend, _ConnectionAdapter):
stack.push_async_callback(backend.close)
yield backend

View file

@ -1,22 +1,34 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from dataclasses import dataclass, field
from enum import Enum
from typing import Final, Generic, TypeAlias, TypeVar
from typing import Final, Generic, NoReturn, Protocol, TypeAlias, TypeVar
from litellm.exceptions import APIError, AuthenticationError, InternalServerError, RateLimitError
from litellm.rust_bridge.bindings import (
UNCHANGED,
NativeBinding,
Unchanged,
native_exception_types,
)
from litellm.rust_bridge.protocols import NativeModule, RustRouteDecline
BindingT = TypeVar("BindingT")
SelectedT = TypeVar("SelectedT")
SelectedSyncT = TypeVar("SelectedSyncT")
SelectedAsyncT = TypeVar("SelectedAsyncT")
NativeT = TypeVar("NativeT")
RequestT = TypeVar("RequestT")
ResultT = TypeVar("ResultT")
SyncBindingT = TypeVar("SyncBindingT")
AsyncBindingT = TypeVar("AsyncBindingT")
class NativeSkipReason(Enum):
DISABLED = "disabled"
INELIGIBLE = "ineligible"
UNAVAILABLE = "unavailable"
DECLINED = "declined"
FAILED = "failed"
class PythonFallbackReason(Enum):
NATIVE_DISABLED = "native_disabled"
NATIVE_UNAVAILABLE = "native_unavailable"
NATIVE_DECLINED = "native_declined"
@dataclass(frozen=True, slots=True)
@ -25,71 +37,494 @@ class Handled(Generic[ResultT]):
@dataclass(frozen=True, slots=True)
class NativeSkipped:
reason: NativeSkipReason
class PythonFallback:
reason: PythonFallbackReason
detail: str | None = None
DispatchResult: TypeAlias = Handled[ResultT] | PythonFallback
@dataclass(frozen=True, slots=True)
class NativeFailed:
error: Exception
class BridgeErrorContext:
provider: str
model: str
DispatchResult: TypeAlias = Handled[ResultT] | NativeSkipped | NativeFailed
class RustEnablement(Protocol):
def __call__(self) -> bool: ...
def _select(load: Callable[[], BindingT | None], enabled: bool, eligible: bool) -> BindingT | NativeSkipped:
if not enabled:
return NativeSkipped(NativeSkipReason.DISABLED)
if not eligible:
return NativeSkipped(NativeSkipReason.INELIGIBLE)
binding: Final = load()
return NativeSkipped(NativeSkipReason.UNAVAILABLE) if binding is None else binding
@dataclass(frozen=True, slots=True)
class EndpointBinding(Generic[BindingT]):
route: str
load: Callable[[], BindingT | None]
enabled: RustEnablement
_native_binding: NativeBinding[BindingT] | None = field(default=None, repr=False)
@staticmethod
def native(
*,
route: str,
select: Callable[[NativeModule], SelectedT],
enabled: RustEnablement,
) -> EndpointBinding[SelectedT]:
binding: Final = NativeBinding(select)
return EndpointBinding(
route=route,
load=binding.load,
enabled=enabled,
_native_binding=binding,
)
def attempt(
*,
load: Callable[[], BindingT | None],
enabled: bool,
eligible: bool,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
) -> DispatchResult[ResultT]:
binding: Final = _select(load, enabled, eligible)
if isinstance(binding, NativeSkipped):
def override(self, value: BindingT | None) -> None:
if self._native_binding is None:
raise RuntimeError("only native Rust bridges support binding overrides")
self._native_binding.override(value)
def reset(self) -> None:
if self._native_binding is None:
raise RuntimeError("only native Rust bridges support binding resets")
self._native_binding.reset()
def _attempt(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> DispatchResult[ResultT]:
binding_or_fallback: Final = self._binding_or_python_fallback(
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return binding_or_fallback
preflight_result: Final = preflight() if preflight is not None else None
if preflight_result is not None:
return preflight_result
return self._attempt_call(
call=lambda: call(binding_or_fallback, prepare()),
adapt=adapt,
error_context=error_context,
)
async def _aattempt(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> DispatchResult[ResultT]:
binding_or_fallback: Final = self._binding_or_python_fallback(
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return binding_or_fallback
preflight_result: Final = preflight() if preflight is not None else None
if preflight_result is not None:
return preflight_result
return await self._attempt_acall(
call=lambda: call(binding_or_fallback, prepare()),
adapt=adapt,
error_context=error_context,
)
def invoke(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], NativeT],
fallback: Callable[[], ResultT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
result: Final = self._attempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
return fallback()
async def ainvoke(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], Awaitable[NativeT]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
result: Final = await self._aattempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
return await fallback()
def assess(
self,
*,
check: Callable[[BindingT], str | None],
) -> PythonFallback | None:
binding: Final = self._binding_or_python_fallback(eligible=True)
if isinstance(binding, PythonFallback):
return binding
reason: Final = check(binding)
return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, reason) if reason is not None else None
def accepts(
self,
*,
check: Callable[[BindingT], str | None],
eligible: bool = True,
) -> bool:
binding_or_fallback: Final = self._binding_or_python_fallback(
eligible=eligible,
)
if isinstance(binding_or_fallback, PythonFallback):
return False
try:
reason: Final = check(binding_or_fallback)
except Exception: # noqa: BLE001 # preflight performs no provider I/O, so Python handoff is safe
return False
return reason is None
def require(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
result: Final = self._attempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
self._raise_required(result)
async def arequire(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
result: Final = await self._aattempt(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
match result:
case Handled(value=value):
return value
case PythonFallback():
self._raise_required(result)
def can_attempt(
self,
*,
eligible: bool = True,
) -> bool:
return not isinstance(
self._binding_or_python_fallback(eligible=eligible),
PythonFallback,
)
def _raise_required(self, fallback: PythonFallback) -> NoReturn:
detail: Final = f": {fallback.detail}" if fallback.detail else ""
reason: Final = _required_reason(fallback.reason)
raise RuntimeError(f"native {self.route} endpoint {reason}{detail}")
def _binding_or_python_fallback(
self,
*,
eligible: bool,
) -> BindingT | PythonFallback:
if not eligible or not self.enabled():
return PythonFallback(PythonFallbackReason.NATIVE_DISABLED)
binding: Final = self.load()
if binding is None:
return PythonFallback(PythonFallbackReason.NATIVE_UNAVAILABLE)
return binding
try:
value: Final = call(binding, prepare())
except Exception as error: # noqa: BLE001 # orchestration applies the endpoint's declared error policy
return NativeFailed(error)
return Handled(adapt(value))
def _attempt_call(
self,
*,
call: Callable[[], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
) -> DispatchResult[ResultT]:
exceptions: Final = native_exception_types()
if exceptions is None:
return Handled(adapt(call()))
declined, upstream = exceptions
try:
value: Final = call()
except declined as error:
return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error))
except upstream as error:
self._raise_upstream(error, error_context)
return Handled(adapt(value))
async def _attempt_acall(
self,
*,
call: Callable[[], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
) -> DispatchResult[ResultT]:
exceptions: Final = native_exception_types()
if exceptions is None:
return Handled(adapt(await call()))
declined, upstream = exceptions
try:
value: Final = await call()
except declined as error:
return PythonFallback(PythonFallbackReason.NATIVE_DECLINED, _error_message(error))
except upstream as error:
self._raise_upstream(error, error_context)
return Handled(adapt(value))
def _raise_upstream(self, error: BaseException, error_context: BridgeErrorContext) -> NoReturn:
args: Final[tuple[object, ...]] = error.args
attribute_status: Final = getattr(error, "status_code", None)
attribute_message: Final = getattr(error, "message", None)
status_value: Final = attribute_status if isinstance(attribute_status, int) else (args[0] if args else 0)
message_value: Final = (
attribute_message if isinstance(attribute_message, str) else (args[1] if len(args) > 1 else str(error))
)
status: Final = status_value if isinstance(status_value, int) else 0
message: Final = message_value if isinstance(message_value, str) else str(message_value)
error_message: Final = f"litellm rust {self.route}: {message}"
if status == 401:
raise AuthenticationError(
message=error_message,
llm_provider=error_context.provider,
model=error_context.model,
) from error
if status == 429:
raise RateLimitError(
message=error_message,
llm_provider=error_context.provider,
model=error_context.model,
) from error
if status == 500:
raise InternalServerError(
message=error_message,
llm_provider=error_context.provider,
model=error_context.model,
) from error
raise APIError(
status_code=status or 500,
message=error_message,
llm_provider=error_context.provider,
model=error_context.model,
) from error
async def aattempt(
*,
load: Callable[[], BindingT | None],
enabled: bool,
eligible: bool,
prepare: Callable[[], RequestT],
call: Callable[[BindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
) -> DispatchResult[ResultT]:
binding: Final = _select(load, enabled, eligible)
if isinstance(binding, NativeSkipped):
return binding
try:
value: Final = await call(binding, prepare())
except Exception as error: # noqa: BLE001 # orchestration applies the endpoint's declared error policy
return NativeFailed(error)
return Handled(adapt(value))
@dataclass(frozen=True, slots=True)
class EndpointDispatch(Generic[SyncBindingT, AsyncBindingT]):
sync: EndpointBinding[SyncBindingT]
asynchronous: EndpointBinding[AsyncBindingT]
@staticmethod
def native(
*,
route: str,
sync: Callable[[NativeModule], SelectedSyncT],
asynchronous: Callable[[NativeModule], SelectedAsyncT],
enabled: RustEnablement,
) -> EndpointDispatch[SelectedSyncT, SelectedAsyncT]:
return EndpointDispatch(
sync=EndpointBinding.native(route=route, select=sync, enabled=enabled),
asynchronous=EndpointBinding.native(
route=route,
select=asynchronous,
enabled=enabled,
),
)
def override(
self,
*,
sync: SyncBindingT | None | Unchanged = UNCHANGED,
asynchronous: AsyncBindingT | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(sync, Unchanged):
self.sync.override(sync)
if not isinstance(asynchronous, Unchanged):
self.asynchronous.override(asynchronous)
def reset(self) -> None:
self.sync.reset()
self.asynchronous.reset()
def invoke(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[SyncBindingT, RequestT], NativeT],
fallback: Callable[[], ResultT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
return self.sync.invoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
async def ainvoke(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]],
fallback: Callable[[], Awaitable[ResultT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
return await self.asynchronous.ainvoke(
prepare=prepare,
call=call,
fallback=fallback,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
def require(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[SyncBindingT, RequestT], NativeT],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
return self.sync.require(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
async def arequire(
self,
*,
prepare: Callable[[], RequestT],
call: Callable[[AsyncBindingT, RequestT], Awaitable[NativeT]],
adapt: Callable[[NativeT], ResultT],
error_context: BridgeErrorContext,
eligible: bool = True,
preflight: Callable[[], PythonFallback | None] | None = None,
) -> ResultT:
return await self.asynchronous.arequire(
prepare=prepare,
call=call,
adapt=adapt,
error_context=error_context,
eligible=eligible,
preflight=preflight,
)
def _error_message(error: BaseException) -> str:
reason: Final[object] = error.args[0] if error.args else str(error)
return reason if isinstance(reason, str) else str(reason)
def _required_reason(reason: PythonFallbackReason) -> str:
match reason:
case PythonFallbackReason.NATIVE_DISABLED:
return "is disabled"
case PythonFallbackReason.NATIVE_UNAVAILABLE:
return "is unavailable"
case PythonFallbackReason.NATIVE_DECLINED:
return "declined the request"
def always_enabled() -> bool:
return True
def identity(value: ResultT) -> ResultT:
return value
def adapt_result(result: DispatchResult[NativeT], adapt: Callable[[NativeT], ResultT]) -> DispatchResult[ResultT]:
if isinstance(result, Handled):
return Handled(adapt(result.value))
return result
async def async_none() -> None:
return None
def assess_route(
binding: EndpointBinding[RustRouteDecline],
model: str,
provider: str,
*,
stream: bool = False,
has_agentic_hook: bool = False,
has_custom_client: bool = False,
request_format: str | None = None,
) -> PythonFallback | None:
return binding.assess(
check=lambda decline: decline(
model,
provider,
stream=stream,
has_agentic_hook=has_agentic_hook,
has_custom_client=has_custom_client,
request_format=request_format,
),
)

View file

@ -1,74 +1,111 @@
from __future__ import annotations
import base64
import json
from collections.abc import Callable, Coroutine
from contextlib import nullcontext
from dataclasses import dataclass
from io import IOBase
from types import MappingProxyType
from typing import Final
import httpx
from pydantic import TypeAdapter
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged
from litellm.rust_bridge.protocols import RustAtranscription, RustTranscription
import litellm
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.rust_bridge.bindings import UNCHANGED, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.protocols import RustAtranscription, RustRouteDecline, RustTranscription
from litellm.rust_bridge.request import (
NativeRequestCapabilities,
NativePreCallDetails,
NativeRequestContext,
NativeRequestOptions,
NativeTranscriptionRequest,
PreparedNativeCall,
bedrock_options,
call_native,
with_capabilities,
)
from litellm.rust_bridge.runtime import DispatchResult, aattempt, attempt, identity
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
EndpointBinding,
EndpointDispatch,
PythonFallback,
always_enabled,
assess_route,
async_none,
identity,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import FileTypes, TranscriptionResponse
_TRANSCRIPTION: Final[NativeBinding[RustTranscription]] = NativeBinding(lambda native: native.transcription)
_ATRANSCRIPTION: Final[NativeBinding[RustAtranscription]] = NativeBinding(lambda native: native.atranscription)
_TRANSCRIPTION: Final[EndpointDispatch[RustTranscription, RustAtranscription]] = EndpointDispatch.native(
route="audio transcription",
sync=lambda native: native.transcription,
asynchronous=lambda native: native.atranscription,
enabled=always_enabled,
)
_PREFLIGHT: Final[EndpointBinding[RustRouteDecline]] = EndpointBinding.native(
route="transcription",
select=lambda native: native.transcription_decline,
enabled=always_enabled,
)
def configure_rust_transcription(
enabled: bool = True,
*,
transcription: RustTranscription | None | Unchanged = UNCHANGED,
atranscription: RustAtranscription | None | Unchanged = UNCHANGED,
decline: RustRouteDecline | None | Unchanged = UNCHANGED,
) -> None:
if not isinstance(decline, Unchanged):
if decline is None:
_PREFLIGHT.reset()
else:
_PREFLIGHT.override(decline)
if not isinstance(transcription, Unchanged):
if transcription is None:
_TRANSCRIPTION.reset()
_TRANSCRIPTION.sync.reset()
else:
_TRANSCRIPTION.override(transcription)
_TRANSCRIPTION.sync.override(transcription)
if not isinstance(atranscription, Unchanged):
if atranscription is None:
_ATRANSCRIPTION.reset()
_TRANSCRIPTION.asynchronous.reset()
else:
_ATRANSCRIPTION.override(atranscription)
_TRANSCRIPTION.asynchronous.override(atranscription)
def load_rust_transcription() -> RustTranscription | None:
return _TRANSCRIPTION.load()
return _TRANSCRIPTION.sync.load()
def load_rust_atranscription() -> RustAtranscription | None:
return _ATRANSCRIPTION.load()
return _TRANSCRIPTION.asynchronous.load()
def transcription(
*,
model: str,
audio: object,
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,
stream: bool = False,
has_custom_client: bool = False,
input_source_kind: str | None = None,
context: NativeRequestContext | None = None,
) -> DispatchResult[dict[str, object]]:
return attempt(
load=_TRANSCRIPTION.load,
enabled=True,
eligible=True,
) -> dict[str, object] | None:
return _TRANSCRIPTION.invoke(
prepare=lambda: PreparedNativeCall(
request=NativeTranscriptionRequest(model=model, audio=audio, optional_params=optional_params),
NativeTranscriptionRequest(
model=model,
audio=audio,
optional_params=optional_params,
),
options=NativeRequestOptions(
api_key=api_key,
api_base=api_base,
@ -77,42 +114,34 @@ def transcription(
timeout_seconds=timeout_to_seconds(timeout),
bedrock=bedrock_options(optional_params),
),
context=with_capabilities(
context or NativeRequestContext(),
NativeRequestCapabilities(
execution_mode="sync",
stream=stream,
has_custom_client=has_custom_client,
input_source_kind=input_source_kind,
),
),
context=NativeRequestContext(),
),
call=call_native,
preflight=lambda: assess_route(_PREFLIGHT, model, custom_llm_provider or ""),
fallback=lambda: None,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
async def atranscription(
*,
model: str,
audio: object,
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,
stream: bool = False,
has_custom_client: bool = False,
input_source_kind: str | None = None,
context: NativeRequestContext | None = None,
) -> DispatchResult[dict[str, object]]:
return await aattempt(
load=_ATRANSCRIPTION.load,
enabled=True,
eligible=True,
) -> dict[str, object] | None:
return await _TRANSCRIPTION.ainvoke(
prepare=lambda: PreparedNativeCall(
request=NativeTranscriptionRequest(model=model, audio=audio, optional_params=optional_params),
NativeTranscriptionRequest(
model=model,
audio=audio,
optional_params=optional_params,
),
options=NativeRequestOptions(
api_key=api_key,
api_base=api_base,
@ -121,16 +150,166 @@ async def atranscription(
timeout_seconds=timeout_to_seconds(timeout),
bedrock=bedrock_options(optional_params),
),
context=with_capabilities(
context or NativeRequestContext(),
NativeRequestCapabilities(
execution_mode="async",
stream=stream,
has_custom_client=has_custom_client,
input_source_kind=input_source_kind,
),
),
context=NativeRequestContext(),
),
call=call_native,
preflight=lambda: assess_route(_PREFLIGHT, model, custom_llm_provider or ""),
fallback=async_none,
adapt=identity,
error_context=BridgeErrorContext(provider=custom_llm_provider or "", model=model),
)
TranscriptionResult = TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]
@dataclass
class _TranscriptionOperation:
model: str
provider: str
file: FileTypes
api_key: str | None
api_base: str | None
headers: dict[str, object] | None
optional_params: dict[str, object]
timeout: float | httpx.Timeout | None
logging: Logging
python: Callable[[FileTypes], TranscriptionResult]
fallback_file: FileTypes | None = None
logged: bool = False
def prepare(self) -> PreparedNativeCall[NativeTranscriptionRequest]:
key: Final = (
self.api_key
or litellm.api_key
or TypeAdapter(str | None).validate_python(getattr(litellm, f"{self.provider}_key", None))
or get_secret_str(f"{self.provider.upper()}_API_KEY")
)
base: Final = (
self.api_base
or litellm.api_base
or get_secret_str(f"{self.provider.upper()}_BASE_URL")
or get_secret_str(f"{self.provider.upper()}_API_BASE")
)
content: Final = self.file[1] if isinstance(self.file, tuple) else self.file
position: Final = content.tell() if isinstance(content, IOBase) and content.seekable() else None
try:
processed: Final = process_audio_file(self.file)
finally:
if position is not None and isinstance(content, IOBase):
content.seek(position)
self.fallback_file = (processed.filename, processed.file_content, processed.content_type)
audio: Final = TypeAdapter(dict[str, object]).validate_python(
MappingProxyType(
{
"data": base64.b64encode(processed.file_content).decode("ascii"),
"format": processed.filename.rsplit(".", 1)[-1].lower() if "." in processed.filename else "wav",
"filename": processed.filename,
}
)
)
log_details: Final[NativePreCallDetails] = {
"api_base": base or "",
"headers": self.headers,
"complete_input_dict": {"model": self.model, **self.optional_params},
}
self.logging.pre_call(input="audio transcription", api_key=key, additional_args=log_details)
self.logged = True
return PreparedNativeCall(
NativeTranscriptionRequest(
model=self.model,
audio=audio,
optional_params=provider_request_params(self.optional_params),
options=NativeRequestOptions(
api_key=key,
api_base=base,
custom_llm_provider=self.provider,
extra_headers=self.headers,
timeout_seconds=timeout_to_seconds(self.timeout),
provider_connection=provider_connection_params(self.optional_params),
),
)
)
def fallback(self) -> TranscriptionResult:
with self.logging.suppress_next_pre_call() if self.logged else nullcontext():
return self.python(self.fallback_file if self.fallback_file is not None else self.file)
async def afallback(self) -> TranscriptionResponse:
with self.logging.suppress_next_pre_call() if self.logged else nullcontext():
result: Final = self.python(self.fallback_file if self.fallback_file is not None else self.file)
return await result if isinstance(result, Coroutine) else result
def adapt(self, response: dict[str, object]) -> TranscriptionResponse:
text: Final = TypeAdapter(str).validate_python(response["text"])
parsed: Final = TranscriptionResponse(text=text)
self.logging.post_call(
input="audio transcription", api_key=self.api_key, original_response=json.dumps(response)
)
return parsed
def dispatch_transcription(
*,
model: str,
provider: str,
file: FileTypes,
api_key: str | None,
api_base: str | None,
headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout: float | httpx.Timeout | None,
logging: Logging,
asynchronous: bool,
has_custom_client: bool,
fallback: Callable[[FileTypes], TranscriptionResult],
) -> TranscriptionResult:
operation: Final = _TranscriptionOperation(
model, provider, file, api_key, api_base, headers, optional_params, timeout, logging, fallback
)
def preflight() -> PythonFallback | None:
return assess_route(
_PREFLIGHT,
model,
provider,
stream=optional_params.get("stream") is True,
has_custom_client=has_custom_client,
)
error_context: Final = BridgeErrorContext(provider=provider, model=model)
if provider == "bedrock":
if asynchronous:
return _TRANSCRIPTION.arequire(
prepare=operation.prepare,
call=call_native,
adapt=operation.adapt,
error_context=error_context,
preflight=preflight,
)
return _TRANSCRIPTION.require(
prepare=operation.prepare,
call=call_native,
adapt=operation.adapt,
error_context=error_context,
preflight=preflight,
)
if asynchronous:
return _TRANSCRIPTION.ainvoke(
prepare=operation.prepare,
call=call_native,
adapt=operation.adapt,
fallback=operation.afallback,
error_context=error_context,
eligible=rust_enabled(),
preflight=preflight,
)
return _TRANSCRIPTION.invoke(
prepare=operation.prepare,
call=call_native,
adapt=operation.adapt,
fallback=operation.fallback,
error_context=error_context,
eligible=rust_enabled(),
preflight=preflight,
)

View file

@ -1,20 +1,14 @@
"""Tests for the optional Rust-backed Anthropic Messages path."""
import importlib
from typing import cast
import httpx
import pytest
from pydantic import ValidationError
import litellm
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import configuration
from litellm.rust_bridge.request import NativeMessagesRequest, NativeRequestContext, NativeRequestOptions
from litellm.rust_bridge.runtime import Handled, NativeFailed, NativeSkipped, NativeSkipReason
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.rust_bridge.request import NativeMessagesRequest, NativeRequestContext
rust_messages = importlib.import_module("litellm.rust_bridge.messages")
rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
@ -39,13 +33,12 @@ REQUEST_BODY: dict[str, object] = {
class RecordingMessages:
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
self.contexts: list[NativeRequestContext] = []
def __call__(
self,
request: NativeMessagesRequest,
*,
options: NativeRequestOptions,
options: object,
context: NativeRequestContext,
) -> dict[str, object]:
self.calls.append(
@ -59,20 +52,18 @@ class RecordingMessages:
"timeout_seconds": options.timeout_seconds,
}
)
self.contexts.append(context)
return dict(FAKE_MESSAGES_RESPONSE)
class RecordingAsyncMessages:
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
self.contexts: list[NativeRequestContext] = []
async def __call__(
self,
request: NativeMessagesRequest,
*,
options: NativeRequestOptions,
options: object,
context: NativeRequestContext,
) -> dict[str, object]:
self.calls.append(
@ -86,7 +77,6 @@ class RecordingAsyncMessages:
"timeout_seconds": options.timeout_seconds,
}
)
self.contexts.append(context)
return dict(FAKE_MESSAGES_RESPONSE)
@ -94,7 +84,9 @@ class ExplodingAsyncMessages:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, *args: object, **kwargs: object) -> dict[str, object]:
async def __call__(
self, request: NativeMessagesRequest, *, options: object, context: NativeRequestContext
) -> dict[str, object]:
self.calls += 1
raise AssertionError("bridge must not be called")
@ -103,18 +95,28 @@ class RaisingAsyncMessages:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, *args: object, **kwargs: object) -> dict[str, object]:
async def __call__(
self, request: NativeMessagesRequest, *, options: object, context: NativeRequestContext
) -> 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)
rust_messages.set_rust_messages(messages=None, amessages=None, decline=None)
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
rust_messages.set_rust_messages(
decline=lambda model, custom_llm_provider, **features: (
"unsupported feature"
if any(features.get(key) for key in ("stream", "has_agentic_hook", "has_custom_client"))
or features.get("request_format") == "native"
else None
)
)
yield
rust_messages.set_rust_messages(messages=None, amessages=None)
rust_messages.set_rust_messages(messages=None, amessages=None, decline=None)
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
@ -126,6 +128,16 @@ def test_load_rust_messages_returns_injected_impl():
assert rust_messages.load_rust_messages() is bridge
def test_bare_rust_still_toggles_ocr():
from litellm.rust_bridge.ocr import rust_ocr_enabled
litellm.rust(True)
assert rust_ocr_enabled() is True
litellm.rust(False)
assert rust_ocr_enabled() is False
def test_load_rust_amessages_returns_injected_impl():
bridge = RecordingAsyncMessages()
litellm.rust(True)
@ -133,7 +145,7 @@ def test_load_rust_amessages_returns_injected_impl():
assert rust_messages.load_rust_amessages() is bridge
def test_messages_wrapper_reports_unavailable(monkeypatch):
def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge.bindings"),
"get_native_bridge",
@ -150,7 +162,7 @@ def test_messages_wrapper_reports_unavailable(monkeypatch):
extra_headers={},
timeout=30.0,
)
assert result == NativeSkipped(NativeSkipReason.UNAVAILABLE)
assert result is None
def test_messages_wrapper_forwards_args_and_converts_timeout():
@ -168,7 +180,7 @@ def test_messages_wrapper_forwards_args_and_converts_timeout():
timeout=httpx.Timeout(600.0, read=42.0),
)
assert response == Handled(FAKE_MESSAGES_RESPONSE)
assert response == FAKE_MESSAGES_RESPONSE
assert bridge.calls[0] == {
"model": "claude-sonnet-4-5",
"body": REQUEST_BODY,
@ -196,331 +208,116 @@ async def test_amessages_wrapper_forwards_args():
timeout=12.5,
)
assert response == Handled(FAKE_MESSAGES_RESPONSE)
assert response == FAKE_MESSAGES_RESPONSE
assert bridge.calls[0]["model"] == "claude-sonnet-4-5"
assert bridge.calls[0]["timeout_seconds"] == 12.5
class PythonMessages:
def __init__(self) -> None:
self.calls = 0
def anthropic_messages_handler(self, **kwargs: object) -> dict[str, object]:
self.calls += 1
return {**FAKE_MESSAGES_RESPONSE, "id": "python"}
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "openai"])
@pytest.mark.asyncio
async def test_amessages_wrapper_preserves_capability_facts():
bridge = RecordingAsyncMessages()
rust_messages.set_rust_messages(amessages=bridge)
await rust_messages.amessages(
model="claude-sonnet-4-5",
body=REQUEST_BODY,
api_key=None,
api_base=None,
custom_llm_provider="anthropic",
extra_headers=None,
timeout=None,
stream=True,
has_custom_client=True,
has_agentic_hook=True,
)
capabilities = bridge.contexts[0].capabilities
assert capabilities.execution_mode == "async"
assert capabilities.stream is True
assert capabilities.has_custom_client is True
assert capabilities.has_agentic_hook is True
def _gate(**overrides):
kwargs = {
"custom_llm_provider": "azure_ai",
"litellm_params": GenericLiteLLMParams(api_key="sk-azure"),
"has_agentic_hook": False,
"stream": False,
"has_custom_client": False,
"model": "claude-sonnet-4-5",
"api_key": "sk-azure",
"api_base": "https://resource.services.ai.azure.com/anthropic",
"headers": {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"},
"request_body": dict(REQUEST_BODY),
"timeout": 30.0,
}
kwargs.update(overrides)
return BaseLLMHTTPHandler._attempt_rust_anthropic_messages(**kwargs)
@pytest.mark.asyncio
async def test_gate_invokes_rust_and_marks_response_header():
bridge = RecordingAsyncMessages()
async def test_public_messages_routes_provider_acceptance(monkeypatch, asynchronous, provider):
bridge = RecordingMessages()
async_bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate()
assert isinstance(response, Handled)
response = response.value
rust_messages.set_rust_messages(messages=bridge, amessages=async_bridge)
if asynchronous:
response = await litellm.anthropic.messages.acreate(
model=f"{provider}/test-model",
max_tokens=64,
messages=[{"role": "user", "content": "hi"}],
api_key="key",
api_base="https://example.test",
)
else:
response = litellm.anthropic.messages.create(
model=f"{provider}/test-model",
max_tokens=64,
messages=[{"role": "user", "content": "hi"}],
api_key="key",
api_base="https://example.test",
)
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
assert response["_hidden_params"]["additional_headers"]["x-litellm-rust"] == "true"
calls = async_bridge.calls if asynchronous else bridge.calls
assert len(calls) == 1
assert calls[0]["custom_llm_provider"] == provider
assert calls[0]["body"]["max_tokens"] == 64
@pytest.mark.asyncio
async def test_gate_reports_failure_to_harness():
bridge = RaisingAsyncMessages()
@pytest.mark.parametrize("condition", ["disabled", "declined", "missing_binding", "missing_preflight", "stream"])
def test_public_messages_fallback_once(monkeypatch, condition):
module = importlib.import_module("litellm.llms.anthropic.experimental_pass_through.messages.handler")
python = PythonMessages()
monkeypatch.setattr(module, "base_llm_http_handler", python)
bridge = RecordingMessages()
litellm.rust(condition != "disabled")
rust_messages.set_rust_messages(messages=bridge)
if condition == "declined":
rust_messages.set_rust_messages(decline=lambda model, custom_llm_provider, **features: "unsupported provider")
elif condition == "missing_binding":
rust_messages._MESSAGES.sync.override(None)
elif condition == "missing_preflight":
rust_messages._PREFLIGHT.override(None)
litellm.anthropic.messages.create(
model="anthropic/test-model",
max_tokens=64,
messages=[{"role": "user", "content": "hi"}],
api_key="key",
stream=condition == "stream",
)
assert python.calls == 1
assert bridge.calls == []
@pytest.mark.parametrize("response", [{}, {"content": "invalid"}])
def test_public_messages_invalid_response_does_not_fallback(monkeypatch, response):
module = importlib.import_module("litellm.llms.anthropic.experimental_pass_through.messages.handler")
python = PythonMessages()
monkeypatch.setattr(module, "base_llm_http_handler", python)
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate()
assert isinstance(response, NativeFailed)
assert bridge.calls == 1
rust_messages.set_rust_messages(messages=lambda request, *, context, callback_adapter=None: response)
with pytest.raises(ValidationError):
litellm.anthropic.messages.create(
model="anthropic/test-model",
max_tokens=64,
messages=[{"role": "user", "content": "hi"}],
api_key="key",
)
assert python.calls == 0
@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 isinstance(response, NativeSkipped)
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 isinstance(response, Handled)
response = response.value
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"},
def test_public_messages_preserves_headers_and_optional_parameters(monkeypatch):
handler = importlib.import_module("litellm.llms.anthropic.experimental_pass_through.messages.handler")
monkeypatch.setattr(handler, "is_reasoning_auto_summary_enabled", lambda: True)
configuration.rust(True)
native = RecordingMessages()
rust_messages.set_rust_messages(messages=native)
response = litellm.anthropic.messages.create(
model="anthropic/claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
max_tokens=2048,
thinking={"type": "enabled", "budget_tokens": 1024},
temperature=0.8,
additional_drop_params=["temperature"],
api_key="key",
headers={"x-source": "forwarded"},
extra_headers={"x-source": "extra"},
provider_specific_header={
"custom_llm_provider": "anthropic",
"extra_headers": {"x-source": "scoped"},
},
)
assert isinstance(response, Handled)
response = response.value
assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
assert bridge.calls[0]["custom_llm_provider"] == "anthropic"
assert bridge.calls[0]["api_key"] == "sk-ant"
@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 isinstance(response, Handled)
response = response.value
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 isinstance(response, NativeSkipped)
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 isinstance(response, NativeSkipped)
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 isinstance(response, NativeSkipped)
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 isinstance(response, Handled)
response = response.value
assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
assert "stream" not in bridge.calls[0]["body"]
assert bridge.calls[0]["body"] == REQUEST_BODY
@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.bindings"),
"get_native_bridge",
lambda: None,
)
litellm.rust(True)
response = await _gate()
assert isinstance(response, NativeSkipped)
@pytest.mark.asyncio
@pytest.mark.parametrize("selection", ("native", "disabled", "failed", "declined", "upstream"))
async def test_messages_handler_runs_selected_backend_once(selection: str, monkeypatch: pytest.MonkeyPatch) -> None:
from datetime import datetime
from types import SimpleNamespace
import httpx
from litellm.exceptions import RateLimitError
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import AnthropicMessagesConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.rust_bridge import bindings
class Declined(Exception):
pass
class Upstream(Exception):
pass
error = (
Upstream(429, "rate limited")
if selection == "upstream"
else Declined("unsupported")
if selection == "declined"
else RuntimeError("native failed")
if selection == "failed"
else None
)
class Native:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, *args: object, **kwargs: object) -> dict[str, object]:
self.calls += 1
if error is not None:
raise error
return dict(FAKE_MESSAGES_RESPONSE)
bridge = Native()
monkeypatch.setattr(
bindings,
"get_native_bridge",
lambda: SimpleNamespace(
RustBridgeDeclined=Declined,
RustUpstreamError=Upstream,
),
)
rust_messages.set_rust_messages(amessages=bridge)
litellm.rust(selection != "disabled")
requests: list[httpx.Request] = []
def respond(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(200, json=FAKE_MESSAGES_RESPONSE)
logging_obj = Logging(
model=FAKE_MESSAGES_RESPONSE["model"],
messages=[],
stream=False,
call_type="anthropic_messages",
start_time=datetime.now(),
litellm_call_id="harness-test",
function_id="harness-test",
)
client = AsyncHTTPHandler()
await client.client.aclose()
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as transport:
client.client = transport
async def run():
return await BaseLLMHTTPHandler().async_anthropic_messages_handler(
model=FAKE_MESSAGES_RESPONSE["model"],
messages=[{"role": "user", "content": "hello"}],
anthropic_messages_provider_config=AnthropicMessagesConfig(),
anthropic_messages_optional_request_params={"max_tokens": 10},
custom_llm_provider="anthropic",
litellm_params=GenericLiteLLMParams(),
logging_obj=logging_obj,
api_key="sk-test",
api_base="https://example.test",
client=client,
)
if selection in ("failed", "upstream"):
with pytest.raises(RateLimitError if selection == "upstream" else RuntimeError) as caught:
await run()
if selection == "upstream":
assert caught.value.__cause__ is error
assert caught.value.llm_provider == "anthropic"
assert caught.value.model == FAKE_MESSAGES_RESPONSE["model"]
else:
assert caught.value is error
else:
response = await run()
assert response["id"] == FAKE_MESSAGES_RESPONSE["id"]
assert len(requests) == (1 if selection in ("disabled", "declined") else 0)
assert bridge.calls == (0 if selection == "disabled" else 1)
assert response["id"] == "msg_123"
assert native.calls[0]["extra_headers"]["x-source"] == "scoped"
assert native.calls[0]["body"]["thinking"]["display"] == "summarized"
assert "temperature" not in native.calls[0]["body"]

View file

@ -2223,398 +2223,3 @@ 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
)
@staticmethod
def _completion_kwargs(**overrides):
from litellm.types.utils import ModelResponse
kwargs = {
"model": "claude-sonnet-4-5",
"messages": [{"role": "user", "content": "hi"}],
"api_base": "https://api.anthropic.com/v1/messages",
"custom_llm_provider": "anthropic",
"custom_prompt_dict": {},
"model_response": ModelResponse(),
"print_verbose": lambda *_args, **_kwargs: None,
"encoding": None,
"api_key": "sk-ant-test",
"logging_obj": MagicMock(),
"optional_params": {"max_tokens": 16},
"timeout": 30.0,
"litellm_params": {},
"acompletion": False,
"headers": {},
"client": None,
}
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(request, *, options, context):
seen["call"].append(
{
"model": request.model,
"messages": request.messages,
"optional_params": request.optional_params,
"api_key": options.api_key,
"api_base": options.api_base,
"context": context,
}
)
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")
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": []}
) as transform, patch.object(
AnthropicChatCompletion, "acompletion_function"
):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(litellm_params={})
)
except Exception:
# The Python path goes on to make an HTTP call; reaching it is
# the assertion, so the network failure below is expected.
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["call"] == []
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(_request, *, options, context):
raise _Declined("blank message text")
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
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("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
async def declining_native(_request, *, options, context):
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(_request, *, options, context):
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("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
def declining_native(_request, *, options, context):
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"] == {
"model": "m",
"messages": [],
}

View file

@ -1,555 +1,69 @@
"""Tests for `BedrockConverseLLM.completion`'s Rust chat completions hook.
The native callables are dependency-injected, so these run without the compiled
extension, and AWS credential resolution is stubbed so nothing reaches STS.
"""
from __future__ import annotations
from unittest.mock import MagicMock, patch
import httpx
import pytest
from botocore.credentials import Credentials
from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
import litellm
from litellm.rust_bridge import chat_completions as bridge
from litellm.types.utils import ModelResponse
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",
token="resolved-token",
)
from litellm.rust_bridge.request import NativeChatCompletionsRequest, NativeRequestContext
@pytest.fixture(autouse=True)
def reset_bridge(monkeypatch):
def native_bridge(monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
bridge.set_rust_chat_completions(decline=lambda **features: None)
yield
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
def _inject(*, decline_reason=None, error: Exception | None = None):
seen: dict[str, list[dict]] = {"gate": [], "call": []}
def test_native_bedrock_receives_explicit_auth_and_endpoint():
requests = []
def gate(**kwargs):
seen["gate"].append(kwargs)
return decline_reason
def native(request: NativeChatCompletionsRequest, *, context: NativeRequestContext, callback_adapter=None):
requests.append(request)
return {
"choices": [{"index": 0, "message": {"role": "assistant", "content": "native"}, "finish_reason": "stop"}]
}
def native(request, *, options, context):
seen["call"].append(
{
"model": request.model,
"messages": request.messages,
"optional_params": request.optional_params,
"api_key": options.api_key,
"api_base": options.api_base,
"context": context,
}
)
if error is not None:
raise error
return dict(RUST_RESPONSE)
bridge.set_rust_chat_completions(decline=gate, chat_completions=native)
return seen
def _completion_kwargs(**overrides):
kwargs = {
"model": "bedrock/us-east-1/anthropic.claude-sonnet-4-5-v1:0",
"messages": [{"role": "user", "content": "hi"}],
"api_base": None,
"custom_prompt_dict": {},
"model_response": ModelResponse(),
"encoding": None,
"logging_obj": MagicMock(),
"optional_params": {"maxTokens": 16},
"acompletion": False,
"timeout": 30.0,
"litellm_params": {},
"extra_headers": None,
"client": None,
"api_key": None,
}
kwargs.update(overrides)
return kwargs
def _run(*, credentials: Credentials | None = RESOLVED_CREDENTIALS, **overrides):
with patch.object(BedrockConverseLLM, "get_credentials", return_value=credentials):
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()
bridge.set_rust_chat_completions(
decline=pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline
)
try:
_run(optional_params={"maxTokens": 16, "stream": True})
except Exception:
pass
assert seen["call"] == []
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("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
async def declining_native(_request, **_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(_request, **_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(_request, **_kwargs):
raise _Declined("blank message text")
logging_obj = MagicMock()
served = []
async def python_path(**kwargs):
served.append(kwargs)
return ModelResponse()
with (
patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()),
patch.object(
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
),
patch.object(
BedrockConverseLLM, "async_completion", side_effect=python_path
),
):
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",
"usage": {"inputTokens": 5, "outputTokens": 2, "totalTokens": 7},
}
async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj):
"""Run the real `async_completion` with a stubbed transport."""
import httpx as _httpx
client = MagicMock()
async def post(**_kwargs):
return _httpx.Response(
200,
json=CONVERSE_RESPONSE,
request=_httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"),
)
client.post = post
client.__class__ = AsyncHTTPHandler
return await BedrockConverseLLM().async_completion(
model="anthropic.claude-sonnet-4-5-v1:0",
bridge.set_rust_chat_completions(chat_completions=native)
response = litellm.completion(
model="bedrock/anthropic.claude-sonnet-4-5-v1:0",
messages=[{"role": "user", "content": "hi"}],
api_base="https://bedrock-runtime.us-west-2.amazonaws.com/model/m/converse",
model_response=ModelResponse(),
timeout=30.0,
encoding=None,
logging_obj=logging_obj,
stream=None,
optional_params={"maxTokens": 16},
litellm_params={"aws_region_name": "us-west-2"},
credentials=RESOLVED_CREDENTIALS,
headers={},
client=client,
skip_pre_call_logging=skip_pre_call_logging,
aws_access_key_id="explicit-id",
aws_secret_access_key="explicit-secret",
aws_session_token="explicit-token",
aws_region_name="us-east-1",
api_base="https://example.test",
max_tokens=7,
)
assert response.choices[0].message.content == "native"
assert len(requests) == 1
assert requests[0].options.provider_connection == {
"aws_access_key_id": "explicit-id",
"aws_secret_access_key": "explicit-secret",
"aws_session_token": "explicit-token",
"aws_region_name": "us-east-1",
}
assert requests[0].options.api_base == "https://example.test"
@pytest.mark.asyncio
async def test_async_completion_honors_the_pre_call_suppression():
logging_obj = MagicMock()
await _drive_async_completion(skip_pre_call_logging=True, logging_obj=logging_obj)
assert logging_obj.pre_call.call_count == 0
@pytest.mark.parametrize("through_environment", [False, True])
def test_native_bedrock_preserves_bearer_auth(monkeypatch, through_environment):
requests = []
def native(request, *, context, callback_adapter=None):
requests.append(request)
return {
"choices": [{"index": 0, "message": {"role": "assistant", "content": "native"}, "finish_reason": "stop"}]
}
@pytest.mark.asyncio
async def test_async_completion_logs_pre_call_by_default():
"""The suppression must be opt-in, so every existing caller keeps its log."""
logging_obj = MagicMock()
await _drive_async_completion(skip_pre_call_logging=False, logging_obj=logging_obj)
assert logging_obj.pre_call.call_count == 1
def _sync_client_returning_converse_response():
client = MagicMock()
client.post.side_effect = lambda **_kwargs: httpx.Response(
200,
json=CONVERSE_RESPONSE,
request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"),
bridge.set_rust_chat_completions(chat_completions=native)
if through_environment:
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-token")
litellm.completion(
model="bedrock/anthropic.claude-sonnet-4-5-v1:0",
messages=[{"role": "user", "content": "hi"}],
api_key=None if through_environment else "bedrock-token",
max_tokens=7,
)
client.__class__ = HTTPHandler
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(_request, **_kwargs):
raise _Declined("blank message text")
logging_obj = MagicMock()
with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()):
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
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")
logging_obj = MagicMock()
response = _run(
logging_obj=logging_obj,
litellm_params={},
client=_sync_client_returning_converse_response(),
)
assert response.choices[0].message.content == "hi"
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(_request, **_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(_request, **_kwargs):
raise _Declined("blank message text")
logging_obj, calls = _recording_logging_obj()
with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()):
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
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")
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
client = _sync_client_returning_converse_response()
response = _run(credentials=None, litellm_params={}, client=client)
assert response.choices[0].message.content == "hi"
sent_headers = client.post.call_args.kwargs["headers"]
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:
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
client = _sync_client_returning_converse_response()
response = BedrockConverseLLM().completion(
**_completion_kwargs(
optional_params={"maxTokens": 16, "aws_profile_name": "litellm-no-such-aws-profile"},
litellm_params={},
client=client,
api_key="bedrock-bearer-token" if configured_through == "api_key" else None,
)
)
assert response.choices[0].message.content == "hi"
assert client.post.call_args.kwargs["headers"]["Authorization"] == "Bearer bedrock-bearer-token"
assert len(requests) == 1
assert requests[0].options.api_key == (None if through_environment else "bedrock-token")

View file

@ -151,7 +151,7 @@ def test_response_api_handler_runs_agentic_hooks_in_sync_path(monkeypatch):
)
logging_obj = Mock()
monkeypatch.setattr(handler, "_has_agentic_completion_hook", Mock(return_value=True))
monkeypatch.setattr(handler, "has_agentic_completion_hook", Mock(return_value=True))
hook_mock = AsyncMock(return_value=final_response)
monkeypatch.setattr(handler, "_call_agentic_completion_hooks", hook_mock)
@ -370,7 +370,7 @@ def test_get_agentic_loop_settings_defaults_and_overrides():
assert fingerprints == ["fp-1", "fp-2"]
def test_has_agentic_completion_hook_detection(monkeypatch):
def testhas_agentic_completion_hook_detection(monkeypatch):
"""The streaming path skips the agentic wrapper only when no callback
overrides async_should_run_agentic_loop. Verify both directions."""
from litellm.integrations.custom_logger import CustomLogger
@ -381,7 +381,7 @@ def test_has_agentic_completion_hook_detection(monkeypatch):
# No callbacks at all -> no agentic hook.
monkeypatch.setattr(litellm, "callbacks", [])
assert handler._has_agentic_completion_hook(logging_obj) is False
assert handler.has_agentic_completion_hook(logging_obj) is False
# A plain CustomLogger that does NOT override the gate -> still no hook
# (so the wrapper is safely skipped).
@ -389,7 +389,7 @@ def test_has_agentic_completion_hook_detection(monkeypatch):
pass
monkeypatch.setattr(litellm, "callbacks", [_PlainLogger()])
assert handler._has_agentic_completion_hook(logging_obj) is False
assert handler.has_agentic_completion_hook(logging_obj) is False
# A logger that overrides the gate (directly) -> hook present.
class _AgenticLogger(CustomLogger):
@ -399,7 +399,7 @@ def test_has_agentic_completion_hook_detection(monkeypatch):
return True, {}
monkeypatch.setattr(litellm, "callbacks", [_AgenticLogger()])
assert handler._has_agentic_completion_hook(logging_obj) is True
assert handler.has_agentic_completion_hook(logging_obj) is True
# Override inherited through an intermediate class is still detected
# (function-identity check, not a leaf __dict__ check).
@ -407,12 +407,12 @@ def test_has_agentic_completion_hook_detection(monkeypatch):
pass
monkeypatch.setattr(litellm, "callbacks", [_DerivedAgenticLogger()])
assert handler._has_agentic_completion_hook(logging_obj) is True
assert handler.has_agentic_completion_hook(logging_obj) is True
# Hook supplied via logging_obj.dynamic_success_callbacks is detected too.
monkeypatch.setattr(litellm, "callbacks", [])
logging_obj.dynamic_success_callbacks = [_AgenticLogger()]
assert handler._has_agentic_completion_hook(logging_obj) is True
assert handler.has_agentic_completion_hook(logging_obj) is True
# String-named callback entry (e.g. "datadog") must be resolved to its
# CustomLogger instance via get_custom_logger_compatible_class -- the same
@ -426,7 +426,7 @@ def test_has_agentic_completion_hook_detection(monkeypatch):
"litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class",
lambda name: agentic_via_string if name == "fake_string_callback" else None,
)
assert handler._has_agentic_completion_hook(logging_obj) is True
assert handler.has_agentic_completion_hook(logging_obj) is True
# Unresolvable string (returns None) is skipped, no false positive.
monkeypatch.setattr(litellm, "callbacks", ["unknown_callback"])
@ -434,7 +434,7 @@ def test_has_agentic_completion_hook_detection(monkeypatch):
"litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class",
lambda name: None,
)
assert handler._has_agentic_completion_hook(logging_obj) is False
assert handler.has_agentic_completion_hook(logging_obj) is False
def test_fingerprint_agentic_tools_is_deterministic():
@ -1181,7 +1181,7 @@ def test_sync_delete_responses_sets_json_content_type():
({}, True, None, None),
],
)
def test_resolve_anthropic_messages_timeout(
def testresolve_anthropic_messages_timeout(
monkeypatch, litellm_params_kwargs, stream, global_timeout, expected
):
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
@ -1203,7 +1203,7 @@ def test_resolve_anthropic_messages_timeout(
"litellm.request_timeout_explicitly_set", True, raising=False
)
resolved = BaseLLMHTTPHandler._resolve_anthropic_messages_timeout(
resolved = BaseLLMHTTPHandler.resolve_anthropic_messages_timeout(
litellm_params=GenericLiteLLMParams(**litellm_params_kwargs),
stream=stream,
custom_llm_provider="anthropic",

View file

@ -210,13 +210,12 @@ def build_prepared_request(
@pytest.fixture(autouse=True)
def _reset_rust_flag():
"""Keep the global toggle isolated between tests."""
rust_bridge._OCR.reset()
rust_bridge._AOCR.reset()
rust_bridge.set_rust_ocr(ocr=None, aocr=None, decline=None)
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
rust_bridge.set_rust_ocr(decline=lambda model, custom_llm_provider, **features: "unsupported feature" if any(features.get(key) for key in ("stream", "has_agentic_hook", "has_custom_client")) or features.get("request_format") == "native" else None)
yield
rust_bridge._OCR.reset()
rust_bridge._AOCR.reset()
rust_bridge.set_rust_ocr(ocr=None, aocr=None, decline=None)
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
@ -226,7 +225,7 @@ def fake_bridge():
"""Enable the Rust path with an injected recording bridge (no native wheel)."""
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
return bridge
@ -235,14 +234,27 @@ def fake_async_bridge():
"""Enable the async Rust path with an injected recording bridge."""
bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge._AOCR.override(bridge)
rust_bridge.set_rust_ocr(aocr=bridge)
return bridge
def test_rust_toggles_flag():
assert rust_bridge.rust_ocr_enabled() is False
litellm.rust(True)
assert rust_bridge.rust_ocr_enabled() is True
litellm.rust(False)
assert rust_bridge.rust_ocr_enabled() is False
def test_env_var_enables_rust_ocr(monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
assert rust_bridge.rust_ocr_enabled() is True
def test_load_rust_ocr_returns_injected_impl():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
assert rust_bridge.load_rust_ocr() is bridge
@ -306,7 +318,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch):
def test_load_rust_aocr_returns_injected_impl():
bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge._AOCR.override(bridge)
rust_bridge.set_rust_ocr(aocr=bridge)
assert rust_bridge.load_rust_aocr() is bridge
@ -315,8 +327,7 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
bridge = RecordingBridge()
async_bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._AOCR.override(async_bridge)
rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge)
litellm.rust(False)
assert rust_bridge.load_rust_ocr() is bridge
@ -335,11 +346,9 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
bridge = RecordingBridge()
async_bridge = RecordingAsyncBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge._AOCR.override(async_bridge)
rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge)
rust_bridge._OCR.override(None)
rust_bridge._AOCR.override(None)
rust_bridge.set_rust_ocr(ocr=None, aocr=None, decline=None)
assert rust_bridge.load_rust_ocr() is None
assert rust_bridge.load_rust_aocr() is None
@ -385,7 +394,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
bridge = RecordingBridge()
logging_obj = RecordingLogging()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
response = rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
@ -423,7 +432,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(api_key=None, timeout=None),
@ -436,7 +445,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
def test_run_rust_ocr_prefers_explicit_key_over_resolver():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
def _resolver(name: str) -> str | None:
raise AssertionError(f"resolver should not be called for {name}")
@ -456,7 +465,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
bridge = RecordingBridge()
resolver_calls = []
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
def _resolver(name):
resolver_calls.append(name)
@ -479,7 +488,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var():
def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
@ -506,7 +515,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
def _resolver(name: str) -> str | None:
return {
@ -530,7 +539,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana
def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
@ -548,7 +557,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint():
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
@ -569,7 +578,7 @@ def test_run_rust_ocr_runs_pre_call_logging():
logging_obj = RecordingLogging()
bridge = RecordingBridge()
litellm.rust(True)
rust_bridge._OCR.override(bridge)
rust_bridge.set_rust_ocr(ocr=bridge)
rust_bridge.attempt_ocr(
prepared_request=build_prepared_request(
@ -633,7 +642,7 @@ def test_ocr_exception_type_uses_resolved_provider_context(
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
litellm.rust(True)
rust_bridge._OCR.override(RaisingBridge())
rust_bridge.set_rust_ocr(ocr=RaisingBridge())
with pytest.raises(CapturedException):
litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
@ -679,7 +688,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context(
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
litellm.rust(True)
rust_bridge._AOCR.override(RaisingAsyncBridge())
rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge())
with pytest.raises(CapturedException):
await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")

View file

@ -2,14 +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.request import (
NativeRequestContext,
NativeRequestOptions,
NativeResponsesWebSocketRequest,
)
from litellm.rust_bridge.runtime import Handled, NativeFailed, NativeSkipped, NativeSkipReason
from litellm.rust_bridge.request import NativeRequestContext, NativeResponsesWebSocketRequest
class _FakeNativeConnection:
@ -33,55 +27,55 @@ class _ClosedNativeConnection:
class _FakeNativeBridge:
contexts: list[NativeRequestContext] = []
@classmethod
async def connect(
cls,
request: NativeResponsesWebSocketRequest,
*,
options: NativeRequestOptions,
options: object,
context: NativeRequestContext,
) -> _FakeNativeConnection:
cls.contexts.append(context)
return _FakeNativeConnection()
@pytest.fixture(autouse=True)
def reset_responses_websocket():
responses_websocket.set_rust_responses_websocket(connection=None)
responses_websocket.set_rust_responses_websocket(connection=None, decline=None)
configuration.reset_rust_configuration()
responses_websocket.set_rust_responses_websocket(
decline=lambda model, custom_llm_provider, **features: (
"unsupported feature"
if any(features.get(key) for key in ("stream", "has_agentic_hook", "has_custom_client"))
or features.get("request_format") == "native"
else None
)
)
yield
responses_websocket.set_rust_responses_websocket(connection=None)
responses_websocket.set_rust_responses_websocket(connection=None, decline=None)
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())
adapter = responses_websocket._ConnectionAdapter(_ClosedNativeConnection())
with pytest.raises(responses_websocket.ConnectionClosedOK):
await adapter.recv()
@pytest.mark.asyncio
async def test_bridge_reports_unavailable(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None:
configuration.rust(True)
responses_websocket._RESPONSES_WEBSOCKET.override(None)
assert await responses_websocket.connect(
url="wss://example.test/responses",
headers={},
timeout=None,
) == NativeSkipped(NativeSkipReason.UNAVAILABLE)
assert (
await responses_websocket.connect(
url="wss://example.test/responses",
headers={},
timeout=None,
)
is None
)
@pytest.mark.asyncio
@ -97,13 +91,10 @@ async def test_enabled_bridge_connects_and_adapts_socket(
timeout=1.0,
)
assert isinstance(connection, Handled)
connection = connection.value
assert connection is not None
await connection.send("response.create")
assert await connection.recv() == "response.completed"
await connection.close()
assert _FakeNativeBridge.contexts[-1].capabilities.websocket_mode == "native"
assert _FakeNativeBridge.contexts[-1].capabilities.requires_connection is True
class _FailingNativeBridge:
@ -112,74 +103,96 @@ class _FailingNativeBridge:
cls,
request: NativeResponsesWebSocketRequest,
*,
options: NativeRequestOptions,
options: object,
context: NativeRequestContext,
) -> _FakeNativeConnection:
raise RuntimeError("connection failed")
@pytest.mark.asyncio
async def test_connection_failure_is_reported_to_orchestration() -> None:
configuration.rust(True)
responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge)
result = await responses_websocket.connect(url="wss://example.test/responses", headers={}, timeout=None)
assert isinstance(result, NativeFailed)
assert str(result.error) == "connection failed"
@pytest.mark.asyncio
async def test_managed_connection_closes_native_socket_on_consumer_failure() -> None:
configuration.rust(True)
socket = _FakeNativeConnection()
class Bridge:
@classmethod
async def connect(
cls,
request: NativeResponsesWebSocketRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> _FakeNativeConnection:
return socket
responses_websocket.set_rust_responses_websocket(connection=Bridge)
result = await responses_websocket.managed_connect(url="wss://example.test/responses", headers={}, timeout=1.0)
assert isinstance(result, Handled)
async def use_connection() -> None:
async with result.value as connection:
await connection.send("hello")
raise ValueError("consumer failed")
with pytest.raises(ValueError, match="consumer failed"):
await use_connection()
assert socket.sent == ["hello"]
assert socket.closed
@pytest.mark.asyncio
async def test_connection_failure_does_not_authorize_python_fallback() -> None:
from contextlib import AbstractAsyncContextManager
from litellm.rust_bridge.dispatch import anative_context, provider_errors
configuration.rust(True)
responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge)
@anative_context(
native=lambda: responses_websocket.managed_connect(
url="wss://example.test/responses", headers={}, timeout=None
),
route="responses_websocket",
errors=lambda: provider_errors("openai", "responses websocket"),
)
def execute() -> AbstractAsyncContextManager[object]:
pytest.fail("unknown native failures must not open a Python connection")
async def run() -> None:
async with execute():
pytest.fail("connection must fail before entering its body")
with pytest.raises(RuntimeError, match="connection failed"):
await responses_websocket.connect(url="wss://example.test/responses", headers={}, timeout=None)
@pytest.mark.parametrize("native", [False, True])
@pytest.mark.parametrize("session_error", [False, True])
@pytest.mark.asyncio
async def test_connection_dispatch_cleans_up_without_reconnecting(native, session_error):
from contextlib import asynccontextmanager
configuration.rust(True)
native_socket = _FakeNativeConnection()
python_socket = _FakeNativeConnection()
connections = []
class Native:
@classmethod
async def connect(cls, request, *, context, callback_adapter=None):
connections.append("native")
assert request.options.custom_llm_provider == "azure"
return native_socket
@asynccontextmanager
async def python():
connections.append("python")
try:
yield python_socket
finally:
await python_socket.close()
responses_websocket.set_rust_responses_websocket(connection=Native)
if not native:
responses_websocket.set_rust_responses_websocket(
decline=lambda model, custom_llm_provider, **features: "declined"
)
async def run():
async with responses_websocket.open_connection(
url="wss://example.test",
headers={},
timeout=1,
model="test-model",
provider="azure",
fallback=python,
):
if session_error:
raise RuntimeError("session failed")
if session_error:
with pytest.raises(RuntimeError, match="session failed"):
await run()
else:
await run()
assert connections == ["native" if native else "python"]
assert native_socket.closed == native
assert python_socket.closed == (not native)
@pytest.mark.asyncio
async def test_missing_acceptance_export_uses_python_connection_once():
from contextlib import asynccontextmanager
calls = []
socket = _FakeNativeConnection()
@asynccontextmanager
async def python():
calls.append("python")
try:
yield socket
finally:
await socket.close()
configuration.rust(True)
responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge)
responses_websocket._PREFLIGHT.override(None)
async with responses_websocket.open_connection(
url="wss://example.test", headers={}, timeout=1, model="model", provider="openai", fallback=python
) as connection:
assert connection is socket
assert calls == ["python"]
assert socket.closed

View file

@ -12,8 +12,12 @@ import pytest
import litellm
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge import chat_completions as bridge
from litellm.rust_bridge.request import NativeChatCompletionsRequest, NativeRequestContext, NativeRequestOptions
from litellm.rust_bridge.runtime import Handled, NativeFailed, NativeSkipped
from litellm.rust_bridge.request import (
NativeBedrockOptions,
NativeRequestCapabilities,
NativeRequestContext,
anthropic_options,
)
from litellm.types.utils import ModelResponse
RUST_RESPONSE = {
@ -96,40 +100,16 @@ class _RecordingCall:
self.result = result if result is not None else dict(RUST_RESPONSE)
self.error = error
self.calls: list[dict] = []
self.contexts: list[NativeRequestContext] = []
def __call__(
self,
request: NativeChatCompletionsRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
):
kwargs = {
"model": request.model,
"messages": request.messages,
"optional_params": request.optional_params,
"api_key": options.api_key,
"api_base": options.api_base,
"custom_llm_provider": options.custom_llm_provider,
"extra_headers": options.extra_headers,
"timeout_seconds": options.timeout_seconds,
}
self.calls.append(kwargs)
self.contexts.append(context)
def __call__(self, request, *, options, context):
self.calls.append({"request": request, "options": options, "context": context})
if self.error is not None:
raise self.error
return self.result
class _RecordingAsyncCall(_RecordingCall):
async def __call__(
self,
request: NativeChatCompletionsRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
):
async def __call__(self, request, *, options, context):
return _RecordingCall.__call__(self, request, options=options, context=context)
@ -282,8 +262,7 @@ class TestSyncCall:
result = bridge.chat_completions(**_call_kwargs(model_response))
assert isinstance(result, Handled)
result = result.value
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"
@ -297,28 +276,16 @@ class TestSyncCall:
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert native.calls[0]["timeout_seconds"] == 30.0
assert native.calls[0]["options"].timeout_seconds == 30.0
def test_preserves_execution_and_client_capabilities(self):
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
bridge.chat_completions(
**_call_kwargs(ModelResponse()),
stream=True,
has_custom_client=True,
)
assert native.contexts[0].capabilities.execution_mode == "sync"
assert native.contexts[0].capabilities.stream is True
assert native.contexts[0].capabilities.has_custom_client is True
def test_reports_unavailable_bridge(self, monkeypatch):
def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
assert isinstance(bridge.chat_completions(**_call_kwargs(ModelResponse())), NativeSkipped)
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
def test_reports_native_decline_to_orchestration(self, monkeypatch):
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 isinstance(bridge.chat_completions(**_call_kwargs(ModelResponse())), NativeFailed)
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
class TestAsyncCall:
@ -326,18 +293,286 @@ class TestAsyncCall:
async def test_builds_a_model_response(self):
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall())
result = await bridge.achat_completions(**_call_kwargs(ModelResponse()))
assert isinstance(result, Handled)
result = result.value
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_reports_unavailable_bridge(self, monkeypatch):
async def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
assert isinstance(await bridge.achat_completions(**_call_kwargs(ModelResponse())), NativeSkipped)
assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None
@pytest.mark.asyncio
async def test_reports_native_decline_to_orchestration(self, monkeypatch):
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 isinstance(await bridge.achat_completions(**_call_kwargs(ModelResponse())), NativeFailed)
assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None
class TestAsyncFallbackWrapper:
@pytest.mark.asyncio
async def test_returns_the_rust_response_without_running_the_fallback(self):
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall())
ran = []
async def fallback():
ran.append(True)
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result.choices[0].message.content == "hello from rust"
assert ran == []
@pytest.mark.asyncio
async def test_runs_the_fallback_when_the_core_declines(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")))
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
@pytest.mark.asyncio
async def test_runs_the_fallback_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
class TestFailureClassification:
"""A failure the provider already saw must not be retried on the Python
path: it would bill the customer for the same work twice."""
@pytest.fixture(autouse=True)
def _native_exceptions(self, monkeypatch):
_fake_native_bridge(monkeypatch)
def test_a_decline_falls_back_because_nothing_was_sent(self):
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming")))
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
def test_an_upstream_failure_is_surfaced_with_its_status(self):
from litellm.exceptions import RateLimitError
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited")))
with pytest.raises(RateLimitError) as raised:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert raised.value.status_code == 429
assert "rate limited" in str(raised.value)
def test_a_transport_failure_with_no_response_surfaces_as_a_500(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset")))
with pytest.raises(APIError) as raised:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert raised.value.status_code == 500
def test_an_unrecognized_error_is_not_swallowed(self):
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=RuntimeError("something else")))
with pytest.raises(RuntimeError):
bridge.chat_completions(**_call_kwargs(ModelResponse()))
@pytest.mark.asyncio
async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self):
from litellm.exceptions import InternalServerError
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom")))
ran = []
async def fallback():
ran.append(True)
return "python"
with pytest.raises(InternalServerError):
await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert ran == [], "a request the provider already served must not be re-issued"
@pytest.mark.asyncio
async def test_the_async_wrapper_falls_back_on_a_decline(self):
bridge.set_rust_chat_completions(
achat_completions=_RecordingAsyncCall(error=_FakeDeclined("blank message text"))
)
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
@pytest.mark.asyncio
async def test_missing_native_exception_types_does_not_authorize_python_fallback(monkeypatch):
_hide_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(
chat_completions=_RecordingCall(error=RuntimeError("connection failed")),
achat_completions=_RecordingAsyncCall(error=RuntimeError("connection failed")),
)
with pytest.raises(RuntimeError, match="connection failed"):
bridge.chat_completions(**_call_kwargs(ModelResponse()))
async def fallback():
pytest.fail("unknown failure must not retry through Python")
with pytest.raises(RuntimeError, match="connection failed"):
await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
def test_provider_credentials_are_separate_from_chat_body_params():
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
configuration.rust(True)
kwargs = _call_kwargs(ModelResponse())
kwargs["optional_params"] = {
"max_tokens": 32,
}
kwargs["bedrock"] = NativeBedrockOptions(
aws_access_key_id="test-access-key",
aws_secret_access_key="test-secret-key",
)
bridge.chat_completions(**kwargs)
request = native.calls[0]["request"]
options = native.calls[0]["options"]
assert request.optional_params == {"max_tokens": 32}
assert options.bedrock.aws_access_key_id == "test-access-key"
assert options.bedrock.aws_secret_access_key == "test-secret-key"
def test_provider_payload_extensions_cross_the_boundary_without_partitioning():
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
configuration.rust(True)
extensions = {
"vendor_object": {"nested": None},
"vendor_array": [1, "two", False],
"vendor_scalar": 0.25,
"extra_body": {"temperature": 0.2, "config": {"replacement": True}},
}
kwargs = _call_kwargs(ModelResponse())
kwargs["optional_params"] = extensions
bridge.chat_completions(**kwargs)
assert native.calls[0]["request"].optional_params == extensions
def test_typed_capability_and_provider_metadata_facts_are_isolated():
context = NativeRequestContext(
capabilities=NativeRequestCapabilities(
stream=True,
has_agentic_hook=True,
has_custom_client=True,
request_format="native",
)
)
anthropic = anthropic_options({"metadata": {"user_id": "user-123", "ignored": object()}})
assert context.capabilities.request_format == "native"
assert context.capabilities.has_agentic_hook is True
assert anthropic.user_id == "user-123"
@pytest.mark.parametrize("provider", ["anthropic", "bedrock", "openai"])
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_public_completion_discovers_any_provider(provider, asynchronous):
native = _RecordingCall()
anative = _RecordingAsyncCall()
gate = _RecordingDecline()
bridge.set_rust_chat_completions(chat_completions=native, achat_completions=anative, decline=gate)
kwargs = {
"model": f"{provider}/test-model",
"messages": MESSAGES,
"api_key": "key",
"max_tokens": 16,
"num_retries": 0,
}
response = await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs)
assert response.choices[0].message.content == "hello from rust"
calls = anative.calls if asynchronous else native.calls
assert len(calls) == 1
assert calls[0]["options"].custom_llm_provider == provider
assert calls[0]["request"].messages == MESSAGES
assert gate.calls[0]["custom_llm_provider"] == provider
assert len(native.calls) + len(anative.calls) == 1
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize(
"failure", ["preflight", "missing_preflight", "decline", "unavailable", "error", "malformed", "cancelled"]
)
@pytest.mark.asyncio
async def test_public_completion_fallback_contract(monkeypatch, asynchronous, failure):
import asyncio
import importlib
from litellm.integrations.custom_logger import CustomLogger
class Recorder(CustomLogger):
def __init__(self):
self.pre = 0
self.post = 0
def log_pre_api_call(self, model, messages, kwargs):
self.pre += 1
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
self.post += 1
recorder = Recorder()
monkeypatch.setattr(litellm, "input_callback", [recorder])
_fake_native_bridge(monkeypatch)
python_calls = []
def python(ctx):
python_calls.append(ctx)
def finish():
ctx.logging.pre_call(input=MESSAGES, api_key="key", additional_args={})
ctx.logging.post_call(input=MESSAGES, api_key="key", original_response="python")
return ModelResponse(choices=[{"message": {"role": "assistant", "content": "python"}}])
async def afinish():
return finish()
return afinish() if ctx.acompletion else finish()
monkeypatch.setattr(importlib.import_module("litellm.main"), "_complete_python", python)
native_error = {
"decline": _FakeDeclined("request unsupported"),
"error": RuntimeError("execution failed"),
"cancelled": asyncio.CancelledError(),
}.get(failure)
native = _RecordingCall(result={} if failure == "malformed" else None, error=native_error)
anative = _RecordingAsyncCall(result={} if failure == "malformed" else None, error=native_error)
bridge.set_rust_chat_completions(
chat_completions=native,
achat_completions=anative,
decline=_RecordingDecline("unsupported" if failure == "preflight" else None),
)
if failure == "missing_preflight":
bridge._CHAT_PREFLIGHT.override(None)
if failure == "unavailable":
bridge._CHAT.sync.override(None)
bridge._CHAT.asynchronous.override(None)
async def run():
kwargs = {"model": "openai/test-model", "messages": MESSAGES, "api_key": "key", "num_retries": 0}
return await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs)
if failure in {"error", "malformed", "cancelled"}:
with pytest.raises(asyncio.CancelledError if failure == "cancelled" else Exception):
await run()
assert python_calls == []
else:
result = await run()
assert result.choices[0].message.content == "python"
assert len(python_calls) == 1
assert recorder.pre == 1
assert recorder.post == 1