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

refactor(native): separate request options and context

Browse source
This commit is contained in:
Yujong Lee 2026-09-05 16:59:57 -07:00 • committed by yujonglee
parent f8d6ac7a76
commit 221ddb3ab6
56 changed files with 1508 additions and 1335 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm-rust/README.md
View file

@ -27,15 +27,15 @@ coverage and production evidence.
## Native request boundary
Native HTTP routes and Responses WebSocket connections accept `native(request, *, context)`
The request carries the endpoint payload and `NativeRequestOptions`: credentials,
provider routing, headers, query parameters, and timeout. `NativeRequestContext`
carries LiteLLM metadata, call identity, and attribution separately from the provider payload
Native HTTP routes and Responses WebSocket connections accept
`native(request, *, options, context)`. The request carries only endpoint payload.
`NativeRequestOptions` carries credentials, typed provider configuration, routing,
headers, query parameters, and timeout. `NativeRequestContext` carries call identity,
attribution, and typed capability facts separately from the provider payload.
Python builds the frozen request dataclasses in `litellm/rust_bridge/request.py` and
PyO3 extracts their fields before execution. Provider connection parameters, such as
AWS credentials and Vertex project/location, belong in `options.provider_connection`
rather than the request body
PyO3 extracts their fields before execution. AWS credentials and metadata policy
belong in `options.bedrock`; Vertex project/location belongs in `options.vertex`.
This boundary preserves existing Python provider preparation, preflight decisions,
fallback, and callbacks

28
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
View file

@ -94,29 +94,29 @@ impl AudioTranscriptionLifecycleHooks {
custom_llm_provider,
audio,
api_key,
provider_connection,
bedrock,
api_base,
extra_headers,
optional_params,
timeout,
..
} = request;
let provider_request =
prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest {
let provider_request = prepare_audio_transcription_provider_call(
CoreAudioTranscriptionRequest {
model: &model,
audio,
optional_params,
options: RequestOptions {
provider_connection,
api_key: (api_key.as_deref()).map(|value| value.to_string()),
api_base: (api_base.as_deref()).map(|value| value.to_string()),
custom_llm_provider: (Some(&custom_llm_provider))
.map(|value| value.to_string()),
extra_headers,
timeout,
..Default::default()
},
})?;
},
RequestOptions {
bedrock: Some(bedrock),
api_key: (api_key.as_deref()).map(|value| value.to_string()),
api_base: (api_base.as_deref()).map(|value| value.to_string()),
custom_llm_provider: (Some(&custom_llm_provider)).map(|value| value.to_string()),
extra_headers,
timeout,
..Default::default()
},
)?;
self.run_during_call_guardrails(provider_request).await
}

4
litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs
View file

@ -3,6 +3,7 @@ use litellm_core::Error;
use litellm_core::audio_transcription::execute_audio_transcription_provider_call;
use litellm_core::call_lifecycle::CallLifecycle;
use litellm_core::request_context::LiteLlmRequestContext;
use litellm_core::request_options::RequestOptions;
use serde_json::Value;
mod hooks;
@ -15,6 +16,7 @@ use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
pub async fn audio_transcription(
request: AudioTranscriptionRequest<'_>,
options: &RequestOptions,
context: &LiteLlmRequestContext,
hooks: RequestHooks,
) -> Result<Value, Error> {
@ -22,7 +24,7 @@ pub async fn audio_transcription(
request,
context,
hooks,
} = prepare_audio_transcription_call(request, context, hooks);
} = prepare_audio_transcription_call(request, options.clone(), context, hooks);
CallLifecycle::default()
.run(
context,

27
litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs
View file

@ -1,6 +1,7 @@
use crate::integrations::types::RequestHooks;
use litellm_core::call_lifecycle::CallLifecycleContext;
use litellm_core::request_context::LiteLlmRequestContext;
use litellm_core::request_options::RequestOptions;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
@ -19,6 +20,7 @@ pub(crate) struct PreparedAudioTranscriptionCall {
pub(crate) fn prepare_audio_transcription_call(
request: AudioTranscriptionRequest<'_>,
options: RequestOptions,
context: &LiteLlmRequestContext,
hooks: RequestHooks,
) -> PreparedAudioTranscriptionCall {
@ -26,14 +28,13 @@ pub(crate) fn prepare_audio_transcription_call(
.litellm_call_id
.clone()
.unwrap_or_else(new_audio_transcription_call_id);
let provider_info = get_custom_llm_provider(
request.model,
request.options.custom_llm_provider.as_deref(),
)
.unwrap_or(CustomLlmProvider {
model: request.model,
custom_llm_provider: "bedrock",
});
let provider_info =
get_custom_llm_provider(request.model, options.custom_llm_provider.as_deref()).unwrap_or(
CustomLlmProvider {
model: request.model,
custom_llm_provider: "bedrock",
},
);
PreparedAudioTranscriptionCall {
context: CallLifecycleContext::new(
"audio_transcription",
@ -45,12 +46,12 @@ pub(crate) fn prepare_audio_transcription_call(
model: provider_info.model.to_string(),
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
audio: request.audio,
provider_connection: request.options.provider_connection,
api_key: request.options.api_key,
api_base: request.options.api_base,
extra_headers: request.options.extra_headers,
bedrock: options.bedrock.unwrap_or_default(),
api_key: options.api_key,
api_base: options.api_base,
extra_headers: options.extra_headers,
optional_params: request.optional_params,
timeout: request.options.timeout,
timeout: options.timeout,
},
hooks: AudioTranscriptionLifecycleHooks::new(
CustomLoggerRunner::new(hooks.callbacks),

32
litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs
View file

@ -1,6 +1,6 @@
use crate::integrations::types::RequestHooks;
use litellm_core::request_context::LiteLlmRequestContext;
use litellm_core::request_options::RequestOptions;
use litellm_core::request_options::{BedrockOptions, RequestOptions};
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
@ -29,27 +29,27 @@ async fn bedrock_request_is_signed_and_contains_audio() {
stream.write_all(response).expect("response");
});
let provider_connection = Map::from_iter([
("aws_access_key_id".to_string(), json!("access-key")),
("aws_secret_access_key".to_string(), json!("secret-key")),
("aws_region_name".to_string(), json!("us-east-1")),
]);
let bedrock = BedrockOptions {
aws_access_key_id: Some("access-key".to_string()),
aws_secret_access_key: Some("secret-key".to_string()),
aws_region_name: Some("us-east-1".to_string()),
..Default::default()
};
let api_base = format!("http://{address}");
let response = audio_transcription(
AudioTranscriptionRequest {
model: "mistral.voxtral-mini-3b-2507",
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
optional_params: Map::new(),
options: RequestOptions {
provider_connection,
api_key: None,
api_base: (Some(&api_base)).map(|value| value.to_string()),
custom_llm_provider: (Some("bedrock")).map(|value| value.to_string()),
extra_headers: None,
timeout: None,
..Default::default()
},
},
&RequestOptions {
bedrock: Some(bedrock),
api_key: None,
api_base: (Some(&api_base)).map(|value| value.to_string()),
custom_llm_provider: (Some("bedrock")).map(|value| value.to_string()),
extra_headers: None,
timeout: None,
..Default::default()
},
&LiteLlmRequestContext {
attribution: Default::default(),

5
litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs
View file

@ -1,4 +1,4 @@
use litellm_core::request_options::RequestOptions;
use litellm_core::request_options::BedrockOptions;
use std::time::Duration;
use serde_json::{Map, Value};
@ -7,14 +7,13 @@ pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
pub audio: Value,
pub optional_params: Map<String, Value>,
pub options: RequestOptions,
}
pub(crate) struct PreparedAudioTranscriptionRequest {
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) audio: Value,
pub(crate) provider_connection: Map<String, Value>,
pub(crate) bedrock: BedrockOptions,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) extra_headers: Option<Map<String, Value>>,

6
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -7,6 +7,7 @@ use litellm_core::Error;
use litellm_core::http_utils::string_headers;
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
use litellm_core::request_context::LiteLlmRequestContext;
use litellm_core::request_options::RequestOptions;
use litellm_core::responses::types::{ResponsesWebSocketRequest, ResponsesWsEvent};
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
use tokio::net::TcpStream;
@ -36,9 +37,10 @@ pub struct ResponsesWebSocketConnection {
impl ResponsesWebSocketConnection {
pub async fn connect(
input: ResponsesWebSocketRequest,
options: &RequestOptions,
_context: &LiteLlmRequestContext,
) -> Result<Self, Error> {
let headers = string_headers("Responses WebSocket", input.options.extra_headers)?;
let headers = string_headers("Responses WebSocket", options.extra_headers.clone())?;
let mut request = input
.url
.as_str()
@ -53,7 +55,7 @@ impl ResponsesWebSocketConnection {
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_async(request);
let result = match input.options.timeout {
let result = match options.timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Network("Responses WebSocket connection timed out".to_string())
})?,

2
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -88,7 +88,7 @@ impl OcrLifecycleHooks {
.optional_params
.clone()
.into_iter()
.chain(request.provider_connection)
.chain(request.vertex.into_map())
.collect();
let url = config.complete_url(
request.api_base.as_deref(),

36
litellm-rust/crates/ai-gateway/src/ocr/mod.rs
View file

@ -2,6 +2,7 @@ use crate::integrations::types::RequestHooks;
use litellm_core::Error;
use litellm_core::call_lifecycle::CallLifecycle;
use litellm_core::request_context::LiteLlmRequestContext;
use litellm_core::request_options::RequestOptions;
use serde_json::Value;
mod common_utils;
@ -18,6 +19,7 @@ use prepare::{PreparedOcrCall, prepare_ocr_call};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn ocr(
request: OcrRequest<'_>,
options: &RequestOptions,
context: &LiteLlmRequestContext,
hooks: RequestHooks,
) -> Result<Value, Error> {
@ -25,7 +27,7 @@ pub async fn ocr(
request,
context,
hooks,
} = prepare_ocr_call(request, context, hooks);
} = prepare_ocr_call(request, options.clone(), context, hooks);
CallLifecycle::default()
.run(context, request, &hooks, |request| {
execute_ocr_provider_call(request, &hooks)
@ -77,16 +79,17 @@ mod tests {
String::from_utf8(request).expect("request is utf8")
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
options: RequestOptions {
fn base_ocr_request(model: &str) -> (OcrRequest<'_>, RequestOptions) {
(
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
},
RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: None,
custom_llm_provider: None,
@ -94,7 +97,7 @@ mod tests {
timeout: None,
..Default::default()
},
}
)
}
#[tokio::test]
@ -132,10 +135,10 @@ mod tests {
(upload_request, parse_request)
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.options.api_base = Some(&api_base).map(|value| value.to_string());
request.options.api_key = None;
request.options.extra_headers = Some(Map::from_iter([
let (mut request, mut options) = base_ocr_request("reducto/parse-v3");
options.api_base = Some(&api_base).map(|value| value.to_string());
options.api_key = None;
options.extra_headers = Some(Map::from_iter([
("Authorization".to_string(), json!("Bearer test-key")),
("x-trace-id".to_string(), json!("trace-1")),
]));
@ -154,6 +157,7 @@ mod tests {
let response = ocr(
request,
&options,
&LiteLlmRequestContext {
..Default::default()
},

38
litellm-rust/crates/ai-gateway/src/ocr/prepare.rs
View file

@ -1,6 +1,7 @@
use crate::integrations::types::RequestHooks;
use litellm_core::call_lifecycle::CallLifecycleContext;
use litellm_core::request_context::LiteLlmRequestContext;
use litellm_core::request_options::RequestOptions;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
@ -22,6 +23,7 @@ pub(crate) struct PreparedOcrCall {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn prepare_ocr_call(
request: OcrRequest<'_>,
options: RequestOptions,
context: &LiteLlmRequestContext,
hooks: RequestHooks,
) -> PreparedOcrCall {
@ -29,14 +31,13 @@ pub(crate) fn prepare_ocr_call(
.litellm_call_id
.clone()
.unwrap_or_else(new_ocr_call_id);
let provider_info = get_custom_llm_provider(
request.model,
request.options.custom_llm_provider.as_deref(),
)
.unwrap_or(CustomLlmProvider {
model: request.model,
custom_llm_provider: "mistral",
});
let provider_info =
get_custom_llm_provider(request.model, options.custom_llm_provider.as_deref()).unwrap_or(
CustomLlmProvider {
model: request.model,
custom_llm_provider: "mistral",
},
);
let model = provider_info.model.to_string();
let custom_llm_provider = provider_info.custom_llm_provider.to_string();
let config = ocr_provider_config(&custom_llm_provider, &model)
@ -72,12 +73,12 @@ pub(crate) fn prepare_ocr_call(
model,
custom_llm_provider,
document: request.document,
provider_connection: request.options.provider_connection,
api_key: request.options.api_key,
api_base: request.options.api_base,
extra_headers: request.options.extra_headers,
vertex: options.vertex.unwrap_or_default(),
api_key: options.api_key,
api_base: options.api_base,
extra_headers: options.extra_headers,
optional_params,
timeout: request.options.timeout,
timeout: options.timeout,
},
hooks: OcrLifecycleHooks::new(
CustomLoggerRunner::new(hooks.callbacks),
@ -135,15 +136,6 @@ mod tests {
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
options: RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
timeout: None,
..Default::default()
},
}
}
@ -157,6 +149,7 @@ mod tests {
fn native_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(
request_with_format("native"),
RequestOptions::default(),
&LiteLlmRequestContext {
..Default::default()
},
@ -173,6 +166,7 @@ mod tests {
fn unknown_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(
request_with_format("raw"),
RequestOptions::default(),
&LiteLlmRequestContext {
..Default::default()
},

5
litellm-rust/crates/ai-gateway/src/ocr/types.rs
View file

@ -1,14 +1,13 @@
use litellm_core::request_options::RequestOptions;
use std::time::Duration;
use litellm_core::ocr::transformation::OcrProviderConfig;
use litellm_core::request_options::VertexOptions;
use serde_json::{Map, Value};
pub struct OcrRequest<'a> {
pub model: &'a str,
pub document: Value,
pub optional_params: Map<String, Value>,
pub options: RequestOptions,
}
pub(crate) struct PreparedOcrRequest {
@ -16,7 +15,7 @@ pub(crate) struct PreparedOcrRequest {
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) document: Value,
pub(crate) provider_connection: Map<String, Value>,
pub(crate) vertex: VertexOptions,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) extra_headers: Option<Map<String, Value>>,

124
litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -216,24 +216,21 @@ impl CustomGuardrail for RecordingOcrGuardrail {
}
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
options: RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
timeout: None,
fn base_ocr_request(model: &str) -> (OcrRequest<'_>, RequestOptions) {
(
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
},
RequestOptions {
api_key: Some("sk-test".to_string()),
..Default::default()
},
}
)
}
#[tokio::test]
@ -244,8 +241,8 @@ async fn reducto_during_call_guardrail_blocks_before_upload() {
let address = listener.local_addr().expect("listener has local address");
let api_base = format!("http://{address}");
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call());
let mut request = base_ocr_request("reducto/parse-v3");
request.options.api_base = Some(&api_base).map(|value| value.to_string());
let (mut request, mut options) = base_ocr_request("reducto/parse-v3");
options.api_base = Some(&api_base).map(|value| value.to_string());
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
@ -257,6 +254,7 @@ async fn reducto_during_call_guardrail_blocks_before_upload() {
let error = ocr(
request,
&options,
&LiteLlmRequestContext {
..Default::default()
},
@ -292,8 +290,8 @@ async fn reducto_upload_error_body_is_truncated() {
.expect("writes upload response");
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.options.api_base = Some(&api_base).map(|value| value.to_string());
let (mut request, mut options) = base_ocr_request("reducto/parse-v3");
options.api_base = Some(&api_base).map(|value| value.to_string());
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
@ -301,6 +299,7 @@ async fn reducto_upload_error_body_is_truncated() {
let error = ocr(
request,
&options,
&LiteLlmRequestContext {
..Default::default()
},
@ -351,15 +350,14 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
options: RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("mistral")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("mistral")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
attribution: RequestAttribution {
@ -430,15 +428,14 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
options: RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("mistral")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("mistral")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
attribution: RequestAttribution::default(),
@ -485,15 +482,14 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
options: RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("mistral")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_millis(100)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("mistral")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_millis(100)),
..Default::default()
},
&LiteLlmRequestContext {
attribution: RequestAttribution::default(),
@ -566,15 +562,14 @@ async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
options: RequestOptions {
api_key: (Some("sk-for-rust-fallback")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("mistral")).map(|value| value.to_string()),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk-for-rust-fallback")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("mistral")).map(|value| value.to_string()),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
attribution: RequestAttribution::default(),
@ -646,15 +641,14 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() {
"document_url": "https://example.com/doc.pdf"
}),
optional_params: Map::new(),
options: RequestOptions {
api_key: (Some("di-key")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("di-key")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
attribution: RequestAttribution::default(),

2
litellm-rust/crates/core/src/audio_transcription/handler.rs
View file

@ -59,7 +59,7 @@ async fn signed_headers(
};
let env_lookup = |key: &str| std::env::var(key).ok();
let credentials = resolve_credentials(
aws_auth_config(&request.provider_connection, &env_lookup),
aws_auth_config(&request.bedrock.into_map(), &env_lookup),
&env_lookup,
)
.await?;

9
litellm-rust/crates/core/src/audio_transcription/mod.rs
View file

@ -1,5 +1,6 @@
use crate::Error;
use crate::request_context::LiteLlmRequestContext;
use crate::request_options::RequestOptions;
mod client;
mod handler;
mod prepare;
@ -15,10 +16,14 @@ pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn audio_transcription(
request: AudioTranscriptionRequest<'_>,
options: &RequestOptions,
_context: &LiteLlmRequestContext,
) -> Result<Value, Error> {
execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?)
.await
execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(
request,
options.clone(),
)?)
.await
}
#[cfg(test)]

52
litellm-rust/crates/core/src/audio_transcription/prepare.rs
View file

@ -2,6 +2,7 @@ use crate::error::Error;
use crate::http_utils::{has_header, string_headers};
#[cfg(feature = "bedrock-auth")]
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use crate::request_options::RequestOptions;
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
@ -20,35 +21,36 @@ fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProv
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn prepare_audio_transcription_provider_call(
request: AudioTranscriptionRequest<'_>,
options: RequestOptions,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let provider_info = get_custom_llm_provider(
request.model,
request.options.custom_llm_provider.as_deref(),
)
.or_else(|| {
request
.options
.custom_llm_provider
.as_deref()
.map(|provider| CustomLlmProvider {
model: request.model,
custom_llm_provider: provider,
let provider_info =
get_custom_llm_provider(request.model, options.custom_llm_provider.as_deref())
.or_else(|| {
options
.custom_llm_provider
.as_deref()
.map(|provider| CustomLlmProvider {
model: request.model,
custom_llm_provider: provider,
})
})
})
.ok_or_else(|| {
Error::InvalidProvider(
"unable to resolve custom_llm_provider for audio transcription request".to_string(),
)
})?;
.ok_or_else(|| {
Error::InvalidProvider(
"unable to resolve custom_llm_provider for audio transcription request"
.to_string(),
)
})?;
let model = provider_info.model.to_string();
let config = provider_config(provider_info.custom_llm_provider)
.ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers("audio transcription", request.options.extra_headers)?;
let auth = config.auth_strategy(&model, &request.options.provider_connection, &env_lookup)?;
let mut headers = string_headers("audio transcription", options.extra_headers)?;
let bedrock = options.bedrock.unwrap_or_default();
let bedrock_options = bedrock.clone().into_map();
let auth = config.auth_strategy(&model, &bedrock_options, &env_lookup)?;
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(&headers, "authorization")
&& let Some(api_key) = request.options.api_key.as_deref()
&& let Some(api_key) = options.api_key.as_deref()
{
headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
}
@ -56,9 +58,9 @@ pub fn prepare_audio_transcription_provider_call(
headers.push(("Content-Type".to_string(), "application/json".to_string()));
}
let url = config.complete_url(
request.options.api_base.as_deref(),
options.api_base.as_deref(),
&model,
&request.options.provider_connection,
&bedrock_options,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
@ -73,7 +75,7 @@ pub fn prepare_audio_transcription_provider_call(
upstream_headers: headers,
auth,
#[cfg(feature = "bedrock-auth")]
provider_connection: request.options.provider_connection,
timeout: request.options.timeout,
bedrock,
timeout: options.timeout,
})
}

31
litellm-rust/crates/core/src/audio_transcription/tests.rs
View file

@ -1,5 +1,5 @@
use crate::request_context::LiteLlmRequestContext;
use crate::request_options::RequestOptions;
use crate::request_options::{BedrockOptions, RequestOptions};
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
@ -29,26 +29,27 @@ async fn bedrock_request_is_signed_and_contains_audio() {
stream.write_all(response).expect("response");
});
let provider_connection = Map::from_iter([
("aws_access_key_id".to_string(), json!("access-key")),
("aws_secret_access_key".to_string(), json!("secret-key")),
("aws_region_name".to_string(), json!("us-east-1")),
]);
let bedrock = BedrockOptions {
aws_access_key_id: Some("access-key".to_string()),
aws_secret_access_key: Some("secret-key".to_string()),
aws_region_name: Some("us-east-1".to_string()),
..Default::default()
};
let api_base = format!("http://{address}");
let response = audio_transcription(
AudioTranscriptionRequest {
model: "mistral.voxtral-mini-3b-2507",
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
optional_params: Map::new(),
options: RequestOptions {
provider_connection,
api_key: None,
api_base: (Some(&api_base)).map(|value| value.to_string()),
custom_llm_provider: (Some("bedrock")).map(|value| value.to_string()),
extra_headers: None,
timeout: None,
..Default::default()
},
},
&RequestOptions {
bedrock: Some(bedrock),
api_key: None,
api_base: (Some(&api_base)).map(|value| value.to_string()),
custom_llm_provider: (Some("bedrock")).map(|value| value.to_string()),
extra_headers: None,
timeout: None,
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()

6
litellm-rust/crates/core/src/audio_transcription/types.rs
View file

@ -1,16 +1,16 @@
use crate::request_options::RequestOptions;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::request_options::BedrockOptions;
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
pub audio: Value,
pub optional_params: Map<String, Value>,
pub options: RequestOptions,
}
#[derive(Clone)]
@ -23,7 +23,7 @@ pub struct ProviderAudioTranscriptionRequest {
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: AudioTranscriptionAuth,
#[cfg(feature = "bedrock-auth")]
pub(super) provider_connection: Map<String, Value>,
pub(super) bedrock: BedrockOptions,
pub(super) timeout: Option<Duration>,
}

11
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -101,15 +101,10 @@ pub(super) async fn signed_headers(
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
// A host with its own resolution chain hands the result down; only fall
// back to deriving credentials here when it supplied none.
let credentials = match host_supplied_credentials(&request.provider_connection) {
let bedrock = request.bedrock.into_map();
let credentials = match host_supplied_credentials(&bedrock) {
Some(credentials) => credentials,
None => {
resolve_credentials(
aws_auth_config(&request.provider_connection, &env_lookup),
&env_lookup,
)
.await?
}
None => resolve_credentials(aws_auth_config(&bedrock, &env_lookup), &env_lookup).await?,
};
let signature = sign_bedrock_post(
&request.url,

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

@ -8,6 +8,7 @@
use crate::Error;
use crate::request_context::LiteLlmRequestContext;
use crate::request_options::RequestOptions;
mod client;
mod common_utils;
pub mod conversation;
@ -26,9 +27,11 @@ use types::{ChatCompletionsRequest, ChatCompletionsResponse};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn chat_completions(
request: ChatCompletionsRequest<'_>,
options: &RequestOptions,
context: &LiteLlmRequestContext,
) -> Result<ChatCompletionsResponse, Error> {
execute_chat_completions_provider_call(resolve_request(request, context)?).await
execute_chat_completions_provider_call(resolve_request(request, options.clone(), context)?)
.await
}
/// Whether the core would accept this request, without resolving credentials or

26
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -1,4 +1,5 @@
use crate::request_context::LiteLlmRequestContext;
use crate::request_options::RequestOptions;
use serde_json::Value;
use crate::error::Error;
@ -40,12 +41,11 @@ pub(super) fn parse_messages(messages: Value) -> Result<Vec<ChatMessage>, Error>
pub(super) fn resolve_request(
request: ChatCompletionsRequest<'_>,
options: RequestOptions,
_context: &LiteLlmRequestContext,
) -> Result<ResolvedChatCompletionsRequest, Error> {
let (model, config) = resolve_provider_config(
request.model,
request.options.custom_llm_provider.as_deref(),
)
let (model, config) =
resolve_provider_config(request.model, options.custom_llm_provider.as_deref())
.map_err(|_| Error::Declined("provider is not on the rust chat completions path"))?;
let messages =
parse_messages(request.messages).map_err(|_| Error::Declined("unreadable message list"))?;
@ -60,7 +60,7 @@ pub(super) fn resolve_request(
config,
messages,
optional_params: request.optional_params,
options: request.options,
options,
})
}
@ -75,7 +75,12 @@ fn validate_environment(
let auth = config.auth(
request.options.api_key.as_deref(),
model,
&request.options.provider_connection,
&request
.options
.bedrock
.clone()
.unwrap_or_default()
.into_map(),
&env_lookup,
)?;
match &auth {
@ -128,7 +133,12 @@ pub(super) fn prepare_provider_request(
let url = config.complete_url(
request.options.api_base.as_deref(),
&model,
&request.options.provider_connection,
&request
.options
.bedrock
.clone()
.unwrap_or_default()
.into_map(),
&env_lookup,
)?;
let transformed =
@ -141,7 +151,7 @@ pub(super) fn prepare_provider_request(
body: transformed.body,
upstream_headers: headers,
auth,
provider_connection: request.options.provider_connection,
bedrock: request.options.bedrock.unwrap_or_default(),
timeout: request.options.timeout,
})
}

91
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -1,5 +1,5 @@
use crate::request_context::LiteLlmRequestContext;
use crate::request_options::RequestOptions;
use crate::request_options::{BedrockOptions, RequestOptions};
use serde_json::{Map, Value, json};
use crate::error::Error;
@ -8,11 +8,17 @@ use super::prepare::{prepare_provider_request, resolve_request};
use super::transformation::ChatCompletionsAuth;
use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest};
struct TestChatCompletionsCall<'a> {
request: ChatCompletionsRequest<'a>,
options: RequestOptions,
}
fn prepare_chat_completions_call(
request: ChatCompletionsRequest<'_>,
call: TestChatCompletionsCall<'_>,
) -> Result<ProviderChatCompletionsRequest, Error> {
prepare_provider_request(resolve_request(
request,
call.request,
call.options,
&LiteLlmRequestContext {
..Default::default()
},
@ -24,15 +30,16 @@ fn request<'a>(
provider: Option<&'a str>,
messages: Value,
optional_params: Value,
) -> ChatCompletionsRequest<'a> {
ChatCompletionsRequest {
model,
messages,
optional_params: match optional_params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
) -> TestChatCompletionsCall<'a> {
TestChatCompletionsCall {
request: ChatCompletionsRequest {
model,
messages,
optional_params: match optional_params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
},
},
options: RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: None,
@ -46,7 +53,7 @@ fn request<'a>(
/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers
/// carry resolved credentials), so unwrap the failure case by hand.
fn decline(request: ChatCompletionsRequest<'_>) -> Error {
fn decline(request: TestChatCompletionsCall<'_>) -> Error {
match prepare_chat_completions_call(request) {
Err(error) => error,
Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url),
@ -325,13 +332,11 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
);
call.options.provider_connection = Map::from_iter([
("aws_access_key_id".to_string(), json!("AKIDEXAMPLE")),
(
"aws_secret_access_key".to_string(),
json!("wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"),
),
]);
call.options.bedrock = Some(BedrockOptions {
aws_access_key_id: Some("AKIDEXAMPLE".to_string()),
aws_secret_access_key: Some("wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY".to_string()),
..Default::default()
});
// A key would resolve to a bearer token and never reach the signer.
call.options.api_key = None;
call.options.extra_headers = Some(Map::from_iter([(
@ -383,13 +388,11 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() {
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
);
call.options.provider_connection = Map::from_iter([
("aws_access_key_id".to_string(), json!("AKIDEXAMPLE")),
(
"aws_secret_access_key".to_string(),
json!("wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"),
),
]);
call.options.bedrock = Some(BedrockOptions {
aws_access_key_id: Some("AKIDEXAMPLE".to_string()),
aws_secret_access_key: Some("wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY".to_string()),
..Default::default()
});
call.options.api_key = None;
call.options.extra_headers =
Some(Map::from_iter([(forwarded.to_string(), json!("forged"))]));
@ -613,7 +616,7 @@ mod round_trip {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use crate::chat_completions::chat_completions;
use crate::chat_completions::chat_completions as run_chat_completions;
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
@ -676,15 +679,16 @@ mod round_trip {
(format!("http://127.0.0.1:{port}/v1/messages"), handle)
}
fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> {
ChatCompletionsRequest {
model: "anthropic/claude-sonnet-4-5",
messages,
optional_params: match params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
fn call(api_base: &str, messages: Value, params: Value) -> TestChatCompletionsCall<'_> {
TestChatCompletionsCall {
request: ChatCompletionsRequest {
model: "anthropic/claude-sonnet-4-5",
messages,
optional_params: match params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
},
},
options: RequestOptions {
api_key: (Some("sk-test")).map(|value| value.to_string()),
api_base: (Some(api_base)).map(|value| value.to_string()),
@ -696,12 +700,19 @@ mod round_trip {
}
}
async fn execute(
call: TestChatCompletionsCall<'_>,
context: &LiteLlmRequestContext,
) -> Result<super::super::types::ChatCompletionsResponse, Error> {
run_chat_completions(call.request, &call.options, context).await
}
const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
#[tokio::test]
async fn round_trip_sends_the_translated_body_and_normalizes_the_response() {
let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await;
let response = chat_completions(
let response = execute(
call(
&api_base,
json!([
@ -751,7 +762,7 @@ mod round_trip {
const NO_USAGE: &str =
r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#;
let (api_base, handle) = serve_once("200 OK", NO_USAGE).await;
let err = chat_completions(
let err = execute(
call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
@ -774,7 +785,7 @@ mod round_trip {
async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() {
const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#;
let (api_base, handle) = serve_once("200 OK", TOOL_USE).await;
let err = chat_completions(
let err = execute(
call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
@ -797,7 +808,7 @@ mod round_trip {
async fn an_upstream_error_status_keeps_its_code() {
let (api_base, handle) =
serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await;
let err = chat_completions(
let err = execute(
call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
@ -823,7 +834,7 @@ mod round_trip {
listener.local_addr().expect("has an address").port()
// Dropped here, so the port is closed and the connect is refused.
};
let err = chat_completions(
let err = execute(
call(
&format!("http://127.0.0.1:{port}/v1/messages"),
json!([{"role": "user", "content": "hi"}]),

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

@ -1,4 +1,4 @@
use crate::request_options::RequestOptions;
use crate::request_options::{BedrockOptions, RequestOptions};
use std::time::Duration;
use serde::{Deserialize, Serialize};
@ -16,7 +16,6 @@ pub struct ChatCompletionsRequest<'a> {
pub model: &'a str,
pub messages: Value,
pub optional_params: Map<String, Value>,
pub options: RequestOptions,
}
pub(super) struct ResolvedChatCompletionsRequest {
@ -35,7 +34,7 @@ pub(super) struct ProviderChatCompletionsRequest {
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: ChatCompletionsAuth,
#[cfg_attr(not(feature = "bedrock-auth"), allow(dead_code))]
pub(super) provider_connection: Map<String, Value>,
pub(super) bedrock: BedrockOptions,
pub(super) timeout: Option<Duration>,
}

6
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -10,8 +10,9 @@ use super::types::{AnthropicMessagesResponse, MessagesRequest};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) async fn execute_messages_provider_call(
request: MessagesRequest<'_>,
options: crate::request_options::RequestOptions,
) -> Result<AnthropicMessagesResponse, Error> {
let request = prepare_provider_request(request)?;
let request = prepare_provider_request(request, options)?;
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
@ -44,8 +45,9 @@ pub(super) async fn execute_messages_provider_call(
pub(super) async fn execute_messages_provider_stream(
request: MessagesRequest<'_>,
options: crate::request_options::RequestOptions,
) -> Result<reqwest::Response, Error> {
let request = prepare_provider_request(request)?;
let request = prepare_provider_request(request, options)?;
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(Error::InvalidRequest(
"streaming messages is not supported for this provider".to_string(),

7
litellm-rust/crates/core/src/messages/mod.rs
View file

@ -9,6 +9,7 @@
use crate::Error;
use crate::request_context::LiteLlmRequestContext;
use crate::request_options::RequestOptions;
mod client;
mod common_utils;
mod handler;
@ -22,16 +23,18 @@ use types::{AnthropicMessagesResponse, MessagesRequest};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn messages(
request: MessagesRequest<'_>,
options: &RequestOptions,
_context: &LiteLlmRequestContext,
) -> Result<AnthropicMessagesResponse, Error> {
execute_messages_provider_call(request).await
execute_messages_provider_call(request, options.clone()).await
}
pub async fn messages_stream(
request: MessagesRequest<'_>,
options: &RequestOptions,
_context: &LiteLlmRequestContext,
) -> Result<reqwest::Response, Error> {
execute_messages_provider_stream(request).await
execute_messages_provider_stream(request, options.clone()).await
}
#[cfg(test)]

43
litellm-rust/crates/core/src/messages/prepare.rs
View file

@ -1,4 +1,5 @@
use crate::error::Error;
use crate::request_options::RequestOptions;
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
@ -8,26 +9,24 @@ use serde_json::{Map, Value};
pub(super) fn prepare_provider_request(
request: MessagesRequest<'_>,
options: RequestOptions,
) -> Result<ProviderMessagesRequest, Error> {
let provider_info = get_custom_llm_provider(
request.model,
request.options.custom_llm_provider.as_deref(),
)
.or_else(|| {
request
.options
.custom_llm_provider
.as_deref()
.map(|provider| CustomLlmProvider {
model: request.model,
custom_llm_provider: provider,
let provider_info =
get_custom_llm_provider(request.model, options.custom_llm_provider.as_deref())
.or_else(|| {
options
.custom_llm_provider
.as_deref()
.map(|provider| CustomLlmProvider {
model: request.model,
custom_llm_provider: provider,
})
})
})
.ok_or_else(|| {
Error::InvalidProvider(
"unable to resolve custom_llm_provider for messages request".to_string(),
)
})?;
.ok_or_else(|| {
Error::InvalidProvider(
"unable to resolve custom_llm_provider for messages request".to_string(),
)
})?;
let model = provider_info.model.to_string();
let provider = provider_info.custom_llm_provider;
@ -37,8 +36,8 @@ pub(super) fn prepare_provider_request(
let headers = validate_environment(
config,
request.options.extra_headers,
request.options.api_key.as_deref(),
options.extra_headers,
options.api_key.as_deref(),
&env_lookup,
)?;
@ -52,7 +51,7 @@ pub(super) fn prepare_provider_request(
))
})?;
let url = config.complete_url(request.options.api_base.as_deref(), &model, &env_lookup)?;
let url = config.complete_url(options.api_base.as_deref(), &model, &env_lookup)?;
Ok(ProviderMessagesRequest {
provider: provider.to_string(),
@ -61,7 +60,7 @@ pub(super) fn prepare_provider_request(
url,
body,
upstream_headers: headers,
timeout: request.options.timeout,
timeout: options.timeout,
})
}

128
litellm-rust/crates/core/src/messages/tests.rs
View file

@ -148,14 +148,14 @@ async fn messages_round_trip_builds_azure_request_and_passes_response_through()
}]
}]
}),
options: RequestOptions {
api_key: (Some("sk-azure")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk-azure")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()
@ -212,14 +212,14 @@ async fn messages_round_trip_builds_native_anthropic_request() {
"max_tokens": 1024,
"messages": [{"role": "user", "content": "hi"}]
}),
options: RequestOptions {
api_key: (Some("sk-ant")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("anthropic")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk-ant")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("anthropic")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()
@ -273,14 +273,14 @@ async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() {
MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
options: RequestOptions {
api_key: (Some("rust-fallback-key")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("rust-fallback-key")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()
@ -335,14 +335,14 @@ async fn messages_forwards_entra_id_bearer_without_requiring_api_key() {
MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
options: RequestOptions {
api_key: None,
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: None,
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()
@ -367,14 +367,14 @@ async fn messages_requires_auth_when_no_key_and_no_header() {
MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
options: RequestOptions {
api_key: None,
api_base: (Some("http://127.0.0.1:1")).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_millis(50)),
..Default::default()
},
},
&RequestOptions {
api_key: None,
api_base: (Some("http://127.0.0.1:1")).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_millis(50)),
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()
@ -413,14 +413,14 @@ async fn messages_ignores_malformed_authorization_and_uses_api_key() {
MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
options: RequestOptions {
api_key: (Some("sk-azure")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk-azure")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()
@ -462,14 +462,14 @@ async fn messages_maps_provider_error_status_to_http_error() {
MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
options: RequestOptions {
api_key: (Some("sk-azure")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk-azure")).map(|value| value.to_string()),
api_base: (Some(&format!("http://{addr}"))).map(|value| value.to_string()),
custom_llm_provider: (Some("azure_ai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()
@ -487,14 +487,14 @@ async fn messages_rejects_unsupported_provider() {
MessagesRequest {
model: "claude-3-5-sonnet",
body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}),
options: RequestOptions {
api_key: (Some("sk")).map(|value| value.to_string()),
api_base: (Some("http://127.0.0.1:1")).map(|value| value.to_string()),
custom_llm_provider: (Some("openai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_millis(50)),
..Default::default()
},
},
&RequestOptions {
api_key: (Some("sk")).map(|value| value.to_string()),
api_base: (Some("http://127.0.0.1:1")).map(|value| value.to_string()),
custom_llm_provider: (Some("openai")).map(|value| value.to_string()),
extra_headers: None,
timeout: Some(Duration::from_millis(50)),
..Default::default()
},
&LiteLlmRequestContext {
..Default::default()

2
litellm-rust/crates/core/src/messages/types.rs
View file

@ -1,4 +1,3 @@
use crate::request_options::RequestOptions;
use std::time::Duration;
use serde::{Deserialize, Serialize};
@ -9,7 +8,6 @@ use super::transformation::AnthropicMessagesProviderConfig;
pub struct MessagesRequest<'a> {
pub model: &'a str,
pub body: Value,
pub options: RequestOptions,
}
pub(super) struct ProviderMessagesRequest {

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

@ -1,5 +1,3 @@
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, PartialEq)]
pub struct RequestAttribution {
pub user_api_key_hash: Option<String>,
@ -7,12 +5,19 @@ pub struct RequestAttribution {
pub user_api_key_team_id: Option<String>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct RequestCapabilities {
pub stream: bool,
pub has_agentic_hook: bool,
pub has_custom_client: bool,
pub request_format: Option<String>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct LiteLlmRequestContext {
pub metadata: Option<Map<String, Value>>,
pub litellm_metadata: Option<Map<String, Value>>,
pub request_metadata_fields: Vec<String>,
pub litellm_call_id: Option<String>,
pub trace_id: Option<String>,
pub request_model: Option<String>,
pub attribution: RequestAttribution,
pub capabilities: RequestCapabilities,
}

71
litellm-rust/crates/core/src/request_options.rs
View file

@ -2,6 +2,73 @@ use std::time::Duration;
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default)]
pub struct BedrockOptions {
pub aws_access_key_id: Option<String>,
pub aws_secret_access_key: Option<String>,
pub aws_session_token: Option<String>,
pub aws_region_name: Option<String>,
pub aws_session_name: Option<String>,
pub aws_profile_name: Option<String>,
pub aws_role_name: Option<String>,
pub aws_web_identity_token: Option<String>,
pub aws_sts_endpoint: Option<String>,
pub aws_external_id: Option<String>,
pub aws_bedrock_runtime_endpoint: Option<String>,
pub request_metadata_fields: Vec<String>,
pub request_metadata: Option<std::collections::BTreeMap<String, String>>,
}
impl BedrockOptions {
pub fn into_map(&self) -> Map<String, Value> {
[
("aws_access_key_id", self.aws_access_key_id.clone()),
("aws_secret_access_key", self.aws_secret_access_key.clone()),
("aws_session_token", self.aws_session_token.clone()),
("aws_region_name", self.aws_region_name.clone()),
("aws_session_name", self.aws_session_name.clone()),
("aws_profile_name", self.aws_profile_name.clone()),
("aws_role_name", self.aws_role_name.clone()),
(
"aws_web_identity_token",
self.aws_web_identity_token.clone(),
),
("aws_sts_endpoint", self.aws_sts_endpoint.clone()),
("aws_external_id", self.aws_external_id.clone()),
(
"aws_bedrock_runtime_endpoint",
self.aws_bedrock_runtime_endpoint.clone(),
),
]
.into_iter()
.filter_map(|(name, value)| value.map(|value| (name.to_string(), Value::String(value))))
.collect()
}
}
#[derive(Clone, Debug, Default)]
pub struct AnthropicOptions {
pub user_id: Option<String>,
}
#[derive(Clone, Debug, Default)]
pub struct VertexOptions {
pub project: Option<String>,
pub location: Option<String>,
}
impl VertexOptions {
pub fn into_map(&self) -> Map<String, Value> {
[
("vertex_project", self.project.clone()),
("vertex_location", self.location.clone()),
]
.into_iter()
.filter_map(|(name, value)| value.map(|value| (name.to_string(), Value::String(value))))
.collect()
}
}
#[derive(Clone, Debug, Default)]
pub struct RequestOptions {
pub api_key: Option<String>,
@ -10,5 +77,7 @@ pub struct RequestOptions {
pub extra_headers: Option<Map<String, Value>>,
pub extra_query: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
pub provider_connection: Map<String, Value>,
pub bedrock: Option<BedrockOptions>,
pub anthropic: Option<AnthropicOptions>,
pub vertex: Option<VertexOptions>,
}

1
litellm-rust/crates/core/src/responses/types.rs
View file

@ -4,7 +4,6 @@ use serde_json::{Map, Value};
#[derive(Clone, Debug)]
pub struct ResponsesWebSocketRequest {
pub url: String,
pub options: crate::request_options::RequestOptions,
}
#[derive(Clone, Debug, PartialEq, Eq)]

27
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -17,7 +17,6 @@ use crate::marshal::{NativeRequestContext, NativeRequestOptions};
#[derive(FromPyObject)]
struct WebSocketConnectRequest {
url: String,
options: NativeRequestOptions,
}
#[pyclass]
@ -28,21 +27,19 @@ struct ResponsesWebSocketConnection {
#[pymethods]
impl ResponsesWebSocketConnection {
#[classmethod]
#[pyo3(signature = (request, *, context))]
#[pyo3(signature = (request, *, options, context))]
fn connect<'py>(
_cls: &Bound<'py, pyo3::types::PyType>,
py: Python<'py>,
request: WebSocketConnectRequest,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<Bound<'py, PyAny>> {
let options: litellm_core::request_options::RequestOptions = request.options.into();
let options: litellm_core::request_options::RequestOptions = options.into();
let context: litellm_core::request_context::LiteLlmRequestContext = context.into();
let request = ResponsesWebSocketRequest {
url: request.url,
options,
};
let request = ResponsesWebSocketRequest { url: request.url };
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let inner = RustResponsesWebSocketConnection::connect(request, &context)
let inner = RustResponsesWebSocketConnection::connect(request, &options, &context)
.await
.map_err(core_error_to_pyerr)?;
Ok(ResponsesWebSocketConnection { inner })
@ -206,14 +203,14 @@ mod tests {
import asyncio
async def exercise():
for request, request_context, field in (
(Request(url=123), context, 'url'),
(Request(url=url, options=Options(extra_headers=[])), context, 'extra_headers'),
(Request(url=url), replace(context, litellm_call_id=123), 'litellm_call_id'),
(Request(url=url), replace(context, attribution=Attribution(user_api_key_user_id=123)), 'user_api_key_user_id'),
for request, request_options, request_context, field in (
(Request(url=123), options, context, 'url'),
(Request(url=url), Options(extra_headers=[]), context, 'extra_headers'),
(Request(url=url), options, replace(context, litellm_call_id=123), 'litellm_call_id'),
(Request(url=url), options, replace(context, attribution=Attribution(user_api_key_user_id=123)), 'user_api_key_user_id'),
):
try:
native.ResponsesWebSocketConnection.connect(request, context=request_context)
native.ResponsesWebSocketConnection.connect(request, options=request_options, context=request_context)
except (TypeError, ValueError) as error:
parts = []
while error is not None:
@ -223,7 +220,7 @@ async def exercise():
else:
raise AssertionError('invalid WebSocket input reached execution')
connection = await native.ResponsesWebSocketConnection.connect(Request(url=url), context=context)
connection = await native.ResponsesWebSocketConnection.connect(Request(url=url), options=options, context=context)
assert type(connection) is native.ResponsesWebSocketConnection
await connection.send_text("from-python")
assert await connection.recv_text() == "from-server"

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

@ -4,6 +4,71 @@ use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde_json::{Map, Value};
#[derive(FromPyObject)]
struct NativeBedrockOptions {
aws_access_key_id: Option<String>,
aws_secret_access_key: Option<String>,
aws_session_token: Option<String>,
aws_region_name: Option<String>,
aws_session_name: Option<String>,
aws_profile_name: Option<String>,
aws_role_name: Option<String>,
aws_web_identity_token: Option<String>,
aws_sts_endpoint: Option<String>,
aws_external_id: Option<String>,
aws_bedrock_runtime_endpoint: Option<String>,
request_metadata_fields: Vec<String>,
request_metadata: Option<std::collections::BTreeMap<String, String>>,
}
impl From<NativeBedrockOptions> for litellm_core::request_options::BedrockOptions {
fn from(input: NativeBedrockOptions) -> Self {
Self {
aws_access_key_id: input.aws_access_key_id,
aws_secret_access_key: input.aws_secret_access_key,
aws_session_token: input.aws_session_token,
aws_region_name: input.aws_region_name,
aws_session_name: input.aws_session_name,
aws_profile_name: input.aws_profile_name,
aws_role_name: input.aws_role_name,
aws_web_identity_token: input.aws_web_identity_token,
aws_sts_endpoint: input.aws_sts_endpoint,
aws_external_id: input.aws_external_id,
aws_bedrock_runtime_endpoint: input.aws_bedrock_runtime_endpoint,
request_metadata_fields: input.request_metadata_fields,
request_metadata: input.request_metadata,
}
}
}
#[derive(FromPyObject)]
struct NativeAnthropicOptions {
user_id: Option<String>,
}
impl From<NativeAnthropicOptions> for litellm_core::request_options::AnthropicOptions {
fn from(input: NativeAnthropicOptions) -> Self {
Self {
user_id: input.user_id,
}
}
}
#[derive(FromPyObject)]
struct NativeVertexOptions {
project: Option<String>,
location: Option<String>,
}
impl From<NativeVertexOptions> for litellm_core::request_options::VertexOptions {
fn from(input: NativeVertexOptions) -> Self {
Self {
project: input.project,
location: input.location,
}
}
}
#[derive(FromPyObject)]
pub(crate) struct NativeRequestOptions {
api_key: Option<String>,
@ -14,8 +79,9 @@ pub(crate) struct NativeRequestOptions {
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_query: Option<Map<String, Value>>,
timeout_seconds: Option<f64>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
provider_connection: Option<Map<String, Value>>,
bedrock: Option<NativeBedrockOptions>,
anthropic: Option<NativeAnthropicOptions>,
vertex: Option<NativeVertexOptions>,
}
impl From<NativeRequestOptions> for litellm_core::request_options::RequestOptions {
@ -27,7 +93,9 @@ impl From<NativeRequestOptions> for litellm_core::request_options::RequestOption
extra_headers: input.extra_headers,
extra_query: input.extra_query,
timeout: optional_timeout(input.timeout_seconds),
provider_connection: input.provider_connection.unwrap_or_default(),
bedrock: input.bedrock.map(Into::into),
anthropic: input.anthropic.map(Into::into),
vertex: input.vertex.map(Into::into),
}
}
}
@ -41,29 +109,38 @@ pub(crate) struct NativeRequestAttribution {
#[derive(FromPyObject)]
pub(crate) struct NativeRequestContext {
#[pyo3(from_py_with = litellm_python_interop::from_py)]
metadata: Option<Map<String, Value>>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
litellm_metadata: Option<Map<String, Value>>,
request_metadata_fields: Vec<String>,
litellm_call_id: Option<String>,
trace_id: Option<String>,
request_model: Option<String>,
attribution: NativeRequestAttribution,
capabilities: NativeRequestCapabilities,
}
#[derive(FromPyObject)]
struct NativeRequestCapabilities {
stream: bool,
has_agentic_hook: bool,
has_custom_client: bool,
request_format: Option<String>,
}
impl From<NativeRequestContext> for litellm_core::request_context::LiteLlmRequestContext {
fn from(input: NativeRequestContext) -> Self {
Self {
metadata: input.metadata,
litellm_metadata: input.litellm_metadata,
request_metadata_fields: input.request_metadata_fields,
litellm_call_id: input.litellm_call_id,
trace_id: input.trace_id,
request_model: input.request_model,
attribution: litellm_core::request_context::RequestAttribution {
user_api_key_hash: input.attribution.user_api_key_hash,
user_api_key_user_id: input.attribution.user_api_key_user_id,
user_api_key_team_id: input.attribution.user_api_key_team_id,
},
capabilities: litellm_core::request_context::RequestCapabilities {
stream: input.capabilities.stream,
has_agentic_hook: input.capabilities.has_agentic_hook,
has_custom_client: input.capabilities.has_custom_client,
request_format: input.capabilities.request_format,
},
}
}
}
@ -107,7 +184,37 @@ class Options:
extra_headers: object = None
extra_query: object = None
timeout_seconds: object = None
provider_connection: object = None
bedrock: object = None
anthropic: object = None
vertex: object = None
@dataclass(frozen=True)
class BedrockOptions:
aws_access_key_id: object = None
aws_secret_access_key: object = None
aws_session_token: object = None
aws_region_name: object = None
aws_session_name: object = None
aws_profile_name: object = None
aws_role_name: object = None
aws_web_identity_token: object = None
aws_sts_endpoint: object = None
aws_external_id: object = None
aws_bedrock_runtime_endpoint: object = None
request_metadata_fields: object = ()
request_metadata: object = None
@dataclass(frozen=True)
class Capabilities:
stream: object = False
has_agentic_hook: object = False
has_custom_client: object = False
request_format: object = None
@dataclass(frozen=True)
class VertexOptions:
project: object = None
location: object = None
@dataclass(frozen=True)
class Attribution:
@ -117,12 +224,11 @@ class Attribution:
@dataclass(frozen=True)
class Context:
metadata: object = None
litellm_metadata: object = None
request_metadata_fields: tuple = ()
litellm_call_id: object = None
trace_id: object = None
request_model: object = None
attribution: Attribution = Attribution()
capabilities: Capabilities = Capabilities()
@dataclass(frozen=True)
class Request:
@ -132,11 +238,11 @@ class Request:
audio: object = None
document: object = None
optional_params: object = None
options: Options = Options()
value: str = ''
url: str = ''
context = Context()
options = Options()
",
Some(&locals),
Some(&locals),

4
litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs
View file

@ -15,11 +15,11 @@ struct AudioTranscriptionInputs {
audio: Value,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Map<String, Value>,
options: NativeRequestOptions,
}
fn prepare_transcription(
input: AudioTranscriptionInputs,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
let context: LiteLlmRequestContext = context.into();
@ -30,8 +30,8 @@ fn prepare_transcription(
model: &input.model,
audio,
optional_params: input.optional_params,
options: input.options.into(),
},
&options.into(),
&context,
)
.await

4
litellm-rust/crates/python-bridge/src/routes/chat_completions.rs
View file

@ -16,11 +16,11 @@ struct ChatCompletionsInputs {
messages: Value,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Map<String, Value>,
options: NativeRequestOptions,
}
fn prepare_chat_completions(
input: ChatCompletionsInputs,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<impl Future<Output = Result<ChatCompletionsResponse, Error>> + Send + 'static> {
let context: LiteLlmRequestContext = context.into();
@ -31,8 +31,8 @@ fn prepare_chat_completions(
model: &input.model,
messages,
optional_params: input.optional_params,
options: input.options.into(),
},
&options.into(),
&context,
)
.await

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

@ -12,24 +12,26 @@ macro_rules! bridge_route {
$(, extra = [$($extra:ident),* $(,)?])? $(,)?
) => {
#[pyfunction]
#[pyo3(signature = (request, *, context))]
#[pyo3(signature = (request, *, options, context))]
fn $sync_name(
py: pyo3::Python<'_>,
request: $inputs,
options: $crate::marshal::NativeRequestOptions,
context: $crate::marshal::NativeRequestContext,
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare(request, context)?;
let future = $prepare(request, options, context)?;
$crate::execution::run_sync(py, future, $map_error)
}
#[pyfunction]
#[pyo3(signature = (request, *, context))]
#[pyo3(signature = (request, *, options, context))]
fn $async_name(
py: pyo3::Python<'_>,
request: $inputs,
options: $crate::marshal::NativeRequestOptions,
context: $crate::marshal::NativeRequestContext,
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare(request, context)?;
let future = $prepare(request, options, context)?;
$crate::execution::run_async(py, future, $map_error)
}
@ -46,24 +48,26 @@ macro_rules! bridge_route {
use super::{$inputs, $map_error, $prepare};
#[pyfunction]
#[pyo3(signature = (request, *, context))]
#[pyo3(signature = (request, *, options, context))]
fn $sync_name(
py: pyo3::Python<'_>,
request: $inputs,
options: $crate::marshal::NativeRequestOptions,
context: $crate::marshal::NativeRequestContext,
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare(request, context)?;
let future = $prepare(request, options, context)?;
$crate::execution::run_sync(py, $crate::function_trace::capture(future), $map_error)
}
#[pyfunction]
#[pyo3(signature = (request, *, context))]
#[pyo3(signature = (request, *, options, context))]
fn $async_name(
py: pyo3::Python<'_>,
request: $inputs,
options: $crate::marshal::NativeRequestOptions,
context: $crate::marshal::NativeRequestContext,
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare(request, context)?;
let future = $prepare(request, options, context)?;
$crate::execution::run_async(py, $crate::function_trace::capture(future), $map_error)
}
@ -141,6 +145,7 @@ mod tests {
fn prepare_echo(
inputs: EchoInputs,
_options: crate::marshal::NativeRequestOptions,
_context: crate::marshal::NativeRequestContext,
) -> PyResult<impl Future<Output = Result<String, Error>> + Send + 'static> {
FUTURE_DROPPED.store(false, Ordering::SeqCst);
@ -182,13 +187,17 @@ mod tests {
let module = PyModule::new(py, "routes").expect("module should be created");
crate::routes::register(&module).expect("routes should register");
let routes = [
("ocr", "aocr", "(request, *, context)"),
("transcription", "atranscription", "(request, *, context)"),
("messages", "amessages", "(request, *, context)"),
("ocr", "aocr", "(request, *, options, context)"),
(
"transcription",
"atranscription",
"(request, *, options, context)",
),
("messages", "amessages", "(request, *, options, context)"),
(
"chat_completions",
"achat_completions",
"(request, *, context)",
"(request, *, options, context)",
),
];
@ -219,17 +228,17 @@ mod tests {
let locals = crate::marshal::request_fixtures(py);
locals.set_item("routes", module).unwrap();
py.run(c"
for names, request, expected in [
(('chat_completions', 'achat_completions'), Request(messages={}, optional_params={}), 'messages must be a list'),
(('messages', 'amessages'), Request(body=[]), 'body must be a dict'),
(('ocr', 'aocr'), Request(document={}, optional_params={}, options=Options(extra_headers=[])), 'extra_headers'),
(('transcription', 'atranscription'), Request(audio={}, optional_params={}, options=Options(timeout_seconds='bad')), 'timeout_seconds'),
(('transcription', 'atranscription'), Request(audio={}, optional_params={}, options=Options(provider_connection=[])), 'provider_connection'),
for names, request, request_options, expected in [
(('chat_completions', 'achat_completions'), Request(messages={}, optional_params={}), options, 'messages must be a list'),
(('messages', 'amessages'), Request(body=[]), options, 'body must be a dict'),
(('ocr', 'aocr'), Request(document={}, optional_params={}), Options(extra_headers=[]), 'extra_headers'),
(('transcription', 'atranscription'), Request(audio={}, optional_params={}), Options(timeout_seconds='bad'), 'timeout_seconds'),
(('transcription', 'atranscription'), Request(audio={}, optional_params={}), Options(bedrock=[]), 'bedrock'),
]:
errors = []
for name in names:
try:
getattr(routes, name)(request, context=context)
getattr(routes, name)(request, options=request_options, context=context)
except (ValueError, TypeError) as error:
parts = []
while error is not None:
@ -241,10 +250,10 @@ for names, request, expected in [
assert errors[0] == errors[1], errors
assert expected in errors[0], (expected, errors)
for field in ('metadata', 'litellm_metadata', 'request_metadata_fields'):
for field in ('litellm_call_id', 'trace_id', 'request_model'):
invalid_context = replace(context, **{field: object()})
try:
routes.chat_completions(Request(messages=[], optional_params={}), context=invalid_context)
routes.chat_completions(Request(messages=[], optional_params={}), options=options, context=invalid_context)
except (ValueError, TypeError) as error:
assert field in str(error)
else:
@ -267,6 +276,7 @@ for field in ('metadata', 'litellm_metadata', 'request_metadata_fields'):
py.eval(c"Request(value=\"sync\")", Some(&locals), Some(&locals))
.and_then(|request| {
let kwargs = PyDict::new(py);
kwargs.set_item("options", locals.get_item("options")?.unwrap())?;
kwargs.set_item("context", locals.get_item("context")?.unwrap())?;
function.call((request,), Some(&kwargs))
})
@ -282,6 +292,7 @@ for field in ('metadata', 'litellm_metadata', 'request_metadata_fields'):
py.eval(c"Request(value=\"error\")", Some(&locals), Some(&locals))
.and_then(|request| {
let kwargs = PyDict::new(py);
kwargs.set_item("options", locals.get_item("options")?.unwrap())?;
kwargs.set_item("context", locals.get_item("context")?.unwrap())?;
function.call((request,), Some(&kwargs))
})
@ -302,17 +313,17 @@ for field in ('metadata', 'litellm_metadata', 'request_metadata_fields'):
import asyncio
async def exercise():
assert await routes.aecho(Request(value="async"), context=context) == "async"
assert await routes.aecho(Request(value="async"), options=options, context=context) == "async"
try:
await routes.aecho(Request(value="error"), context=context)
await routes.aecho(Request(value="error"), options=options, context=context)
except LookupError as error:
assert str(error) == "invalid request: synthetic error"
else:
raise AssertionError("mapped error was not raised")
try:
await routes.aecho(Request(value="panic"), context=context)
await routes.aecho(Request(value="panic"), options=options, context=context)
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "synthetic panic"
@ -320,14 +331,14 @@ async def exercise():
raise AssertionError("panic was not raised")
try:
await routes.aecho(Request(value="map_panic"), context=context)
await routes.aecho(Request(value="map_panic"), options=options, context=context)
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "synthetic mapper panic"
else:
raise AssertionError("mapper panic was not raised")
task = asyncio.ensure_future(routes.aecho(Request(value="pending"), context=context))
task = asyncio.ensure_future(routes.aecho(Request(value="pending"), options=options, context=context))
await asyncio.sleep(0)
task.cancel()
try:
@ -365,7 +376,7 @@ asyncio.run(exercise())
.expect("module should enter Python locals");
let code = CString::new(
r#"
result = routes.echo(Request(value="traced"), context=context)
result = routes.echo(Request(value="traced"), options=options, context=context)
assert result == {
"response": "traced",
"trace": [{"function": "execute_echo", "depth": 0}],

4
litellm-rust/crates/python-bridge/src/routes/messages.rs
View file

@ -13,11 +13,11 @@ struct MessagesInputs {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
body: Value,
options: NativeRequestOptions,
}
fn prepare_messages(
input: MessagesInputs,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<impl Future<Output = Result<AnthropicMessagesResponse, Error>> + Send + 'static> {
let context: LiteLlmRequestContext = context.into();
@ -27,8 +27,8 @@ fn prepare_messages(
MessagesRequest {
model: &input.model,
body,
options: input.options.into(),
},
&options.into(),
&context,
)
.await

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

@ -16,11 +16,11 @@ struct OcrInputs {
document: Value,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Map<String, Value>,
options: NativeRequestOptions,
}
fn prepare_ocr(
input: OcrInputs,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
let context: LiteLlmRequestContext = context.into();
@ -31,8 +31,8 @@ fn prepare_ocr(
model: &input.model,
document,
optional_params: input.optional_params,
options: input.options.into(),
},
&options.into(),
&context,
RequestHooks {
callbacks: Vec::new(),

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

@ -27,8 +27,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
from litellm.rust_bridge.dispatch import anative_first, native_first, provider_errors
from litellm.rust_bridge.runtime import DispatchResult
from litellm.rust_bridge.request import anthropic_options
from litellm.types.llms.anthropic import (
ContentBlockDelta,
ContentBlockStart,
@ -370,7 +369,15 @@ 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.
Shared by the normal path and by the Rust path's fallback, which
builds it only when the Rust call did not serve the request.
"""
request_data: Final = config.transform_request(
model=model,
messages=messages,
@ -378,29 +385,12 @@ 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
# 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
@ -417,26 +407,68 @@ class AnthropicChatCompletion(BaseLLM):
litellm_params=litellm_params,
stream=stream,
)
rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict
"complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent
"model": model,
"messages": messages,
**rust_optional_params,
},
"api_base": api_base,
"headers": headers,
}
if serves_via_rust:
rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict
"complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent
"model": model,
"messages": messages,
**rust_optional_params,
},
"api_base": api_base,
"headers": headers,
}
logging_obj.pre_call(input=messages, api_key=api_key, additional_args=rust_logging_args)
log_rust_post_call: Final = rust_chat_completions_bridge.response_logger(
logging_obj=logging_obj,
messages=messages,
api_key=api_key,
additional_args=rust_logging_args,
)
log_rust_post_call: Final = rust_chat_completions_bridge.response_logger(
logging_obj=logging_obj,
messages=messages,
api_key=api_key,
additional_args=rust_logging_args,
)
if acompletion is True:
def native_completion() -> DispatchResult[ModelResponse]:
return rust_chat_completions_bridge.chat_completions(
async def python_fallback() -> "ModelResponse | CustomStreamWrapper":
# pre_call already fired for this request above. The Rust
# path only declines before the provider is called, so this
# is the same attempt continuing, not a second one.
fallback_headers, fallback_data = build_request()
return await self.acompletion_function(
model=model,
messages=messages,
data=fallback_data,
api_base=api_base,
custom_prompt_dict=custom_prompt_dict,
model_response=model_response,
print_verbose=print_verbose,
encoding=encoding,
api_key=api_key,
provider_config=config,
logging_obj=logging_obj,
optional_params=optional_params,
stream=stream,
_is_function_call=_is_function_call,
litellm_params=litellm_params,
logger_fn=logger_fn,
headers=fallback_headers,
client=client,
json_mode=json_mode,
timeout=timeout,
)
return rust_chat_completions_bridge.achat_completions_or_fallback(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
python_fallback=python_fallback,
anthropic=anthropic_options(litellm_params),
)
rust_response: Final = rust_chat_completions_bridge.chat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
@ -447,37 +479,35 @@ class AnthropicChatCompletion(BaseLLM):
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
eligible=serves_via_rust,
anthropic=anthropic_options(litellm_params),
)
if rust_response is not None:
return rust_response
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,
headers, data = build_request()
## LOGGING
# Reaching here with `serves_via_rust` set means the Rust attempt
# declined at call time, before the provider was called, and already
# logged this request. That is the same attempt continuing.
if not serves_via_rust:
logging_obj.pre_call(
input=messages,
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,
eligible=serves_via_rust,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": headers,
},
)
@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,
@ -499,7 +529,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,
@ -521,14 +551,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
@ -560,12 +583,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),
@ -586,21 +610,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

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

@ -10,6 +10,10 @@ from litellm.anthropic_beta_headers_manager import (
update_headers_with_filtered_beta,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
from litellm.llms.bedrock.request_metadata import (
get_bedrock_request_metadata_fields,
resolve_bedrock_request_metadata,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
@ -18,8 +22,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
from litellm.rust_bridge.dispatch import anative_first, native_first, provider_errors
from litellm.rust_bridge.runtime import DispatchResult
from litellm.rust_bridge.request import NativeBedrockOptions
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
@ -395,11 +398,15 @@ class BedrockConverseLLM(BaseAWSLLM):
# resolved so both paths sign as the same principal. Bearer-token auth
# resolves no SigV4 principal at all, and each path reads that token
# itself.
rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy
**optional_params,
**_sigv4_principal(credentials),
"aws_region_name": aws_region_name,
}
rust_optional_params: Final = optional_params
rust_bedrock_options: Final = NativeBedrockOptions(
aws_access_key_id=None if credentials is None else credentials.access_key,
aws_secret_access_key=None if credentials is None else credentials.secret_key,
aws_session_token=None if credentials is None else credentials.token,
aws_region_name=aws_region_name,
request_metadata_fields=get_bedrock_request_metadata_fields(),
request_metadata=resolve_bedrock_request_metadata(litellm_params, optional_params.get("requestMetadata")),
)
serves_via_rust: Final = rust_chat_completions_accepts(
model=model,
messages=messages,
@ -408,25 +415,55 @@ class BedrockConverseLLM(BaseAWSLLM):
litellm_params=litellm_params,
stream=stream,
)
rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict
"complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent
"messages": messages,
**optional_params,
},
"api_base": proxy_endpoint_url,
"headers": headers,
}
if serves_via_rust:
rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict
"complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent
"messages": messages,
**optional_params,
},
"api_base": proxy_endpoint_url,
"headers": headers,
}
logging_obj.pre_call(input=messages, api_key="", additional_args=rust_logging_args)
log_rust_post_call: Final = rust_chat_completions_bridge.response_logger(
logging_obj=logging_obj,
messages=messages,
api_key="",
additional_args=rust_logging_args,
)
def native_completion() -> DispatchResult[ModelResponse]:
return rust_chat_completions_bridge.chat_completions(
log_rust_post_call: Final = rust_chat_completions_bridge.response_logger(
logging_obj=logging_obj,
messages=messages,
api_key="",
additional_args=rust_logging_args,
)
if acompletion:
return rust_chat_completions_bridge.achat_completions_or_fallback(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=proxy_endpoint_url,
custom_llm_provider="bedrock",
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
bedrock=rust_bedrock_options,
python_fallback=lambda: self.async_completion(
model=model,
messages=messages,
api_base=proxy_endpoint_url,
model_response=model_response,
encoding=encoding,
logging_obj=logging_obj,
optional_params=optional_params,
stream=stream,
litellm_params=litellm_params,
logger_fn=logger_fn,
headers=headers,
timeout=timeout,
client=client,
credentials=credentials,
api_key=api_key,
skip_pre_call_logging=True,
),
)
rust_response: Final = rust_chat_completions_bridge.chat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
@ -437,33 +474,17 @@ class BedrockConverseLLM(BaseAWSLLM):
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
eligible=serves_via_rust,
bedrock=rust_bedrock_options,
)
if rust_response is not None:
return rust_response
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,
eligible=serves_via_rust,
)
@anative_first(
native=native_acompletion,
route="chat_completions",
errors=lambda: provider_errors("bedrock", model),
)
async def execute_async() -> ModelResponse | CustomStreamWrapper:
python_client: Final = None if isinstance(client, HTTPHandler) else client
### 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,
@ -476,7 +497,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,
@ -484,7 +505,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,
@ -497,112 +518,108 @@ 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(
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,
)
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,
## LOGGING
# Reaching here with `serves_via_rust` set means the synchronous Rust
# attempt declined at call time, before the provider was called, and
# already logged this request. That is the same attempt continuing.
# The asynchronous branch above returns before this point, and hands
# its own fallback `skip_pre_call_logging=True` for the same reason.
if not serves_via_rust:
logging_obj.pre_call(
input=messages,
api_key="",
data=data,
messages=messages,
optional_params=optional_params,
encoding=encoding,
additional_args={
"complete_input_dict": data,
"api_base": proxy_endpoint_url,
"headers": prepped.headers,
},
)
sync_transformed_response.set_provider_response_headers(response.headers)
return sync_transformed_response
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
return execute_async() if acompletion else execute_sync()
if stream is not None and stream is True:
completion_stream, response_headers = make_sync_call(
client=(client if client is not None and isinstance(client, HTTPHandler) else None),
api_base=proxy_endpoint_url,
headers=prepped.headers,
data=data,
model=model,
messages=messages,
logging_obj=logging_obj,
json_mode=json_mode,
fake_stream=fake_stream,
stream_chunk_size=stream_chunk_size,
)
streaming_response: Final = CustomStreamWrapper(
completion_stream=completion_stream,
model=model,
custom_llm_provider="bedrock",
logging_obj=logging_obj,
_response_headers=response_headers,
)
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

6
litellm/llms/bedrock/request_metadata.py
View file

@ -43,7 +43,7 @@ def _text_pairs(source: object) -> tuple[tuple[str, str], ...]:
return tuple((key, value) for key, value in source.items() if isinstance(key, str) and isinstance(value, str))
def _allowed_fields() -> tuple[str, ...]:
def get_bedrock_request_metadata_fields() -> tuple[str, ...]:
"""
The operator allow-list, deduplicated so a field repeated in config cannot consume a second
reserved slot and shrink the client budget for nothing. First occurrence wins, which keeps
@ -121,7 +121,7 @@ def resolve_bedrock_request_metadata(
been validated (and rejected with a 400) by the Converse transformation, so it is only
filtered here for the reserved identity prefix and the remaining slot budget.
"""
allowed_fields: Final = _allowed_fields()
allowed_fields: Final = get_bedrock_request_metadata_fields()
if not allowed_fields:
return None
sources: Final = _metadata_sources(litellm_params)
@ -146,7 +146,7 @@ def bedrock_request_metadata_is_owned() -> bool:
"fall back to whatever the caller supplied", or the reserved-prefix guarantee is bypassable
by anyone who can make the resolver produce nothing.
"""
return bool(_allowed_fields())
return bool(get_bedrock_request_metadata_fields())
def bedrock_request_metadata_headers(

25
litellm/ocr/main.py
View file

@ -32,8 +32,7 @@ from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.request import (
NativeRequestOptions,
PreparedNativeCall,
provider_connection_params,
provider_request_params,
vertex_options,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.router import GenericLiteLLMParams
@ -275,18 +274,18 @@ def _prepare_rust_ocr_call(
request=rust_ocr_bridge.NativeOCRRequest(
model=prepared_request.model,
document=prepared_request.document,
optional_params=provider_request_params(rust_optional_params),
options=NativeRequestOptions(
provider_connection=provider_connection_params(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),
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),
),
)

67
litellm/rust_bridge/chat_completions.py
View file

@ -32,13 +32,13 @@ from litellm.rust_bridge.protocols import (
RustChatCompletionsDecline,
)
from litellm.rust_bridge.request import (
NativeAnthropicOptions,
NativeBedrockOptions,
NativeChatCompletionsRequest,
NativeRequestContext,
NativeRequestOptions,
PreparedNativeCall,
call_native,
provider_connection_params,
provider_request_params,
)
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
@ -238,6 +238,8 @@ def chat_completions(
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
bedrock: NativeBedrockOptions | None = None,
anthropic: NativeAnthropicOptions | None = None,
) -> ModelResponse | None:
def adapt(rust_response: Mapping[str, object]) -> ModelResponse:
on_response(rust_response)
@ -248,15 +250,16 @@ def chat_completions(
NativeChatCompletionsRequest(
model=model,
messages=messages,
optional_params=provider_request_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),
provider_connection=provider_connection_params(optional_params),
),
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(),
),
@ -279,6 +282,8 @@ async def achat_completions(
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
bedrock: NativeBedrockOptions | None = None,
anthropic: NativeAnthropicOptions | None = None,
) -> ModelResponse | None:
def adapt(rust_response: Mapping[str, object]) -> ModelResponse:
on_response(rust_response)
@ -289,15 +294,16 @@ async def achat_completions(
NativeChatCompletionsRequest(
model=model,
messages=messages,
optional_params=provider_request_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),
provider_connection=provider_connection_params(optional_params),
),
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(),
),
@ -321,6 +327,8 @@ async def achat_completions_or_fallback(
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.
@ -340,15 +348,16 @@ async def achat_completions_or_fallback(
NativeChatCompletionsRequest(
model=model,
messages=messages,
optional_params=provider_request_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),
provider_connection=provider_connection_params(optional_params),
),
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(),
),

28
litellm/rust_bridge/messages.py
View file

@ -72,13 +72,13 @@ def messages(
NativeMessagesRequest(
model=model,
body=body,
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),
),
),
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),
),
context=NativeRequestContext(),
),
@ -104,13 +104,13 @@ async def amessages(
NativeMessagesRequest(
model=model,
body=body,
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),
),
),
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),
),
context=NativeRequestContext(),
),

2
litellm/rust_bridge/protocols.py
View file

@ -9,6 +9,7 @@ from .request import (
NativeMessagesRequest,
NativeOCRRequest,
NativeRequestContext,
NativeRequestOptions,
NativeResponsesWebSocketRequest,
NativeTranscriptionRequest,
)
@ -47,6 +48,7 @@ class RustResponsesWebSocketConnection(Protocol):
cls,
request: NativeResponsesWebSocketRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> RustResponsesWebSocket: ...

131
litellm/rust_bridge/request.py
View file

@ -2,7 +2,70 @@ from __future__ import annotations
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Generic, Protocol, TypeVar
from typing import Generic, Protocol, TypeVar
@dataclass(frozen=True, slots=True)
class NativeBedrockOptions:
aws_access_key_id: str | None = None
aws_secret_access_key: str | None = None
aws_session_token: str | None = None
aws_region_name: str | None = None
aws_session_name: str | None = None
aws_profile_name: str | None = None
aws_role_name: str | None = None
aws_web_identity_token: str | None = None
aws_sts_endpoint: str | None = None
aws_external_id: str | None = None
aws_bedrock_runtime_endpoint: str | None = None
request_metadata_fields: tuple[str, ...] = ()
request_metadata: Mapping[str, str] | None = None
@dataclass(frozen=True, slots=True)
class NativeAnthropicOptions:
user_id: str | None = None
@dataclass(frozen=True, slots=True)
class NativeVertexOptions:
project: str | None = None
location: str | None = None
def bedrock_options(params: Mapping[str, object]) -> NativeBedrockOptions:
def string(name: str) -> str | None:
value = params.get(name)
return value if isinstance(value, str) else None
return NativeBedrockOptions(
aws_access_key_id=string("aws_access_key_id"),
aws_secret_access_key=string("aws_secret_access_key"),
aws_session_token=string("aws_session_token"),
aws_region_name=string("aws_region_name"),
aws_session_name=string("aws_session_name"),
aws_profile_name=string("aws_profile_name"),
aws_role_name=string("aws_role_name"),
aws_web_identity_token=string("aws_web_identity_token"),
aws_sts_endpoint=string("aws_sts_endpoint"),
aws_external_id=string("aws_external_id"),
aws_bedrock_runtime_endpoint=string("aws_bedrock_runtime_endpoint"),
)
def anthropic_options(litellm_params: Mapping[str, object] | None) -> NativeAnthropicOptions:
metadata = None if litellm_params is None else litellm_params.get("metadata")
user_id = metadata.get("user_id") if isinstance(metadata, Mapping) else None
return NativeAnthropicOptions(user_id=user_id if isinstance(user_id, str) else None)
def vertex_options(params: Mapping[str, object]) -> NativeVertexOptions:
project = params.get("vertex_project") or params.get("vertex_ai_project")
location = params.get("vertex_location") or params.get("vertex_ai_location")
return NativeVertexOptions(
project=project if isinstance(project, str) else None,
location=location if isinstance(location, str) else None,
)
@dataclass(frozen=True, slots=True)
@ -13,7 +76,9 @@ class NativeRequestOptions:
extra_headers: Mapping[str, object] | None = None
extra_query: Mapping[str, object] | None = None
timeout_seconds: float | None = None
provider_connection: Mapping[str, object] | None = None
bedrock: NativeBedrockOptions | None = None
anthropic: NativeAnthropicOptions | None = None
vertex: NativeVertexOptions | None = None
@dataclass(frozen=True, slots=True)
@ -23,14 +88,21 @@ class RequestAttribution:
user_api_key_team_id: str | None = None
@dataclass(frozen=True, slots=True)
class NativeRequestCapabilities:
stream: bool = False
has_agentic_hook: bool = False
has_custom_client: bool = False
request_format: str | None = None
@dataclass(frozen=True, slots=True)
class NativeRequestContext:
metadata: Mapping[str, object] | None = None
litellm_metadata: Mapping[str, object] | None = None
request_metadata_fields: tuple[str, ...] = ()
litellm_call_id: str | None = None
trace_id: str | None = None
request_model: str | None = None
attribution: RequestAttribution = RequestAttribution()
capabilities: NativeRequestCapabilities = NativeRequestCapabilities()
RequestT = TypeVar("RequestT")
@ -41,48 +113,22 @@ ResultT = TypeVar("ResultT", covariant=True)
@dataclass(frozen=True, slots=True)
class PreparedNativeCall(Generic[RequestT]):
request: RequestT
options: NativeRequestOptions = NativeRequestOptions()
context: NativeRequestContext = NativeRequestContext()
class NativeFunction(Protocol[RequestContraT, ResultT]):
def __call__(self, request: RequestContraT, *, context: NativeRequestContext) -> ResultT: ...
def __call__(
self,
request: RequestContraT,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> ResultT: ...
def call_native(native: NativeFunction[RequestT, ResultT], prepared: PreparedNativeCall[RequestT]) -> ResultT:
return native(prepared.request, context=prepared.context)
_PROVIDER_CONNECTION_FIELDS: Final = frozenset(
(
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
"aws_region_name",
"aws_session_name",
"aws_profile_name",
"aws_role_name",
"aws_web_identity_token",
"aws_sts_endpoint",
"aws_external_id",
"aws_bedrock_runtime_endpoint",
"vertex_project",
"vertex_ai_project",
"vertex_location",
"vertex_ai_location",
)
)
def provider_connection_params(params: Mapping[str, object]) -> dict[str, object]:
return { # mutable-ok: PyO3 boundary payload
key: value for key, value in params.items() if key in _PROVIDER_CONNECTION_FIELDS
}
def provider_request_params(params: Mapping[str, object]) -> dict[str, object]:
return { # mutable-ok: PyO3 boundary payload
key: value for key, value in params.items() if key not in _PROVIDER_CONNECTION_FIELDS
}
return native(prepared.request, options=prepared.options, context=prepared.context)
@dataclass(frozen=True, slots=True)
@ -90,14 +136,12 @@ class NativeChatCompletionsRequest:
model: str
messages: Sequence[object]
optional_params: Mapping[str, object]
options: NativeRequestOptions
@dataclass(frozen=True, slots=True)
class NativeMessagesRequest:
model: str
body: dict[str, object]
options: NativeRequestOptions
@dataclass(frozen=True, slots=True)
@ -105,7 +149,6 @@ class NativeOCRRequest:
model: str
document: dict[str, object]
optional_params: dict[str, object]
options: NativeRequestOptions
@dataclass(frozen=True, slots=True)
@ -113,10 +156,8 @@ class NativeTranscriptionRequest:
model: str
audio: dict[str, object]
optional_params: dict[str, object]
options: NativeRequestOptions
@dataclass(frozen=True, slots=True)
class NativeResponsesWebSocketRequest:
url: str
options: NativeRequestOptions

2
litellm/rust_bridge/responses_websocket.py
View file

@ -73,8 +73,8 @@ async def connect(
prepare=lambda: PreparedNativeCall(
NativeResponsesWebSocketRequest(
url=url,
options=NativeRequestOptions(extra_headers=headers, timeout_seconds=timeout_to_seconds(timeout)),
),
options=NativeRequestOptions(extra_headers=headers, timeout_seconds=timeout_to_seconds(timeout)),
context=NativeRequestContext(),
),
call=lambda connection_type, request: call_native(connection_type.connect, request),

39
litellm/rust_bridge/transcription.py
View file

@ -11,9 +11,8 @@ from litellm.rust_bridge.request import (
NativeRequestOptions,
NativeTranscriptionRequest,
PreparedNativeCall,
bedrock_options,
call_native,
provider_connection_params,
provider_request_params,
)
from litellm.rust_bridge.runtime import (
BridgeErrorContext,
@ -74,15 +73,15 @@ def transcription(
NativeTranscriptionRequest(
model=model,
audio=audio,
optional_params=provider_request_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),
provider_connection=provider_connection_params(optional_params),
),
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_options(optional_params),
),
context=NativeRequestContext(),
),
@ -109,15 +108,15 @@ async def atranscription(
NativeTranscriptionRequest(
model=model,
audio=audio,
optional_params=provider_request_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),
provider_connection=provider_connection_params(optional_params),
),
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_options(optional_params),
),
context=NativeRequestContext(),
),

49
tests/rust-python-harness/strategies/trace_parity/sdk/execution.py
View file

@ -69,8 +69,8 @@ def _native_kwargs(route: str, kwargs: dict[str, object]) -> dict[str, object]:
from litellm.rust_bridge.request import (
NativeRequestContext,
NativeRequestOptions,
provider_connection_params,
provider_request_params,
bedrock_options,
vertex_options,
)
from litellm.rust_bridge.transcription import NativeTranscriptionRequest
@ -88,25 +88,36 @@ def _native_kwargs(route: str, kwargs: dict[str, object]) -> dict[str, object]:
"timeout_seconds",
)
},
"provider_connection": provider_connection_params(params),
"bedrock": bedrock_options(params),
"vertex": vertex_options(params),
}
)
payload: Final = {
key: value
for key, value in kwargs.items()
if key not in {"api_key", "api_base", "custom_llm_provider", "extra_headers", "timeout_seconds"}
}
request_type: Final = {
"chat_completions": NativeChatCompletionsRequest,
"messages": NativeMessagesRequest,
"ocr": NativeOCRRequest,
"transcription": NativeTranscriptionRequest,
"audio_transcription": NativeTranscriptionRequest,
}[route]
request: Final = TypeAdapter(request_type).validate_python(
{**payload, "optional_params": provider_request_params(params), "options": options}
)
return {"request": request, "context": NativeRequestContext()}
if route == "chat_completions":
request: Final = NativeChatCompletionsRequest(
model=TypeAdapter(str).validate_python(kwargs.get("model")),
messages=TypeAdapter(list[object]).validate_python(kwargs.get("messages")),
optional_params=params,
)
elif route == "messages":
request = NativeMessagesRequest(
model=TypeAdapter(str).validate_python(kwargs.get("model")),
body=TypeAdapter(dict[str, object]).validate_python(kwargs.get("body")),
)
elif route == "ocr":
request = NativeOCRRequest(
model=TypeAdapter(str).validate_python(kwargs.get("model")),
document=TypeAdapter(dict[str, object]).validate_python(kwargs.get("document")),
optional_params=params,
)
elif route in {"transcription", "audio_transcription"}:
request = NativeTranscriptionRequest(
model=TypeAdapter(str).validate_python(kwargs.get("model")),
audio=TypeAdapter(dict[str, object]).validate_python(kwargs.get("audio")),
optional_params=params,
)
else:
raise ValueError(f"unsupported native trace route: {route}")
return {"request": request, "options": options, "context": NativeRequestContext()}
def collect_trace(

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

@ -43,17 +43,18 @@ class RecordingMessages:
self,
request: NativeMessagesRequest,
*,
options: object,
context: NativeRequestContext,
) -> dict[str, object]:
self.calls.append(
{
"model": request.model,
"body": request.body,
"api_key": request.options.api_key,
"api_base": request.options.api_base,
"custom_llm_provider": request.options.custom_llm_provider,
"extra_headers": request.options.extra_headers,
"timeout_seconds": request.options.timeout_seconds,
"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,
}
)
return dict(FAKE_MESSAGES_RESPONSE)
@ -67,17 +68,18 @@ class RecordingAsyncMessages:
self,
request: NativeMessagesRequest,
*,
options: object,
context: NativeRequestContext,
) -> dict[str, object]:
self.calls.append(
{
"model": request.model,
"body": request.body,
"api_key": request.options.api_key,
"api_base": request.options.api_base,
"custom_llm_provider": request.options.custom_llm_provider,
"extra_headers": request.options.extra_headers,
"timeout_seconds": request.options.timeout_seconds,
"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,
}
)
return dict(FAKE_MESSAGES_RESPONSE)
@ -87,7 +89,9 @@ class ExplodingAsyncMessages:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, request: NativeMessagesRequest, *, context: NativeRequestContext) -> 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")
@ -96,7 +100,9 @@ class RaisingAsyncMessages:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, request: NativeMessagesRequest, *, context: NativeRequestContext) -> 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")

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

@ -23,9 +23,7 @@ async def test_make_call_passes_logging_obj_to_client_post():
mock_client = AsyncMock()
mock_response = MagicMock()
mock_response.aiter_lines = MagicMock(
return_value=iter(
[b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']
)
return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n'])
)
mock_client.post.return_value = mock_response
@ -94,9 +92,7 @@ def test_redacted_thinking_content_block_delta():
"data": "EuoBCoYBGAIiQJ/SxkPAgqxhKok29YrpJHRUJ0OT8ahCHKAwyhmRuUhtdmDX9+mn4gDzKNv3fVpQdB01zEPMzNY3QuTCd+1bdtEqQK6JuKHqdndbwpr81oVWb4wxd1GqF/7Jkw74IlQa27oobX+KuRkopr9Dllt/RDe7Se0sI1IkU7tJIAQCoP46OAwSDF51P09q67xhHlQ3ihoM2aOVlkghq/X0w8NlIjBMNvXYNbjhyrOcIg6kPFn2ed/KK7Cm5prYAtXCwkb4Wr5tUSoSHu9T5hKdJRbr6WsqEc7Lle7FULqMLZGkhqXyc3BA",
},
}
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
model_response = model_response_iterator.chunk_parser(chunk=chunk)
print(f"\n\nmodel_response: {model_response}\n\n")
assert model_response.choices[0].delta.thinking_blocks is not None
@ -104,19 +100,14 @@ def test_redacted_thinking_content_block_delta():
print(
f"\n\nmodel_response.choices[0].delta.thinking_blocks[0]: {model_response.choices[0].delta.thinking_blocks[0]}\n\n"
)
assert (
model_response.choices[0].delta.thinking_blocks[0]["type"]
== "redacted_thinking"
)
assert model_response.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking"
assert model_response.choices[0].delta.provider_specific_fields is not None
assert "thinking_blocks" in model_response.choices[0].delta.provider_specific_fields
def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
chunks = [
{
"type": "content_block_start",
@ -140,17 +131,12 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
expected_delta_blocks = (
{"type": "thinking", "thinking": "Step 1. "},
@ -164,18 +150,12 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
assert reasoning_content == "Step 1. Step 2."
assert thinking_blocks == (*expected_delta_blocks, expected_thinking_block)
assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {
"thinking_blocks": [expected_delta_blocks[0]]
}
assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {
"thinking_blocks": [expected_thinking_block]
}
assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_delta_blocks[0]]}
assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_thinking_block]}
def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
chunks = [
{
"type": "content_block_start",
@ -195,17 +175,12 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
{"type": "content_block_stop", "index": 0},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
assert reasoning_content == "Step 1. Step 2."
@ -216,9 +191,7 @@ def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
chunks = [
{
"type": "content_block_start",
@ -237,17 +210,12 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
},
]
parsed_chunks = [
model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
]
parsed_chunks = [model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks]
reasoning_content = "".join(
getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
for chunk in parsed_chunks
getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in parsed_chunks
)
thinking_blocks = tuple(
block
for chunk in parsed_chunks
for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
block for chunk in parsed_chunks for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
)
assert reasoning_content == "Step 1. Step 2."
@ -258,9 +226,7 @@ def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
def test_handle_json_mode_chunk_response_format_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
response_format_tool = ChatCompletionToolCallChunk(
id="tool_123",
type="function",
@ -271,9 +237,7 @@ def test_handle_json_mode_chunk_response_format_tool():
index=0,
)
text, tool_use = model_response_iterator._handle_json_mode_chunk(
"", response_format_tool
)
text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool)
print(f"\n\nresponse_format_tool text: {text}\n\n")
print(f"\n\nresponse_format_tool tool_use: {tool_use}\n\n")
@ -282,15 +246,11 @@ def test_handle_json_mode_chunk_response_format_tool():
def test_handle_json_mode_chunk_regular_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
regular_tool = ChatCompletionToolCallChunk(
id="tool_456",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name="get_weather", arguments='{"location": "San Francisco, CA"}'
),
function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'),
index=0,
)
@ -304,17 +264,13 @@ def test_handle_json_mode_chunk_regular_tool():
def test_handle_json_mode_chunk_streaming_response_format_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
# First chunk: response_format tool with id and name, but no arguments
first_chunk = ChatCompletionToolCallChunk(
id="tool_123",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=RESPONSE_FORMAT_TOOL_NAME, arguments=""
),
function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments=""),
index=0,
)
@ -322,9 +278,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
second_chunk = ChatCompletionToolCallChunk(
id=None,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=None, arguments='{"question": "What is the weather?"'
),
function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"question": "What is the weather?"'),
index=0,
)
@ -332,9 +286,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
third_chunk = ChatCompletionToolCallChunk(
id=None,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=None, arguments=', "answer": "It is sunny"}'
),
function=ChatCompletionToolCallFunctionChunk(name=None, arguments=', "answer": "It is sunny"}'),
index=0,
)
@ -365,9 +317,7 @@ def test_handle_json_mode_chunk_streaming_response_format_tool():
def test_handle_json_mode_chunk_streaming_regular_tool():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
# First chunk: regular tool with id and name, but no arguments
first_chunk = ChatCompletionToolCallChunk(
@ -381,9 +331,7 @@ def test_handle_json_mode_chunk_streaming_regular_tool():
second_chunk = ChatCompletionToolCallChunk(
id=None,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=None, arguments='{"location": "San Francisco, CA"}'
),
function=ChatCompletionToolCallFunctionChunk(name=None, arguments='{"location": "San Francisco, CA"}'),
index=0,
)
@ -408,27 +356,19 @@ def test_handle_json_mode_chunk_streaming_regular_tool():
def test_response_format_tool_finish_reason():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
# First chunk: response_format tool
response_format_tool = ChatCompletionToolCallChunk(
id="tool_123",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}'
),
function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments='{"answer": "test"}'),
index=0,
)
# Process the tool call (should set converted_response_format_tool flag)
text, tool_use = model_response_iterator._handle_json_mode_chunk(
"", response_format_tool
)
print(
f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n"
)
text, tool_use = model_response_iterator._handle_json_mode_chunk("", response_format_tool)
print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n")
# Simulate message_delta chunk with tool_use stop_reason
message_delta_chunk = {
@ -447,25 +387,19 @@ def test_response_format_tool_finish_reason():
def test_regular_tool_finish_reason():
model_response_iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=True
)
model_response_iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=True)
# First chunk: regular tool (not response_format)
regular_tool = ChatCompletionToolCallChunk(
id="tool_456",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name="get_weather", arguments='{"location": "San Francisco, CA"}'
),
function=ChatCompletionToolCallFunctionChunk(name="get_weather", arguments='{"location": "San Francisco, CA"}'),
index=0,
)
# Process the tool call (should NOT set converted_response_format_tool flag)
text, tool_use = model_response_iterator._handle_json_mode_chunk("", regular_tool)
print(
f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n"
)
print(f"\n\nconverted_response_format_tool flag: {model_response_iterator.converted_response_format_tool}\n\n")
# Simulate message_delta chunk with tool_use stop_reason
message_delta_chunk = {
@ -525,9 +459,7 @@ def test_text_only_streaming_has_index_zero():
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
if parsed.choices:
assert (
parsed.choices[0].index == 0
), f"Expected index=0, got {parsed.choices[0].index}"
assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}"
def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage():
@ -704,9 +636,7 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin
]
self._write_response(
content_type="text/event-stream",
body="".join(
f"data: {json.dumps(event)}\n\n" for event in events
).encode("utf-8"),
body="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode("utf-8"),
)
return
@ -787,13 +717,9 @@ def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinkin
assert content_chunks == [answer_text]
assert stream_usage is not None
stream_completion_details = stream_usage["completion_tokens_details"]
assert (
stream_completion_details["reasoning_tokens"]
== non_stream_details.reasoning_tokens
)
assert stream_completion_details["reasoning_tokens"] == non_stream_details.reasoning_tokens
assert stream_completion_details["text_tokens"] == (
stream_usage["completion_tokens"]
- stream_completion_details["reasoning_tokens"]
stream_usage["completion_tokens"] - stream_completion_details["reasoning_tokens"]
)
assert requests_seen == [
{
@ -885,9 +811,9 @@ def test_text_and_tool_streaming_has_index_zero():
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
if parsed.choices:
assert (
parsed.choices[0].index == 0
), f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}"
assert parsed.choices[0].index == 0, (
f"Expected index=0 for chunk type {chunk.get('type')}, got {parsed.choices[0].index}"
)
def test_multiple_tools_streaming_has_index_zero():
@ -940,15 +866,11 @@ def test_multiple_tools_streaming_has_index_zero():
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
if parsed.choices:
assert (
parsed.choices[0].index == 0
), f"Expected index=0, got {parsed.choices[0].index}"
assert parsed.choices[0].index == 0, f"Expected index=0, got {parsed.choices[0].index}"
def test_streaming_chunks_have_stable_ids():
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
first_chunk = {
"type": "content_block_delta",
"index": 0,
@ -973,9 +895,7 @@ def test_partial_json_chunk_accumulation():
This tests the fix for https://github.com/BerriAI/litellm/issues/17473
where network fragmentation can cause SSE data to arrive in partial chunks.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
partial_chunk_1 = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel'
partial_chunk_2 = 'lo"}}'
@ -983,31 +903,21 @@ def test_partial_json_chunk_accumulation():
# First partial chunk should return None (still accumulating)
result1 = iterator._parse_sse_data(f"data:{partial_chunk_1}")
assert result1 is None, "First partial chunk should return None while accumulating"
assert (
iterator.chunk_type == "accumulated_json"
), "Should switch to accumulated_json mode"
assert (
iterator.accumulated_json == partial_chunk_1
), "Should have accumulated first part"
assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
assert iterator.accumulated_json == partial_chunk_1, "Should have accumulated first part"
# Second partial chunk should complete the JSON and return a parsed result
result2 = iterator._parse_sse_data(f"data:{partial_chunk_2}")
assert result2 is not None, "Second chunk should return parsed result"
assert (
iterator.accumulated_json == ""
), "Buffer should be cleared after successful parse"
assert (
result2.choices[0].delta.content == "Hello"
), f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
assert iterator.accumulated_json == "", "Buffer should be cleared after successful parse"
assert result2.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
def test_complete_json_chunk_no_accumulation():
"""
Test that complete JSON chunks are parsed immediately without accumulation.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
complete_chunk = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}'
@ -1015,18 +925,14 @@ def test_complete_json_chunk_no_accumulation():
assert result is not None, "Complete chunk should return parsed result immediately"
assert iterator.chunk_type == "valid_json", "Should remain in valid_json mode"
assert iterator.accumulated_json == "", "Buffer should remain empty"
assert (
result.choices[0].delta.content == "Hello"
), f"Expected 'Hello', got '{result.choices[0].delta.content}'"
assert result.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result.choices[0].delta.content}'"
def test_multiple_partial_chunks_accumulation():
"""
Test that multiple partial chunks can be accumulated across several iterations.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Split a JSON chunk into three parts
part1 = '{"type":"content_block_del'
@ -1054,17 +960,11 @@ def test_accumulated_json_partial_fragment_returns_none_without_parsing():
unlike Vertex which already deferred parsing until the buffer could close.
A fragment that can't close a JSON value must not trigger a decode attempt.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
iterator.chunk_type = "accumulated_json"
with patch.object(
json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode
) as spy:
result = iterator._handle_accumulated_json_chunk(
'{"type":"content_block_delta","index":0,"delta":'
)
with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy:
result = iterator._handle_accumulated_json_chunk('{"type":"content_block_delta","index":0,"delta":')
assert result is None
assert spy.call_count == 0, "incomplete buffer should not be parsed"
@ -1076,21 +976,15 @@ def test_accumulated_json_does_not_reparse_every_fragment():
fragment.
"""
text = "x" * 200_000
blob = json.dumps(
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}
)
blob = json.dumps({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}})
fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)]
assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug"
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
iterator.chunk_type = "accumulated_json"
parsed = None
with patch.object(
json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode
) as spy:
with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy:
for fragment in fragments:
out = iterator._handle_accumulated_json_chunk(fragment)
if out is not None:
@ -1114,9 +1008,7 @@ def test_accumulated_json_concatenated_envelopes_do_not_wedge():
and keeps the remainder, so both values surface across two calls.
"""
obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}'
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
iterator.chunk_type = "accumulated_json"
first = iterator._handle_accumulated_json_chunk(obj + obj)
@ -1137,9 +1029,7 @@ def test_accumulated_json_heuristic_passes_but_value_still_incomplete():
heuristic must let the parse attempt through, and pop_next_value
finding nothing must propagate as None rather than raising.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
iterator.chunk_type = "accumulated_json"
result = iterator._handle_accumulated_json_chunk('{"type": {"nested": 1}')
@ -1153,9 +1043,7 @@ def test_accumulated_json_setter_and_sync_end_of_stream_drain():
underlying stream ends, instead of being silently dropped.
"""
obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}'
iterator = ModelResponseIterator(
streaming_response=iter([]), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False)
iterator.chunk_type = "accumulated_json"
iterator.accumulated_json = obj # exercises the setter
@ -1170,9 +1058,7 @@ def test_accumulated_json_async_end_of_stream_drain():
import asyncio
obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}'
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
iterator.chunk_type = "accumulated_json"
iterator.accumulated_json = obj
mock_async_iterator = MagicMock()
@ -1194,9 +1080,7 @@ def test_web_search_tool_result_no_extra_tool_calls():
The issue was that web_search_tool_result blocks have input_json_delta events with {}
that were incorrectly being converted to tool calls.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Simulate the streaming sequence:
# 1. server_tool_use block starts (web_search)
@ -1271,9 +1155,7 @@ def test_web_search_tool_result_no_extra_tool_calls():
# Should have exactly 2 tool calls:
# 1. From content_block_start (server_tool_use) with id and name
# 2. From content_block_delta with the actual query
assert (
len(tool_calls_emitted) == 2
), f"Expected 2 tool calls, got {len(tool_calls_emitted)}"
assert len(tool_calls_emitted) == 2, f"Expected 2 tool calls, got {len(tool_calls_emitted)}"
# First tool call should have the id and name
assert tool_calls_emitted[0]["id"] == "srvtoolu_01ABC123"
@ -1289,9 +1171,7 @@ def test_current_content_block_type_tracking():
"""
Test that current_content_block_type is properly tracked and reset.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Initially should be None
assert iterator.current_content_block_type is None
@ -1344,9 +1224,7 @@ def test_web_search_tool_result_captured_in_provider_specific_fields():
The web_search_tool_result content comes ALL AT ONCE in content_block_start,
not in deltas, so we need to capture it there.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Simulate the streaming sequence with web_search_tool_result
chunks = [
@ -1417,23 +1295,15 @@ def test_web_search_tool_result_captured_in_provider_specific_fields():
and parsed.choices[0].delta.provider_specific_fields
and "web_search_results" in parsed.choices[0].delta.provider_specific_fields
):
web_search_results = parsed.choices[0].delta.provider_specific_fields[
"web_search_results"
]
web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"]
# Verify web_search_results was captured
assert web_search_results is not None, "web_search_results should be captured"
assert len(web_search_results) == 1, "Should have 1 web_search_tool_result block"
assert (
web_search_results[0]["type"] == "web_search_tool_result"
), "Block type should be web_search_tool_result"
assert (
web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123"
), "tool_use_id should match"
assert web_search_results[0]["type"] == "web_search_tool_result", "Block type should be web_search_tool_result"
assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match"
assert len(web_search_results[0]["content"]) == 2, "Should have 2 search results"
assert (
web_search_results[0]["content"][0]["title"] == "Fun Otter Facts"
), "First result title should match"
assert web_search_results[0]["content"][0]["title"] == "Fun Otter Facts", "First result title should match"
def test_web_fetch_tool_result_captured_in_provider_specific_fields():
@ -1447,9 +1317,7 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields():
The web_fetch_tool_result content comes ALL AT ONCE in content_block_start,
not in deltas, so we need to capture it there.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Simulate the streaming sequence with web_fetch_tool_result
chunks = [
@ -1520,25 +1388,15 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields():
and parsed.choices[0].delta.provider_specific_fields
and "web_search_results" in parsed.choices[0].delta.provider_specific_fields
):
web_search_results = parsed.choices[0].delta.provider_specific_fields[
"web_search_results"
]
web_search_results = parsed.choices[0].delta.provider_specific_fields["web_search_results"]
# Verify web_fetch_tool_result was captured (stored in web_search_results list)
assert web_search_results is not None, "web_search_results should be captured"
assert len(web_search_results) == 1, "Should have 1 web_fetch_tool_result block"
assert (
web_search_results[0]["type"] == "web_fetch_tool_result"
), "Block type should be web_fetch_tool_result"
assert (
web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123"
), "tool_use_id should match"
assert (
web_search_results[0]["content"]["url"] == "https://example.com"
), "URL should match"
assert (
web_search_results[0]["content"]["content"]["title"] == "Example Page"
), "Title should match"
assert web_search_results[0]["type"] == "web_fetch_tool_result", "Block type should be web_fetch_tool_result"
assert web_search_results[0]["tool_use_id"] == "srvtoolu_01ABC123", "tool_use_id should match"
assert web_search_results[0]["content"]["url"] == "https://example.com", "URL should match"
assert web_search_results[0]["content"]["content"]["title"] == "Example Page", "Title should match"
def test_web_fetch_tool_result_no_extra_tool_calls():
@ -1551,9 +1409,7 @@ def test_web_fetch_tool_result_no_extra_tool_calls():
The issue was that web_fetch_tool_result blocks have input_json_delta events with {}
that were incorrectly being converted to tool calls.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# to verify it doesn't emit tool calls
chunks = [
@ -1597,9 +1453,9 @@ def test_web_fetch_tool_result_no_extra_tool_calls():
tool_call_count += 1
# Should have 0 tool calls - web_fetch_tool_result should not emit tool calls
assert (
tool_call_count == 0
), f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls"
assert tool_call_count == 0, (
f"Expected 0 tool calls, got {tool_call_count}. web_fetch_tool_result should not emit tool calls"
)
def test_container_in_provider_specific_fields_streaming():
@ -1609,9 +1465,7 @@ def test_container_in_provider_specific_fields_streaming():
When container with skills is used, the container field should be present in
the provider_specific_fields of the message_delta chunk.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
# Simulate streaming chunks
chunks = [
@ -1679,20 +1533,12 @@ def test_container_in_provider_specific_fields_streaming():
and parsed.choices[0].delta.provider_specific_fields
and "container" in parsed.choices[0].delta.provider_specific_fields
):
container_field = parsed.choices[0].delta.provider_specific_fields[
"container"
]
container_field = parsed.choices[0].delta.provider_specific_fields["container"]
# Verify container was captured
assert (
container_field is not None
), "container should be captured in provider_specific_fields"
assert (
container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p"
), "container id should match"
assert (
container_field["expires_at"] == "2025-12-16T04:57:16.913181Z"
), "expires_at should match"
assert container_field is not None, "container should be captured in provider_specific_fields"
assert container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p", "container id should match"
assert container_field["expires_at"] == "2025-12-16T04:57:16.913181Z", "expires_at should match"
assert len(container_field["skills"]) == 1, "Should have 1 skill"
assert container_field["skills"][0]["skill_id"] == "pptx", "skill_id should be pptx"
assert container_field["skills"][0]["version"] == "20251013", "version should match"
@ -1705,9 +1551,7 @@ def test_container_in_provider_specific_fields_non_streaming():
When container with skills is used in non-streaming, the container field should be
present in the provider_specific_fields of the response.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
# Simulate a message_delta chunk with container (as it would appear in non-streaming)
message_delta_chunk = {
@ -1743,21 +1587,13 @@ def test_container_in_provider_specific_fields_non_streaming():
# Verify container is in provider_specific_fields
assert model_response.choices[0].delta.provider_specific_fields is not None
assert "container" in model_response.choices[0].delta.provider_specific_fields
container_field = model_response.choices[0].delta.provider_specific_fields[
"container"
]
container_field = model_response.choices[0].delta.provider_specific_fields["container"]
assert container_field["id"] == "container_abc123xyz", "container id should match"
assert (
container_field["expires_at"] == "2025-12-20T10:30:00.000000Z"
), "expires_at should match"
assert container_field["expires_at"] == "2025-12-20T10:30:00.000000Z", "expires_at should match"
assert len(container_field["skills"]) == 2, "Should have 2 skills"
assert (
container_field["skills"][0]["skill_id"] == "code_execution"
), "First skill_id should be code_execution"
assert (
container_field["skills"][1]["skill_id"] == "pptx"
), "Second skill_id should be pptx"
assert container_field["skills"][0]["skill_id"] == "code_execution", "First skill_id should be code_execution"
assert container_field["skills"][1]["skill_id"] == "pptx", "Second skill_id should be pptx"
def test_container_absent_when_not_provided():
@ -1766,9 +1602,7 @@ def test_container_absent_when_not_provided():
This ensures we don't add empty or None container fields.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=False, json_mode=False
)
iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False)
# message_delta without container
message_delta_chunk = {
@ -1787,9 +1621,9 @@ def test_container_absent_when_not_provided():
# Verify container is NOT in provider_specific_fields when not provided
if model_response.choices[0].delta.provider_specific_fields:
assert (
"container" not in model_response.choices[0].delta.provider_specific_fields
), "container should not be present when not provided in delta"
assert "container" not in model_response.choices[0].delta.provider_specific_fields, (
"container should not be present when not provided in delta"
)
def test_streaming_code_execution_produces_code_interpreter_results():
@ -1985,8 +1819,7 @@ def test_streaming_multiple_code_executions_no_duplicates():
# Second (final) emission: cumulative list with BOTH results
# This is what stream_chunk_builder will pick as "last value wins"
assert len(emissions[1]) == 2, (
f"Expected final emission to have 2 results, got {len(emissions[1])}. "
f"IDs: {[r.id for r in emissions[1]]}"
f"Expected final emission to have 2 results, got {len(emissions[1])}. IDs: {[r.id for r in emissions[1]]}"
)
assert emissions[1][0].id == "srvtoolu_01AAA"
assert emissions[1][0].code == "echo first"
@ -2150,9 +1983,7 @@ def test_empty_output_produces_null_outputs():
assert code_results is not None, "No code_interpreter_results emitted"
assert len(code_results) == 1
assert code_results[0].id == "srvtoolu_01AAA"
assert (
code_results[0].outputs is None
), f"Expected outputs=None for empty execution, got {code_results[0].outputs}"
assert code_results[0].outputs is None, f"Expected outputs=None for empty execution, got {code_results[0].outputs}"
def test_non_bash_tool_result_skipped():
@ -2215,12 +2046,10 @@ def test_non_bash_tool_result_skipped():
code_results = psf["code_interpreter_results"]
# code_interpreter_results should be emitted but empty (no bash results)
assert (
code_results is not None
), "Expected code_interpreter_results key to be emitted"
assert (
len(code_results) == 0
), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"
assert code_results is not None, "Expected code_interpreter_results key to be emitted"
assert len(code_results) == 0, (
f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"
)
class TestRustChatCompletionsHook:
@ -2257,13 +2086,9 @@ class TestRustChatCompletionsHook:
from litellm.rust_bridge import chat_completions as bridge
monkeypatch.setenv("LITELLM_RUST", "1")
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
yield
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
@staticmethod
def _completion_kwargs(**overrides):
@ -2309,17 +2134,17 @@ class TestRustChatCompletionsHook:
seen["gate"].append(kwargs)
return decline_reason
def native(request, *, context):
def native(request, *, options, context):
seen["call"].append(
{
"model": request.model,
"messages": request.messages,
"optional_params": {**request.optional_params, **(request.options.provider_connection or {})},
"api_key": request.options.api_key,
"api_base": request.options.api_base,
"custom_llm_provider": request.options.custom_llm_provider,
"extra_headers": request.options.extra_headers,
"timeout_seconds": request.options.timeout_seconds,
"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,
}
)
if sync_error is not None:
@ -2372,9 +2197,7 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
seen = self._inject()
AnthropicChatCompletion().completion(
**self._completion_kwargs(optional_params={"max_tokens": 7})
)
AnthropicChatCompletion().completion(**self._completion_kwargs(optional_params={"max_tokens": 7}))
assert seen["call"][0]["optional_params"]["max_tokens"] == 7
def test_without_the_opt_in_the_core_is_never_consulted(self, monkeypatch):
@ -2383,15 +2206,14 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
) as transform, patch.object(
AnthropicChatCompletion, "acompletion_function"
with (
patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
) as transform,
patch.object(AnthropicChatCompletion, "acompletion_function"),
):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(litellm_params={})
)
AnthropicChatCompletion().completion(**self._completion_kwargs(litellm_params={}))
except Exception:
# The Python path goes on to make an HTTP call; reaching it is
# the assertion, so the network failure below is expected.
@ -2405,9 +2227,7 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject(decline_reason="unrecognized request parameter")
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(**self._completion_kwargs())
except Exception:
@ -2420,9 +2240,7 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(optional_params={"max_tokens": 16, "stream": True})
@ -2436,9 +2254,7 @@ class TestRustChatCompletionsHook:
seen = self._inject()
logging_obj = MagicMock()
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
assert logging_obj.pre_call.call_count == 1
assert len(seen["call"]) == 1
@ -2452,9 +2268,7 @@ class TestRustChatCompletionsHook:
self._inject()
logging_obj = MagicMock()
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
assert logging_obj.post_call.call_count == 1
logged = logging_obj.post_call.call_args.kwargs["original_response"]
@ -2475,22 +2289,16 @@ class TestRustChatCompletionsHook:
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
def declining_native(request, *, context):
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
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native)
logging_obj, calls = self._recording_logging_obj()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
except Exception:
# The Python path goes on to make an HTTP call; the log count is
# the assertion, so a failure past this point is expected.
@ -2512,24 +2320,18 @@ class TestRustChatCompletionsHook:
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
async def declining_native(request, *, context):
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
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native)
sentinel = object()
async def python_path(**_kwargs):
return sentinel
with patch.object(
AnthropicChatCompletion, "acompletion_function", side_effect=python_path
) as python_call:
result = await AnthropicChatCompletion().completion(
**self._completion_kwargs(acompletion=True)
)
with patch.object(AnthropicChatCompletion, "acompletion_function", side_effect=python_path) as python_call:
result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True))
assert result is sentinel
assert python_call.called, "a failing rust call must re-enter the python path"
@ -2539,23 +2341,18 @@ class TestRustChatCompletionsHook:
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.rust_bridge import chat_completions as bridge
async def native(request, *, context):
async def native(request, *, options, context):
return dict(self.RUST_RESPONSE)
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, achat_completions=native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native)
with patch.object(AnthropicChatCompletion, "acompletion_function") as python_call:
result = await AnthropicChatCompletion().completion(
**self._completion_kwargs(acompletion=True)
)
result = await AnthropicChatCompletion().completion(**self._completion_kwargs(acompletion=True))
assert result.choices[0].message.content == "hello from rust"
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
assert not python_call.called
def test_pre_call_logging_fires_once_when_the_sync_rust_call_declines(self, monkeypatch):
"""One request, one pre_call, on the synchronous path too. Without the
suppression the Python path logs a second time for the same attempt."""
@ -2572,30 +2369,22 @@ class TestRustChatCompletionsHook:
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
def declining_native(request, *, context):
def declining_native(request, *, options, context):
raise _Declined("blank message text")
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native)
logging_obj, calls = self._recording_logging_obj()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(logging_obj=logging_obj)
)
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
except Exception:
# The Python path goes on to make an HTTP call; the log count is
# the assertion, so a failure past this point is expected.
pass
assert len(calls["pre_call"]) == 1
assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == (
"claude-sonnet-4-5"
)
assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == ("claude-sonnet-4-5")
def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch):
"""The suppression must not swallow the log on the ordinary path."""
@ -2605,9 +2394,7 @@ class TestRustChatCompletionsHook:
self._inject()
logging_obj, calls = self._recording_logging_obj()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
):
with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}):
try:
AnthropicChatCompletion().completion(
**self._completion_kwargs(litellm_params={}, logging_obj=logging_obj)

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

@ -10,8 +10,8 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
from botocore.credentials import Credentials
from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.rust_bridge import chat_completions as bridge
@ -49,13 +49,9 @@ RESOLVED_CREDENTIALS = Credentials(
@pytest.fixture(autouse=True)
def reset_bridge(monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
yield
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
def _inject(*, decline_reason=None, error: Exception | None = None):
@ -65,17 +61,18 @@ def _inject(*, decline_reason=None, error: Exception | None = None):
seen["gate"].append(kwargs)
return decline_reason
def native(request, *, context):
def native(request, *, options, context):
seen["call"].append(
{
"model": request.model,
"messages": request.messages,
"optional_params": {**request.optional_params, **(request.options.provider_connection or {})},
"api_key": request.options.api_key,
"api_base": request.options.api_base,
"custom_llm_provider": request.options.custom_llm_provider,
"extra_headers": request.options.extra_headers,
"timeout_seconds": request.options.timeout_seconds,
"optional_params": request.optional_params,
"bedrock": options.bedrock,
"api_key": options.api_key,
"api_base": options.api_base,
"custom_llm_provider": options.custom_llm_provider,
"extra_headers": options.extra_headers,
"timeout_seconds": options.timeout_seconds,
}
)
if error is not None:
@ -137,20 +134,18 @@ def test_the_core_receives_the_credentials_this_handler_already_resolved():
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"
bedrock = seen["call"][0]["bedrock"]
assert bedrock.aws_access_key_id == "AKIARESOLVED"
assert bedrock.aws_secret_access_key == "resolved-secret"
assert bedrock.aws_session_token == "resolved-token"
assert bedrock.aws_region_name == "us-east-1"
def test_the_core_receives_the_converse_url_this_handler_already_built():
seen = _inject()
_run()
assert seen["call"][0]["api_base"].endswith(
"/model/anthropic.claude-sonnet-4-5-v1%3A0/converse"
)
assert seen["call"][0]["api_base"].endswith("/model/anthropic.claude-sonnet-4-5-v1%3A0/converse")
assert "bedrock-runtime.us-east-1.amazonaws.com" in seen["call"][0]["api_base"]
@ -218,12 +213,10 @@ async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch):
monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative())
async def declining_native(request, *, context):
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
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native)
sentinel = object()
@ -231,16 +224,10 @@ async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch):
return sentinel
with (
patch.object(
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
),
patch.object(
BedrockConverseLLM, "async_completion", side_effect=python_path
) as python_call,
patch.object(BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS),
patch.object(BedrockConverseLLM, "async_completion", side_effect=python_path) as python_call,
):
result = await BedrockConverseLLM().completion(
**_completion_kwargs(acompletion=True)
)
result = await BedrockConverseLLM().completion(**_completion_kwargs(acompletion=True))
assert result is sentinel
assert python_call.called, "a failing rust call must re-enter the python path"
@ -248,22 +235,16 @@ async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch):
@pytest.mark.asyncio
async def test_the_async_path_serves_the_rust_response_without_the_fallback():
async def native(request, *, context):
async def native(request, *, options, context):
return dict(RUST_RESPONSE)
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, achat_completions=native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native)
with (
patch.object(
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
),
patch.object(BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS),
patch.object(BedrockConverseLLM, "async_completion") as python_call,
):
result = await BedrockConverseLLM().completion(
**_completion_kwargs(acompletion=True)
)
result = await BedrockConverseLLM().completion(**_completion_kwargs(acompletion=True))
assert result.choices[0].message.content == "hello from rust"
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
@ -282,7 +263,7 @@ async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines():
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
async def declining_native(request, *, context):
async def declining_native(request, *, options, context):
raise _Declined("blank message text")
logging_obj = MagicMock()
@ -294,19 +275,11 @@ async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines():
with (
patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()),
patch.object(
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
),
patch.object(
BedrockConverseLLM, "async_completion", side_effect=python_path
),
patch.object(BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS),
patch.object(BedrockConverseLLM, "async_completion", side_effect=python_path),
):
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, achat_completions=declining_native
)
await BedrockConverseLLM().completion(
**_completion_kwargs(acompletion=True, logging_obj=logging_obj)
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=declining_native)
await BedrockConverseLLM().completion(**_completion_kwargs(acompletion=True, logging_obj=logging_obj))
assert logging_obj.pre_call.call_count == 1
assert served and served[0]["skip_pre_call_logging"] is True
@ -395,15 +368,13 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines():
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
def declining_native(request, *, context):
def declining_native(request, *, options, context):
raise _Declined("blank message text")
logging_obj = MagicMock()
with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()):
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native)
response = _run(
logging_obj=logging_obj,
client=_sync_client_returning_converse_response(),
@ -449,20 +420,14 @@ async def test_post_call_logging_fires_on_the_async_rust_path():
cannot drift apart the way the pre_call suppression once did."""
import json
async def native(request, *, context):
async def native(request, *, options, context):
return dict(RUST_RESPONSE)
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, achat_completions=native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, achat_completions=native)
logging_obj = MagicMock()
with patch.object(
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
):
await BedrockConverseLLM().completion(
**_completion_kwargs(acompletion=True, logging_obj=logging_obj)
)
with patch.object(BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS):
await BedrockConverseLLM().completion(**_completion_kwargs(acompletion=True, logging_obj=logging_obj))
assert logging_obj.post_call.call_count == 1
logged = logging_obj.post_call.call_args.kwargs["original_response"]
@ -481,15 +446,13 @@ def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines():
RustBridgeDeclined = _Declined
RustUpstreamError = type("_Upstream", (Exception,), {})
def declining_native(request, *, context):
def declining_native(request, *, options, context):
raise _Declined("blank message text")
logging_obj, calls = _recording_logging_obj()
with patch("litellm.rust_bridge.bindings.get_native_bridge", lambda: _FakeNative()):
bridge.set_rust_chat_completions(
decline=lambda **_kwargs: None, chat_completions=declining_native
)
bridge.set_rust_chat_completions(decline=lambda **_kwargs: None, chat_completions=declining_native)
response = _run(
logging_obj=logging_obj,
client=_sync_client_returning_converse_response(),
@ -523,9 +486,11 @@ def test_the_rust_opt_in_needs_no_sigv4_principal():
response = _run(credentials=None, api_key="bedrock-bearer-token")
assert response.choices[0].message.content == "hello from rust"
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"
bedrock = seen["call"][0]["bedrock"]
assert bedrock.aws_access_key_id is None
assert bedrock.aws_secret_access_key is None
assert bedrock.aws_session_token is None
assert bedrock.aws_region_name == "us-east-1"
assert seen["call"][0]["api_key"] == "bedrock-bearer-token"

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

@ -12,7 +12,13 @@ import pytest
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge import configuration
from litellm.rust_bridge.request import NativeOCRRequest, NativeRequestContext, NativeRequestOptions, PreparedNativeCall
from litellm.rust_bridge.request import (
NativeOCRRequest,
NativeRequestContext,
NativeRequestOptions,
NativeVertexOptions,
PreparedNativeCall,
)
from litellm.rust_bridge.timeouts import timeout_to_seconds
# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr`
@ -51,18 +57,20 @@ class RecordingBridge:
self,
request: NativeOCRRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> dict[str, object]:
self.calls.append(
{
"model": request.model,
"document": request.document,
"api_key": request.options.api_key,
"api_base": request.options.api_base,
"custom_llm_provider": request.options.custom_llm_provider,
"extra_headers": request.options.extra_headers,
"optional_params": {**request.optional_params, **(request.options.provider_connection or {})},
"timeout_seconds": request.options.timeout_seconds,
"api_key": options.api_key,
"api_base": options.api_base,
"custom_llm_provider": options.custom_llm_provider,
"extra_headers": options.extra_headers,
"optional_params": request.optional_params,
"vertex": options.vertex,
"timeout_seconds": options.timeout_seconds,
}
)
return dict(FAKE_OCR_RESPONSE)
@ -78,18 +86,20 @@ class RecordingAsyncBridge:
self,
request: NativeOCRRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> dict[str, object]:
self.calls.append(
{
"model": request.model,
"document": request.document,
"api_key": request.options.api_key,
"api_base": request.options.api_base,
"custom_llm_provider": request.options.custom_llm_provider,
"extra_headers": request.options.extra_headers,
"optional_params": {**request.optional_params, **(request.options.provider_connection or {})},
"timeout_seconds": request.options.timeout_seconds,
"api_key": options.api_key,
"api_base": options.api_base,
"custom_llm_provider": options.custom_llm_provider,
"extra_headers": options.extra_headers,
"optional_params": request.optional_params,
"vertex": options.vertex,
"timeout_seconds": options.timeout_seconds,
}
)
return dict(FAKE_OCR_RESPONSE)
@ -100,6 +110,7 @@ class RaisingBridge:
self,
request: NativeOCRRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> dict[str, object]:
raise RuntimeError("bridge failed")
@ -110,6 +121,7 @@ class RaisingAsyncBridge:
self,
request: NativeOCRRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> dict[str, object]:
raise RuntimeError("bridge failed")
@ -381,13 +393,13 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response():
model="mistral-ocr-latest",
document=DOCUMENT,
optional_params={"include_image_base64": True, "pages": [0]},
options=NativeRequestOptions(
api_key="sk-test",
api_base="https://proxy.internal",
custom_llm_provider="mistral",
extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"},
timeout_seconds=12.5,
),
),
options=NativeRequestOptions(
api_key="sk-test",
api_base="https://proxy.internal",
custom_llm_provider="mistral",
extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"},
timeout_seconds=12.5,
),
),
fallback=lambda: pytest.fail("unexpected Python fallback"),
@ -410,6 +422,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response():
"x-trace-id": "trace-1",
},
"optional_params": {"include_image_base64": True, "pages": [0]},
"vertex": None,
"timeout_seconds": 12.5,
}
@ -431,11 +444,11 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response():
model="mistral-ocr-maas",
document=DOCUMENT,
optional_params={},
options=NativeRequestOptions(
custom_llm_provider="vertex_ai",
provider_connection={"vertex_project": "project-1"},
timeout_seconds=42.0,
),
),
options=NativeRequestOptions(
custom_llm_provider="vertex_ai",
vertex=NativeVertexOptions(project="project-1"),
timeout_seconds=42.0,
),
),
fallback=unexpected_fallback,
@ -453,7 +466,8 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response():
"api_base": None,
"custom_llm_provider": "vertex_ai",
"extra_headers": None,
"optional_params": {"vertex_project": "project-1"},
"optional_params": {},
"vertex": NativeVertexOptions(project="project-1"),
"timeout_seconds": 42.0,
}
@ -489,6 +503,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response():
"x-trace-id": "trace-1",
},
"optional_params": {"include_image_base64": True},
"vertex": NativeVertexOptions(),
"timeout_seconds": 12.5,
}
@ -573,11 +588,8 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
resolve_api_key=lambda _name: None,
)
assert bridge.calls[0]["optional_params"] == {
"include_image_base64": True,
"vertex_project": "project-1",
"vertex_location": "us-central1",
}
assert bridge.calls[0]["optional_params"] == {"include_image_base64": True}
assert bridge.calls[0]["vertex"] == NativeVertexOptions(project="project-1", location="us-central1")
def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager():
@ -601,8 +613,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana
resolve_api_key=_resolver,
)
assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret"
assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5"
assert bridge.calls[0]["vertex"] == NativeVertexOptions(project="project-from-secret", location="us-east5")
def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():

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

@ -33,6 +33,7 @@ class _FakeNativeBridge:
cls,
request: NativeResponsesWebSocketRequest,
*,
options: object,
context: NativeRequestContext,
) -> _FakeNativeConnection:
return _FakeNativeConnection()
@ -103,6 +104,7 @@ class _FailingNativeBridge:
cls,
request: NativeResponsesWebSocketRequest,
*,
options: object,
context: NativeRequestContext,
) -> _FakeNativeConnection:
raise RuntimeError("connection failed")

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

@ -183,34 +183,74 @@ def _record(name: str, fields: dict[str, object]) -> object:
def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
inputs: Final = _route_inputs(route, api_base, outcome)
params: Final = inputs.get("optional_params", {})
connection_fields: Final = frozenset(("aws_access_key_id", "aws_secret_access_key", "aws_region_name"))
options: Final = _record("RequestOptions", {
"api_key": inputs.get("api_key"),
"api_base": inputs.get("api_base"),
"custom_llm_provider": inputs.get("custom_llm_provider"),
"extra_headers": inputs.get("extra_headers"),
"extra_query": None,
"timeout_seconds": inputs.get("timeout_seconds"),
"provider_connection": {key: value for key, value in params.items() if key in connection_fields},
})
request: Final = _record("Request", {
**{key: value for key, value in inputs.items() if key in {"model", "document", "audio", "body", "messages"}},
**({"optional_params": {key: value for key, value in params.items() if key not in connection_fields}} if route != "messages" else {}),
"options": options,
})
context: Final = _record("RequestContext", {
"metadata": None,
"litellm_metadata": {"internal_marker": "must-not-reach-provider"},
"request_metadata_fields": (),
"litellm_call_id": "native-wheel-call",
"request_model": inputs["model"],
"attribution": _record("Attribution", {
"user_api_key_hash": None,
"user_api_key_user_id": "native-user",
"user_api_key_team_id": None,
}),
})
return {"request": request, "context": context}
bedrock: Final = _record(
"BedrockOptions",
{
"aws_access_key_id": params.get("aws_access_key_id"),
"aws_secret_access_key": params.get("aws_secret_access_key"),
"aws_session_token": None,
"aws_region_name": params.get("aws_region_name"),
"aws_session_name": None,
"aws_profile_name": None,
"aws_role_name": None,
"aws_web_identity_token": None,
"aws_sts_endpoint": None,
"aws_external_id": None,
"aws_bedrock_runtime_endpoint": None,
"request_metadata_fields": (),
"request_metadata": None,
},
)
options: Final = _record(
"RequestOptions",
{
"api_key": inputs.get("api_key"),
"api_base": inputs.get("api_base"),
"custom_llm_provider": inputs.get("custom_llm_provider"),
"extra_headers": inputs.get("extra_headers"),
"extra_query": None,
"timeout_seconds": inputs.get("timeout_seconds"),
"bedrock": bedrock,
"anthropic": None,
"vertex": None,
},
)
request_params: Final = {"language": params.get("language")} if route == "transcription" else params
request: Final = _record(
"Request",
{
**{
key: value for key, value in inputs.items() if key in {"model", "document", "audio", "body", "messages"}
},
**({"optional_params": request_params} if route != "messages" else {}),
},
)
context: Final = _record(
"RequestContext",
{
"litellm_call_id": "native-wheel-call",
"trace_id": None,
"request_model": inputs["model"],
"attribution": _record(
"Attribution",
{
"user_api_key_hash": None,
"user_api_key_user_id": "native-user",
"user_api_key_team_id": None,
},
),
"capabilities": _record(
"Capabilities",
{
"stream": False,
"has_agentic_hook": False,
"has_custom_client": False,
"request_format": None,
},
),
},
)
return {"request": request, "options": options, "context": context}
def assert_success(route: str, response: object) -> None:
@ -270,12 +310,7 @@ async def exercise_async(native: object, api_base: str) -> None:
async def exercise_async_concurrency(native: object, api_base: str) -> None:
responses: Final = await asyncio.wait_for(
asyncio.gather(
*(
native.amessages(**route_kwargs("messages", api_base, "success"))
for _ in range(32)
)
),
asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))),
timeout=15,
)
for response in responses:

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

@ -12,6 +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 (
NativeBedrockOptions,
NativeRequestCapabilities,
NativeRequestContext,
anthropic_options,
)
from litellm.types.utils import ModelResponse
RUST_RESPONSE = {
@ -95,16 +101,16 @@ class _RecordingCall:
self.error = error
self.calls: list[dict] = []
def __call__(self, request, *, context):
self.calls.append({"request": request, "context": 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, *, context):
return _RecordingCall.__call__(self, request, context=context)
async def __call__(self, request, *, options, context):
return _RecordingCall.__call__(self, request, options=options, context=context)
def _accepts(**overrides) -> bool:
@ -264,7 +270,7 @@ class TestSyncCall:
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert native.calls[0]["request"].options.timeout_seconds == 30.0
assert native.calls[0]["options"].timeout_seconds == 30.0
def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
@ -420,13 +426,48 @@ def test_provider_credentials_are_separate_from_chat_body_params():
kwargs = _call_kwargs(ModelResponse())
kwargs["optional_params"] = {
"max_tokens": 32,
"aws_access_key_id": "test-access-key",
"aws_secret_access_key": "test-secret-key",
}
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 request.options.provider_connection == {
"aws_access_key_id": "test-access-key",
"aws_secret_access_key": "test-secret-key",
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"

23
tests/test_litellm/test_audio_transcription_rust_bridge.py
View file

@ -4,7 +4,11 @@ import pytest
import litellm
from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch
from litellm.rust_bridge.request import NativeRequestContext, NativeTranscriptionRequest
from litellm.rust_bridge.request import (
NativeRequestContext,
NativeRequestOptions,
NativeTranscriptionRequest,
)
rust_bridge = importlib.import_module("litellm.rust_bridge.transcription")
@ -17,9 +21,17 @@ class SyncBridge:
self,
request: NativeTranscriptionRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> dict[str, object]:
self.calls.append({"model": request.model, "audio": request.audio, "optional_params": {**request.optional_params, **(request.options.provider_connection or {})}})
self.calls.append(
{
"model": request.model,
"audio": request.audio,
"optional_params": request.optional_params,
"bedrock": options.bedrock,
}
)
return {"text": "hello"}
@ -28,6 +40,7 @@ class AsyncBridge:
self,
request: NativeTranscriptionRequest,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> dict[str, object]:
return {"text": "async"}
@ -111,7 +124,7 @@ async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPat
def test_bedrock_transcription_uses_rust_only_path() -> None:
rust_bridge.configure_rust_transcription(
transcription=lambda request, *, context: {"text": "rust"},
transcription=lambda request, *, options, context: {"text": "rust"},
atranscription=None,
)
try:
@ -127,7 +140,9 @@ def test_bedrock_transcription_uses_rust_only_path() -> None:
@pytest.mark.asyncio
async def test_bedrock_atranscription_uses_rust_only_path() -> None:
async def rust_response(request: NativeTranscriptionRequest, *, context: NativeRequestContext) -> dict[str, object]:
async def rust_response(
request: NativeTranscriptionRequest, *, options: object, context: NativeRequestContext
) -> dict[str, object]:
return {"text": "rust"}
rust_bridge.configure_rust_transcription(transcription=None, atranscription=rust_response)