From 221ddb3ab62a297e9c3afdfaa8da6db91c6a80d9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 16:59:57 -0700 Subject: [PATCH] refactor(native): separate request options and context --- litellm-rust/README.md | 14 +- .../src/audio_transcription/hooks.rs | 28 +- .../ai-gateway/src/audio_transcription/mod.rs | 4 +- .../src/audio_transcription/prepare.rs | 27 +- .../src/audio_transcription/tests.rs | 32 +- .../src/audio_transcription/types.rs | 5 +- .../crates/ai-gateway/src/io/responses_ws.rs | 6 +- .../crates/ai-gateway/src/ocr/hooks.rs | 2 +- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 36 +- .../crates/ai-gateway/src/ocr/prepare.rs | 38 +- .../crates/ai-gateway/src/ocr/types.rs | 5 +- .../crates/ai-gateway/tests/ocr_lifecycle.rs | 124 +++-- .../core/src/audio_transcription/handler.rs | 2 +- .../core/src/audio_transcription/mod.rs | 9 +- .../core/src/audio_transcription/prepare.rs | 52 +- .../core/src/audio_transcription/tests.rs | 31 +- .../core/src/audio_transcription/types.rs | 6 +- .../core/src/chat_completions/handler.rs | 11 +- .../crates/core/src/chat_completions/mod.rs | 5 +- .../core/src/chat_completions/prepare.rs | 26 +- .../crates/core/src/chat_completions/tests.rs | 91 ++-- .../crates/core/src/chat_completions/types.rs | 5 +- .../crates/core/src/messages/handler.rs | 6 +- litellm-rust/crates/core/src/messages/mod.rs | 7 +- .../crates/core/src/messages/prepare.rs | 43 +- .../crates/core/src/messages/tests.rs | 128 ++--- .../crates/core/src/messages/types.rs | 2 - .../crates/core/src/request_context.rs | 15 +- .../crates/core/src/request_options.rs | 71 ++- .../crates/core/src/responses/types.rs | 1 - litellm-rust/crates/python-bridge/src/lib.rs | 27 +- .../crates/python-bridge/src/marshal.rs | 138 ++++- .../src/routes/audio_transcription.rs | 4 +- .../src/routes/chat_completions.rs | 4 +- .../python-bridge/src/routes/definition.rs | 65 ++- .../python-bridge/src/routes/messages.rs | 4 +- .../crates/python-bridge/src/routes/ocr.rs | 4 +- litellm/llms/anthropic/chat/handler.py | 197 +++---- litellm/llms/bedrock/chat/converse_handler.py | 315 ++++++------ litellm/llms/bedrock/request_metadata.py | 6 +- litellm/ocr/main.py | 25 +- litellm/rust_bridge/chat_completions.py | 67 +-- litellm/rust_bridge/messages.py | 28 +- litellm/rust_bridge/protocols.py | 2 + litellm/rust_bridge/request.py | 131 +++-- litellm/rust_bridge/responses_websocket.py | 2 +- litellm/rust_bridge/transcription.py | 39 +- .../strategies/trace_parity/sdk/execution.py | 49 +- .../test_rust_bridge_messages.py | 30 +- .../chat/test_anthropic_chat_handler.py | 485 +++++------------- .../chat/test_bedrock_converse_handler.py | 123 ++--- tests/test_litellm/ocr/test_rust_bridge.py | 77 +-- .../responses/test_rust_bridge_websocket.py | 2 + .../rust_bridge/native_route_wheel_test.py | 103 ++-- .../rust_bridge/test_chat_completions.py | 61 ++- .../test_audio_transcription_rust_bridge.py | 23 +- 56 files changed, 1508 insertions(+), 1335 deletions(-) diff --git a/litellm-rust/README.md b/litellm-rust/README.md index 5720a463455..c76b33ef5a2 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -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 diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index ee4a99ca094..3b1201b1666 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -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 } diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs index de385cb667d..642e2c2b168 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs @@ -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 { @@ -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, diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs index cb4884cceb6..f551b632793 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs @@ -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), diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs index d54399b9f96..ea7c6fb0074 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs @@ -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(), diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs index 8656b839c09..aee24cf4bb0 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs @@ -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, - 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, + pub(crate) bedrock: BedrockOptions, pub(crate) api_key: Option, pub(crate) api_base: Option, pub(crate) extra_headers: Option>, diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 9d22ec76e5e..8409cdcbaec 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -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 { - 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()) })?, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 0a018b19a0e..ff9b47b4ccc 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -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(), diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index 520fd88c8a5..328b42be352 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -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 { @@ -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() }, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index cd0d2899bc7..1fe884489ad 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -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() }, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index d6be6ba0d78..9be628a2c81 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -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, - 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, + pub(crate) vertex: VertexOptions, pub(crate) api_key: Option, pub(crate) api_base: Option, pub(crate) extra_headers: Option>, diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index abbd1424525..f68ac32333a 100644 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -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(), diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index e9d91d7cd3b..ca29c40d4f5 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -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?; diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index c7331a6ada0..798be28803a 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -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 { - 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)] diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 28d975bd178..3c382315864 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -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 { - 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, }) } diff --git a/litellm-rust/crates/core/src/audio_transcription/tests.rs b/litellm-rust/crates/core/src/audio_transcription/tests.rs index 224141776c2..98cb93da5df 100644 --- a/litellm-rust/crates/core/src/audio_transcription/tests.rs +++ b/litellm-rust/crates/core/src/audio_transcription/tests.rs @@ -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() diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index f988202cd14..5ecec7df642 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -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, - 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, + pub(super) bedrock: BedrockOptions, pub(super) timeout: Option, } diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index efa7351abcb..ee35d26443d 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -101,15 +101,10 @@ pub(super) async fn signed_headers( let unsigned: BTreeMap = 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, diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 8cb9ffe3246..b6139a57b70 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -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 { - 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 diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index f4e74b8563b..d22a66ed6c0 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -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, Error> pub(super) fn resolve_request( request: ChatCompletionsRequest<'_>, + options: RequestOptions, _context: &LiteLlmRequestContext, ) -> Result { - 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, }) } diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index d25d45b9de0..a02904db716 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -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 { 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 { + 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"}]), diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index b096df4664d..e9c459b6108 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -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, - 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, + pub(super) bedrock: BedrockOptions, pub(super) timeout: Option, } diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 61ff81bcdc8..1897d550b39 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -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 { - 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 { - 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(), diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 10bcb7ab451..61360bb17be 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -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 { - 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 { - execute_messages_provider_stream(request).await + execute_messages_provider_stream(request, options.clone()).await } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 27ea0f14f87..2a75dfc30cf 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -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 { - 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, }) } diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index fd7d8f8c453..9c74bbb7034 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -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() diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index 7ad0a36cc3e..608e3f77fff 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -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 { diff --git a/litellm-rust/crates/core/src/request_context.rs b/litellm-rust/crates/core/src/request_context.rs index ea8e6096a4f..2796edf9993 100644 --- a/litellm-rust/crates/core/src/request_context.rs +++ b/litellm-rust/crates/core/src/request_context.rs @@ -1,5 +1,3 @@ -use serde_json::{Map, Value}; - #[derive(Clone, Debug, Default, PartialEq)] pub struct RequestAttribution { pub user_api_key_hash: Option, @@ -7,12 +5,19 @@ pub struct RequestAttribution { pub user_api_key_team_id: Option, } +#[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, +} + #[derive(Clone, Debug, Default, PartialEq)] pub struct LiteLlmRequestContext { - pub metadata: Option>, - pub litellm_metadata: Option>, - pub request_metadata_fields: Vec, pub litellm_call_id: Option, + pub trace_id: Option, pub request_model: Option, pub attribution: RequestAttribution, + pub capabilities: RequestCapabilities, } diff --git a/litellm-rust/crates/core/src/request_options.rs b/litellm-rust/crates/core/src/request_options.rs index 9ab0a81c998..66c7ed89366 100644 --- a/litellm-rust/crates/core/src/request_options.rs +++ b/litellm-rust/crates/core/src/request_options.rs @@ -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, + pub aws_secret_access_key: Option, + pub aws_session_token: Option, + pub aws_region_name: Option, + pub aws_session_name: Option, + pub aws_profile_name: Option, + pub aws_role_name: Option, + pub aws_web_identity_token: Option, + pub aws_sts_endpoint: Option, + pub aws_external_id: Option, + pub aws_bedrock_runtime_endpoint: Option, + pub request_metadata_fields: Vec, + pub request_metadata: Option>, +} + +impl BedrockOptions { + pub fn into_map(&self) -> Map { + [ + ("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, +} + +#[derive(Clone, Debug, Default)] +pub struct VertexOptions { + pub project: Option, + pub location: Option, +} + +impl VertexOptions { + pub fn into_map(&self) -> Map { + [ + ("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, @@ -10,5 +77,7 @@ pub struct RequestOptions { pub extra_headers: Option>, pub extra_query: Option>, pub timeout: Option, - pub provider_connection: Map, + pub bedrock: Option, + pub anthropic: Option, + pub vertex: Option, } diff --git a/litellm-rust/crates/core/src/responses/types.rs b/litellm-rust/crates/core/src/responses/types.rs index 009310e3734..65a6951df0f 100644 --- a/litellm-rust/crates/core/src/responses/types.rs +++ b/litellm-rust/crates/core/src/responses/types.rs @@ -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)] diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 0ff434b4107..5dbbb100ec9 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -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> { - 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" diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 8e4543b7997..1b6b37f04de 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -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, + aws_secret_access_key: Option, + aws_session_token: Option, + aws_region_name: Option, + aws_session_name: Option, + aws_profile_name: Option, + aws_role_name: Option, + aws_web_identity_token: Option, + aws_sts_endpoint: Option, + aws_external_id: Option, + aws_bedrock_runtime_endpoint: Option, + request_metadata_fields: Vec, + request_metadata: Option>, +} + +impl From 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, +} + +impl From for litellm_core::request_options::AnthropicOptions { + fn from(input: NativeAnthropicOptions) -> Self { + Self { + user_id: input.user_id, + } + } +} + +#[derive(FromPyObject)] +struct NativeVertexOptions { + project: Option, + location: Option, +} + +impl From 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, @@ -14,8 +79,9 @@ pub(crate) struct NativeRequestOptions { #[pyo3(from_py_with = litellm_python_interop::from_py)] extra_query: Option>, timeout_seconds: Option, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - provider_connection: Option>, + bedrock: Option, + anthropic: Option, + vertex: Option, } impl From for litellm_core::request_options::RequestOptions { @@ -27,7 +93,9 @@ impl From 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>, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - litellm_metadata: Option>, - request_metadata_fields: Vec, litellm_call_id: Option, + trace_id: Option, request_model: Option, attribution: NativeRequestAttribution, + capabilities: NativeRequestCapabilities, +} + +#[derive(FromPyObject)] +struct NativeRequestCapabilities { + stream: bool, + has_agentic_hook: bool, + has_custom_client: bool, + request_format: Option, } impl From 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), diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 0bb93bbf132..18d351b3791 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -15,11 +15,11 @@ struct AudioTranscriptionInputs { audio: Value, #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Map, - options: NativeRequestOptions, } fn prepare_transcription( input: AudioTranscriptionInputs, + options: NativeRequestOptions, context: NativeRequestContext, ) -> PyResult> + 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 diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 96dad4c6328..aad6c3c737f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -16,11 +16,11 @@ struct ChatCompletionsInputs { messages: Value, #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Map, - options: NativeRequestOptions, } fn prepare_chat_completions( input: ChatCompletionsInputs, + options: NativeRequestOptions, context: NativeRequestContext, ) -> PyResult> + 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 diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 3957fb9976f..cf38fef9d55 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -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> { - 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> { - 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> { - 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> { - 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> + 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}], diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index 2ecb3690c15..2ee126b9807 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -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> + 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 diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index f4a3c8ddc6e..fceb4b1b865 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -16,11 +16,11 @@ struct OcrInputs { document: Value, #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Map, - options: NativeRequestOptions, } fn prepare_ocr( input: OcrInputs, + options: NativeRequestOptions, context: NativeRequestContext, ) -> PyResult> + 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(), diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index bcc2235e11e..cc9430f6bf4 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -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 diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 8d6cf23f47e..19ffb695af1 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -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 diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py index 1f4e5886508..599afec42e8 100644 --- a/litellm/llms/bedrock/request_metadata.py +++ b/litellm/llms/bedrock/request_metadata.py @@ -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( diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 9c77fc39909..bebb4b44589 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -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), + ), ) diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index 8cb5c80f2fd..a46a1f541bd 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -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(), ), diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index d9ea26333bd..f41cff140f8 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -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(), ), diff --git a/litellm/rust_bridge/protocols.py b/litellm/rust_bridge/protocols.py index 76b76504a6c..dc779b346ab 100644 --- a/litellm/rust_bridge/protocols.py +++ b/litellm/rust_bridge/protocols.py @@ -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: ... diff --git a/litellm/rust_bridge/request.py b/litellm/rust_bridge/request.py index ba650f22372..4239b9763a1 100644 --- a/litellm/rust_bridge/request.py +++ b/litellm/rust_bridge/request.py @@ -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 diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index a48405a7884..b7098aaddb4 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -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), diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 7c89ae449c9..0300cd84883 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -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(), ), diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index e3a032c8d94..5bd0358eeb2 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -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( diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 5f83e9c9950..59787822ab6 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -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") diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 6751cf7f06f..019022ed919 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -23,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) diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index c372b295d64..c614e040d4c 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -10,8 +10,8 @@ from unittest.mock import MagicMock, patch import httpx import pytest - from botocore.credentials import Credentials + from 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" diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 8720e73e77d..2a9b314ac85 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -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(): diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index cbf183dea91..aea5137db66 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -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") diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 5bc754bb724..a24d5ba1ae9 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -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: diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 4b07099cc78..f81b7050134 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -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" diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index 72c60463fc5..57c936c170b 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -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)