From 533870bb0d24cecc86729ed017857256c6f8d738 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 23:59:03 +0000 Subject: [PATCH 01/12] feat(messages): route Azure Anthropic /messages through Rust behind rust:true Adds an opt-in Rust path for non-streaming Azure Anthropic Messages. A deployment sets rust: true in litellm_params to route litellm.messages() and the proxy /v1/messages endpoint through the native Rust bridge; a missing flag or rust: false keeps the existing Python path, and non-Azure providers, streaming, an unavailable bridge, or a None result all fall back to Python. Rust-backed responses carry an x-litellm-rust: true response header so callers can see which path served the request. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- basedpyright-code-budget.json | 2 +- .../crates/ai-gateway/src/constants.rs | 12 + .../crates/ai-gateway/src/io/messages.rs | 1 + litellm-rust/crates/ai-gateway/src/io/mod.rs | 1 + litellm-rust/crates/ai-gateway/src/lib.rs | 1 + .../crates/ai-gateway/src/messages/client.rs | 15 + .../ai-gateway/src/messages/common_utils.rs | 50 +++ .../crates/ai-gateway/src/messages/handler.rs | 46 +++ .../crates/ai-gateway/src/messages/mod.rs | 21 ++ .../crates/ai-gateway/src/messages/prepare.rs | 64 ++++ .../crates/ai-gateway/src/messages/tests.rs | 259 +++++++++++++ .../crates/ai-gateway/src/messages/types.rs | 23 ++ litellm-rust/crates/core/src/lib.rs | 1 + litellm-rust/crates/core/src/messages/mod.rs | 2 + .../core/src/messages/transformation.rs | 72 ++++ .../crates/core/src/messages/types.rs | 18 + .../src/providers/azure_ai/messages/mod.rs | 1 + .../azure_ai/messages/transformation.rs | 345 ++++++++++++++++++ .../crates/core/src/providers/azure_ai/mod.rs | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 89 +++++ litellm/llms/custom_httpx/llm_http_handler.py | 87 +++++ litellm/rust_bridge/messages.py | 153 ++++++++ litellm/rust_bridge/ocr.py | 43 ++- litellm/rust_bridge/timeouts.py | 15 + tests/test_litellm/rust_bridge/__init__.py | 0 .../test_litellm/rust_bridge/test_messages.py | 274 ++++++++++++++ type-discipline-budget.json | 2 +- 27 files changed, 1579 insertions(+), 19 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/src/io/messages.rs create mode 100644 litellm-rust/crates/ai-gateway/src/messages/client.rs create mode 100644 litellm-rust/crates/ai-gateway/src/messages/common_utils.rs create mode 100644 litellm-rust/crates/ai-gateway/src/messages/handler.rs create mode 100644 litellm-rust/crates/ai-gateway/src/messages/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/messages/prepare.rs create mode 100644 litellm-rust/crates/ai-gateway/src/messages/tests.rs create mode 100644 litellm-rust/crates/ai-gateway/src/messages/types.rs create mode 100644 litellm-rust/crates/core/src/messages/mod.rs create mode 100644 litellm-rust/crates/core/src/messages/transformation.rs create mode 100644 litellm-rust/crates/core/src/messages/types.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/messages/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs create mode 100644 litellm/rust_bridge/messages.py create mode 100644 litellm/rust_bridge/timeouts.py create mode 100644 tests/test_litellm/rust_bridge/__init__.py create mode 100644 tests/test_litellm/rust_bridge/test_messages.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index edfb3536ad37..c05f98d22ae8 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 10397 + "limit": 10393 }, "reportFunctionMemberAccess": { "limit": 11 diff --git a/litellm-rust/crates/ai-gateway/src/constants.rs b/litellm-rust/crates/ai-gateway/src/constants.rs index 109b648f5db7..557fe5d53d44 100644 --- a/litellm-rust/crates/ai-gateway/src/constants.rs +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -28,3 +28,15 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500; /// Provider attributed to realtime sessions in the logging payload. #[cfg(feature = "server")] pub(crate) const DEFAULT_PROVIDER: &str = "openai"; + +/// Full-request timeout ceiling for Anthropic Messages provider calls, in +/// seconds. Mirrors the Python Anthropic Messages default. The per-request +/// timeout from `litellm_params` still overrides this on the request builder. +pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600; + +/// Connect timeout for Anthropic Messages provider calls, in seconds. +pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10; + +/// Max characters of an upstream error body echoed across the host boundary +/// before truncation, so provider bodies are bounded and data-minimized. +pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256; diff --git a/litellm-rust/crates/ai-gateway/src/io/messages.rs b/litellm-rust/crates/ai-gateway/src/io/messages.rs new file mode 100644 index 000000000000..b784d2b62a19 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/messages.rs @@ -0,0 +1 @@ +pub use crate::messages::{messages, MessagesRequest}; diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index 3b566027646a..7bc9642d1924 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -1,3 +1,4 @@ +pub mod messages; pub mod ocr; pub mod realtime; pub mod realtime_pool; diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index d8ef7bb5ba17..db4c8211a5a0 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -12,6 +12,7 @@ //! for the load-time config reader. pub mod io; +pub mod messages; pub mod ocr; /// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and diff --git a/litellm-rust/crates/ai-gateway/src/messages/client.rs b/litellm-rust/crates/ai-gateway/src/messages/client.rs new file mode 100644 index 000000000000..6281270b9640 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/messages/client.rs @@ -0,0 +1,15 @@ +use std::sync::OnceLock; +use std::time::Duration; + +use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS}; + +pub(super) fn http_client() -> &'static reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::Client::builder() + .timeout(Duration::from_secs(MESSAGES_TIMEOUT_SECS)) + .connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS)) + .build() + .unwrap_or_else(|_| reqwest::Client::new()) + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs new file mode 100644 index 000000000000..fe4ac4cf26f0 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs @@ -0,0 +1,50 @@ +use litellm_core::error::{json_type_name, CoreError}; +use litellm_core::messages::transformation::AnthropicMessagesProviderConfig; +use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; +use litellm_core::CoreResult; +use serde_json::{Map, Value}; + +use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS; + +pub(super) fn truncate_error_body(body: &str) -> String { + if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS { + return body.to_string(); + } + let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect(); + format!("{truncated}... (truncated)") +} + +pub(super) fn messages_provider_config( + provider: &str, +) -> Option<&'static dyn AnthropicMessagesProviderConfig> { + match provider { + "azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG), + _ => None, + } +} + +pub(super) fn string_headers( + extra_headers: Option>, +) -> CoreResult> { + extra_headers + .unwrap_or_default() + .into_iter() + .map(|(key, value)| { + value + .as_str() + .map(|value| (key.clone(), value.to_string())) + .ok_or_else(|| { + CoreError::InvalidRequest(format!( + "messages extra_headers.{key} must be a string, got {}", + json_type_name(&value) + )) + }) + }) + .collect() +} + +pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { + headers + .iter() + .any(|(key, _)| key.eq_ignore_ascii_case(name)) +} diff --git a/litellm-rust/crates/ai-gateway/src/messages/handler.rs b/litellm-rust/crates/ai-gateway/src/messages/handler.rs new file mode 100644 index 000000000000..ad58d74a87f6 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/messages/handler.rs @@ -0,0 +1,46 @@ +use litellm_core::error::CoreError; +use litellm_core::CoreResult; +use serde_json::Value; + +use super::client::http_client; +use super::common_utils::truncate_error_body; +use super::types::ProviderMessagesRequest; + +pub(super) async fn execute_messages_provider_call( + request: ProviderMessagesRequest, +) -> CoreResult { + 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); + } + if let Some(duration) = request.timeout { + request_builder = request_builder.timeout(duration); + } + + let response = request_builder + .send() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + + let status = response.status(); + let text = response + .text() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + + if !status.is_success() { + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + + let response_json: Value = serde_json::from_str(&text).map_err(|err| { + CoreError::InvalidResponse(format!("invalid messages response JSON: {err}")) + })?; + + Ok(request + .config + .transform_response(&request.model, response_json)? + .into_json()) +} diff --git a/litellm-rust/crates/ai-gateway/src/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/messages/mod.rs new file mode 100644 index 000000000000..7ed81474c47a --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/messages/mod.rs @@ -0,0 +1,21 @@ +use litellm_core::CoreResult; +use serde_json::Value; + +mod client; +mod common_utils; +mod handler; +mod prepare; +mod types; + +pub use types::MessagesRequest; + +use handler::execute_messages_provider_call; +use prepare::prepare_messages_call; + +pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { + let prepared = prepare_messages_call(request)?; + execute_messages_provider_call(prepared).await +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs new file mode 100644 index 000000000000..72e321972679 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs @@ -0,0 +1,64 @@ +use litellm_core::messages::transformation::MessagesAuthStrategy; +use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; +use litellm_core::CoreError; +use litellm_core::CoreResult; + +use super::common_utils::{has_header, messages_provider_config, string_headers}; +use super::types::{MessagesRequest, ProviderMessagesRequest}; + +pub(super) fn prepare_messages_call( + request: MessagesRequest<'_>, +) -> CoreResult { + let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) + .or_else(|| { + request + .custom_llm_provider + .map(|provider| CustomLlmProvider { + model: request.model, + custom_llm_provider: provider, + }) + }) + .ok_or_else(|| { + CoreError::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; + + let config = messages_provider_config(provider) + .ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?; + let env_lookup = |key: &str| std::env::var(key).ok(); + + let mut headers = string_headers(request.extra_headers)?; + + let auth_strategy = config.auth_strategy(); + if !has_header(&headers, auth_strategy.header_name()) { + let api_key = config.resolve_api_key(request.api_key, &env_lookup)?; + let auth_header = match auth_strategy { + MessagesAuthStrategy::Bearer => { + ("authorization".to_string(), format!("Bearer {api_key}")) + } + MessagesAuthStrategy::Header(name) => (name.to_string(), api_key), + }; + headers.push(auth_header); + } + + for (name, value) in config.default_headers() { + if !has_header(&headers, name) { + headers.push((name.to_string(), value.to_string())); + } + } + + let url = config.complete_url(request.api_base, &model, &env_lookup)?; + let body = config.transform_request(request.body)?.body; + + Ok(ProviderMessagesRequest { + model, + config, + url, + body, + upstream_headers: headers, + timeout: request.timeout, + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/messages/tests.rs b/litellm-rust/crates/ai-gateway/src/messages/tests.rs new file mode 100644 index 000000000000..30f6642400ea --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/messages/tests.rs @@ -0,0 +1,259 @@ +use std::time::Duration; + +use litellm_core::error::CoreError; +use serde_json::{json, Map, Value}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; + +use super::common_utils::{ + has_header, messages_provider_config, string_headers, truncate_error_body, +}; +use super::{messages, MessagesRequest}; + +async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") +} + +fn write_response(body: &str) -> String { + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ) +} + +#[test] +fn provider_config_only_resolves_azure_ai() { + assert!(messages_provider_config("azure_ai").is_some()); + assert!(messages_provider_config("anthropic").is_none()); + assert!(messages_provider_config("openai").is_none()); +} + +#[test] +fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(400); + let truncated = truncate_error_body(&body); + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, 256); +} + +#[test] +fn string_headers_rejects_non_string_values() { + let headers = json!({"x-count": 3}).as_object().unwrap().clone(); + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert!(matches!(err, CoreError::InvalidRequest(_))); +} + +#[test] +fn has_header_is_case_insensitive() { + let headers = vec![("X-Api-Key".to_string(), "secret".to_string())]; + assert!(has_header(&headers, "x-api-key")); + assert!(!has_header(&headers, "authorization")); +} + +#[tokio::test] +async fn messages_round_trip_builds_azure_request_and_passes_response_through() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#; + socket + .write_all(write_response(response_body).as_bytes()) + .await + .expect("writes response"); + request + }); + + let response = messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 1024, + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": "hi", + "cache_control": {"type": "ephemeral", "scope": "global"} + }] + }] + }), + api_key: Some("sk-azure"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("messages request succeeds"); + + assert_eq!(response["content"][0]["text"], "hi"); + assert_eq!(response["stop_reason"], "end_turn"); + + let request = server.await.expect("server task completes"); + let (head, body) = request.split_once("\r\n\r\n").expect("has body"); + assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}"); + let head_lower = head.to_ascii_lowercase(); + assert!(head_lower.contains("x-api-key: sk-azure"), "{head}"); + assert!( + head_lower.contains("anthropic-version: 2023-06-01"), + "{head}" + ); + assert!( + head_lower.contains("content-type: application/json"), + "{head}" + ); + + let sent_body: Value = serde_json::from_str(body).expect("body is json"); + assert_eq!( + sent_body["messages"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); +} + +#[tokio::test] +async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = + r#"{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"m"}"#; + socket + .write_all(write_response(response_body).as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "x-api-key".to_string(), + Value::String("from-python".to_string()), + ); + headers.insert( + "anthropic-beta".to_string(), + Value::String("token-efficient-tools-2025-02-19".to_string()), + ); + + messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: Some("rust-fallback-key"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: Some(headers), + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("messages request succeeds"); + + let request = server.await.expect("server task completes"); + let head = request + .split_once("\r\n\r\n") + .expect("has body") + .0 + .to_ascii_lowercase(); + let api_key_count = head + .lines() + .filter(|line| line.starts_with("x-api-key:")) + .count(); + assert_eq!(api_key_count, 1, "{head}"); + assert!(head.contains("x-api-key: from-python"), "{head}"); + assert!( + head.contains("anthropic-beta: token-efficient-tools-2025-02-19"), + "{head}" + ); + assert!(!head.contains("rust-fallback-key"), "{head}"); +} + +#[tokio::test] +async fn messages_maps_provider_error_status_to_http_error() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + let body = "unauthorized"; + let response = format!( + "HTTP/1.1 401 Unauthorized\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + }); + + let err = messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: Some("sk-azure"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect_err("provider error propagates"); + + assert!(matches!(err, CoreError::Http { status: 401, .. })); +} + +#[tokio::test] +async fn messages_rejects_unsupported_provider() { + let err = messages(MessagesRequest { + model: "claude-3-5-sonnet", + body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}), + api_key: Some("sk"), + api_base: Some("http://127.0.0.1:1"), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_millis(50)), + }) + .await + .expect_err("unsupported provider errors"); + + assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "anthropic")); +} diff --git a/litellm-rust/crates/ai-gateway/src/messages/types.rs b/litellm-rust/crates/ai-gateway/src/messages/types.rs new file mode 100644 index 000000000000..6840ff57cc43 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/messages/types.rs @@ -0,0 +1,23 @@ +use std::time::Duration; + +use litellm_core::messages::transformation::AnthropicMessagesProviderConfig; +use serde_json::{Map, Value}; + +pub struct MessagesRequest<'a> { + pub model: &'a str, + pub body: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: Option<&'a str>, + pub extra_headers: Option>, + pub timeout: Option, +} + +pub(crate) struct ProviderMessagesRequest { + pub(crate) model: String, + pub(crate) config: &'static dyn AnthropicMessagesProviderConfig, + pub(crate) url: String, + pub(crate) body: Value, + pub(crate) upstream_headers: Vec<(String, String)>, + pub(crate) timeout: Option, +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 555a04ce8533..27154f5a08b6 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,5 +1,6 @@ pub mod call_lifecycle; pub mod error; +pub mod messages; pub mod ocr; pub mod providers; pub mod realtime; diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs new file mode 100644 index 000000000000..ec2fbb969a62 --- /dev/null +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -0,0 +1,2 @@ +pub mod transformation; +pub mod types; diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs new file mode 100644 index 000000000000..448f7dd6dc83 --- /dev/null +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -0,0 +1,72 @@ +use serde_json::Value; + +use crate::error::{json_type_name, CoreError, CoreResult}; + +use super::types::{MessagesRequestData, MessagesResponseData}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MessagesAuthStrategy { + Bearer, + Header(&'static str), +} + +impl MessagesAuthStrategy { + pub fn header_name(self) -> &'static str { + match self { + Self::Bearer => "authorization", + Self::Header(header_name) => header_name, + } + } +} + +pub trait AnthropicMessagesProviderConfig: Sync { + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn auth_strategy(&self) -> MessagesAuthStrategy { + MessagesAuthStrategy::Header("x-api-key") + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + &[ + ("anthropic-version", "2023-06-01"), + ("content-type", "application/json"), + ] + } + + fn transform_request(&self, body: Value) -> CoreResult { + if !body.is_object() { + return Err(CoreError::InvalidType { + expected: "object", + actual: json_type_name(&body), + }); + } + Ok(MessagesRequestData { body }) + } + + fn transform_response( + &self, + _model: &str, + response_json: Value, + ) -> CoreResult { + if !response_json.is_object() { + return Err(CoreError::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + }); + } + Ok(MessagesResponseData { + body: response_json, + }) + } +} diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs new file mode 100644 index 000000000000..e56e924e9104 --- /dev/null +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -0,0 +1,18 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct MessagesRequestData { + pub body: Value, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct MessagesResponseData { + pub body: Value, +} + +impl MessagesResponseData { + pub fn into_json(self) -> Value { + self.body + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/mod.rs new file mode 100644 index 000000000000..f239b6921fac --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs new file mode 100644 index 000000000000..4ac1e521d455 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -0,0 +1,345 @@ +use serde_json::{Map, Value}; + +use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; +use crate::messages::types::MessagesRequestData; + +const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY"; +const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; + +pub struct AzureAnthropicMessagesConfig; + +pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig = + AzureAnthropicMessagesConfig; + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +pub fn resolve_azure_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + non_empty(api_key) + .map(str::to_string) + .or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| { + CoreError::Auth( + "Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable" + .to_string(), + ) + }) +} + +pub fn complete_azure_anthropic_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let api_base = non_empty(api_base) + .map(str::to_string) + .or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| { + CoreError::Auth( + "Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \ + Expected format: https://.services.ai.azure.com/anthropic" + .to_string(), + ) + })?; + + let api_base = api_base.trim_end_matches('/'); + + if api_base.ends_with("/v1/messages") || api_base.ends_with("/anthropic/v1/messages") { + return Ok(api_base.to_string()); + } + + let with_anthropic = match api_base.split_once("/anthropic") { + Some((prefix, _)) => format!("{prefix}/anthropic"), + None => format!("{api_base}/anthropic"), + }; + Ok(format!("{with_anthropic}/v1/messages")) +} + +fn remove_scope_from_content_blocks(content: &mut [Value]) { + for item in content.iter_mut() { + if let Some(cache_control) = item + .as_object_mut() + .and_then(|block| block.get_mut("cache_control")) + .and_then(Value::as_object_mut) + { + cache_control.remove("scope"); + } + } +} + +fn remove_scope_from_cache_control(body: &mut Map) { + if let Some(Value::Array(system)) = body.get_mut("system") { + remove_scope_from_content_blocks(system); + } + if let Some(Value::Array(messages)) = body.get_mut("messages") { + for message in messages.iter_mut() { + if let Some(Value::Array(content)) = message + .as_object_mut() + .and_then(|message| message.get_mut("content")) + { + remove_scope_from_content_blocks(content); + } + } + } +} + +impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_azure_anthropic_url(api_base, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_azure_api_key(api_key, env_lookup) + } + + fn auth_strategy(&self) -> MessagesAuthStrategy { + MessagesAuthStrategy::Header("x-api-key") + } + + fn transform_request(&self, body: Value) -> CoreResult { + let mut body = match body { + Value::Object(body) => body, + other => { + return Err(CoreError::InvalidType { + expected: "object", + actual: json_type_name(&other), + }) + } + }; + remove_scope_from_cache_control(&mut body); + Ok(MessagesRequestData { + body: Value::Object(body), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn url_appends_anthropic_and_messages_suffix() { + let url = + complete_azure_anthropic_url(Some("https://resource.services.ai.azure.com"), &|_| None) + .expect("url builds"); + assert_eq!( + url, + "https://resource.services.ai.azure.com/anthropic/v1/messages" + ); + } + + #[test] + fn url_keeps_existing_anthropic_segment() { + let url = complete_azure_anthropic_url( + Some("https://resource.services.ai.azure.com/anthropic"), + &|_| None, + ) + .expect("url builds"); + assert_eq!( + url, + "https://resource.services.ai.azure.com/anthropic/v1/messages" + ); + } + + #[test] + fn url_leaves_complete_messages_endpoint_untouched() { + for base in [ + "https://resource.services.ai.azure.com/anthropic/v1/messages", + "https://resource.services.ai.azure.com/v1/messages", + ] { + assert_eq!( + complete_azure_anthropic_url(Some(base), &|_| None).expect("url builds"), + base + ); + } + } + + #[test] + fn url_trims_trailing_slash_and_truncates_after_anthropic() { + let url = complete_azure_anthropic_url( + Some("https://resource.services.ai.azure.com/anthropic/extra/"), + &|_| None, + ) + .expect("url builds"); + assert_eq!( + url, + "https://resource.services.ai.azure.com/anthropic/v1/messages" + ); + } + + #[test] + fn url_falls_back_to_env_then_errors_when_absent() { + let with_env = |key: &str| { + (key == AZURE_API_BASE_ENV).then(|| "https://env.services.ai.azure.com".to_string()) + }; + assert_eq!( + complete_azure_anthropic_url(None, &with_env).expect("url builds"), + "https://env.services.ai.azure.com/anthropic/v1/messages" + ); + let err = complete_azure_anthropic_url(Some(" "), &|_| None).expect_err("missing base"); + assert!(matches!(err, CoreError::Auth(_))); + } + + #[test] + fn resolve_api_key_prefers_param_then_env() { + assert_eq!( + resolve_azure_api_key(Some("sk-param"), &|_| None).unwrap(), + "sk-param" + ); + let with_env = |key: &str| (key == AZURE_API_KEY_ENV).then(|| "sk-env".to_string()); + assert_eq!( + resolve_azure_api_key(Some(" "), &with_env).unwrap(), + "sk-env" + ); + assert!(matches!( + resolve_azure_api_key(None, &|_| None).expect_err("missing key"), + CoreError::Auth(_) + )); + } + + #[test] + fn auth_strategy_is_x_api_key() { + assert_eq!( + AZURE_ANTHROPIC_MESSAGES_CONFIG + .auth_strategy() + .header_name(), + "x-api-key" + ); + } + + #[test] + fn default_headers_match_python() { + assert_eq!( + AZURE_ANTHROPIC_MESSAGES_CONFIG.default_headers(), + &[ + ("anthropic-version", "2023-06-01"), + ("content-type", "application/json"), + ] + ); + } + + #[test] + fn transform_request_strips_scope_from_system_and_messages() { + let body = json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 1024, + "system": [ + { + "type": "text", + "text": "sys", + "cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"} + } + ], + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "hi", + "cache_control": {"type": "ephemeral", "scope": "global"} + }, + {"type": "text", "text": "no cache control"} + ] + } + ] + }); + + let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(body) + .expect("request transforms") + .body; + + assert_eq!( + transformed["system"][0]["cache_control"], + json!({"type": "ephemeral", "ttl": "1h"}) + ); + assert_eq!( + transformed["messages"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); + assert_eq!( + transformed["messages"][0]["content"][1], + json!({"type": "text", "text": "no cache control"}) + ); + } + + #[test] + fn transform_request_is_idempotent_and_preserves_string_system() { + let body = json!({ + "model": "claude-sonnet-4-5", + "system": "plain string system", + "messages": [{"role": "user", "content": "hi"}] + }); + let once = AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(body) + .expect("request transforms") + .body; + let twice = AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(once.clone()) + .expect("request transforms") + .body; + assert_eq!(once, twice); + assert_eq!(once["system"], json!("plain string system")); + } + + #[test] + fn transform_request_rejects_non_object_body() { + let err = AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(json!("bad")) + .expect_err("non-object body should error"); + assert_eq!( + err, + CoreError::InvalidType { + expected: "object", + actual: "string", + } + ); + } + + #[test] + fn transform_response_passes_through_object() { + let response = json!({ + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hello"}], + "model": "claude-sonnet-4-5", + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 2} + }); + let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_response("claude-sonnet-4-5", response.clone()) + .expect("response transforms") + .into_json(); + assert_eq!(transformed, response); + } + + #[test] + fn transform_response_rejects_non_object() { + let err = AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_response("claude-sonnet-4-5", json!([1, 2, 3])) + .expect_err("array response should error"); + assert_eq!( + err, + CoreError::InvalidType { + expected: "object", + actual: "array", + } + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs index 3621ff6a2fdb..5d13fa93e00f 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs @@ -1 +1,2 @@ +pub mod messages; pub mod ocr; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 946a99f990c1..80662b66d707 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,5 +1,6 @@ use std::time::Duration; +use litellm_ai_gateway::io::messages::{messages as run_messages, MessagesRequest}; use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest}; use litellm_core::error::CoreError; use pyo3::exceptions::{PyRuntimeError, PyValueError}; @@ -171,6 +172,92 @@ fn aocr( }) } +type MarshaledMessagesInputs = (Value, Option>, Option); + +fn marshal_messages_inputs( + py: Python<'_>, + body: Py, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult { + let body = py_to_json(py, body.bind(py))?; + if !body.is_object() { + return Err(PyValueError::new_err("body must be a dict")); + } + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + Ok((body, extra_headers, optional_timeout(timeout_seconds))) +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn messages( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (body, extra_headers, timeout) = + marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; + + let result = gil::release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest { + model: &model, + body, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + })) + }); + + match result { + Ok(value) => json_to_py(py, value), + Err(err) => Err(core_error_to_pyerr(err)), + } +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn amessages( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (body, extra_headers, timeout) = + marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let value = run_messages(MessagesRequest { + model: &model, + body, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + }) + .await + .map_err(core_error_to_pyerr)?; + + Python::with_gil(|py| json_to_py(py, value)) + }) +} + #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); @@ -182,6 +269,8 @@ fn gil_stats(py: Python<'_>) -> PyResult> { fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(ocr, module)?)?; module.add_function(wrap_pyfunction!(aocr, module)?)?; + module.add_function(wrap_pyfunction!(messages, module)?)?; + module.add_function(wrap_pyfunction!(amessages, module)?)?; module.add_function(wrap_pyfunction!(gil_stats, module)?)?; Ok(()) } diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b47fc50e1965..e58641a5d8da 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2091,6 +2091,34 @@ async def async_anthropic_messages_handler( }, ) + rust_messages_response = await self._maybe_rust_anthropic_messages( + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + stream=stream or False, + model=model, + api_key=api_key, + api_base=api_base, + headers=headers, + request_body=request_body, + timeout=self._resolve_anthropic_messages_timeout( + litellm_params=litellm_params, + stream=stream or False, + custom_llm_provider=custom_llm_provider, + ), + ) + if rust_messages_response is not None: + return await self._finalize_anthropic_messages_response( + initial_response=rust_messages_response, + model=model, + messages=messages, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + api_key=api_key, + kwargs=kwargs, + ) + response = await self._async_post_anthropic_messages_with_http_error_retry( async_httpx_client=async_httpx_client, request_url=request_url, @@ -2165,6 +2193,31 @@ async def async_anthropic_messages_handler( logging_obj=logging_obj, ) + return await self._finalize_anthropic_messages_response( + initial_response=initial_response, + model=model, + messages=messages, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + api_key=api_key, + kwargs=kwargs, + ) + + async def _finalize_anthropic_messages_response( + self, + *, + initial_response: AnthropicMessagesResponse, + model: str, + messages: list[dict], + anthropic_messages_provider_config: BaseAnthropicMessagesConfig, + anthropic_messages_optional_request_params: dict, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str, + api_key: str | None, + kwargs: dict, + ) -> AnthropicMessagesResponse | AsyncIterator: # Inject api_key into kwargs so follow-up calls in agentic hooks can # authenticate. api_key is a named param here (not in kwargs), so # _prepare_followup_kwargs would miss it otherwise. @@ -2188,6 +2241,40 @@ async def async_anthropic_messages_handler( "anthropic_messages", ) + @staticmethod + async def _maybe_rust_anthropic_messages( + *, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + stream: bool, + model: str, + api_key: str | None, + api_base: str | None, + headers: dict, + request_body: dict, + timeout: float | httpx.Timeout | None, + ) -> AnthropicMessagesResponse | None: + if stream or custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True: + return None + + from litellm.rust_bridge import messages as rust_messages_bridge + + rust_response = await rust_messages_bridge.amessages( + model=model, + body=request_body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + ) + if rust_response is None: + return None + + response_obj = cast(AnthropicMessagesResponse, dict(rust_response)) + response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}} + return response_obj + def anthropic_messages_handler( self, model: str, diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py new file mode 100644 index 000000000000..d51689f5f3b6 --- /dev/null +++ b/litellm/rust_bridge/messages.py @@ -0,0 +1,153 @@ +"""Thin Python wrapper for the native Rust Anthropic Messages bridge.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Awaitable, Final, Protocol, Union, cast + +import httpx + +from litellm.rust_bridge.timeouts import timeout_to_seconds + + +class RustMessages(Protocol): + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> dict[str, object]: + raise NotImplementedError + + +class RustAmessages(Protocol): + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> Awaitable[dict[str, object]]: + raise NotImplementedError + + +class _Unset: + pass + + +_UNSET: Final[_Unset] = _Unset() + + +def _env_enables_rust_messages() -> bool: + return os.getenv("LITELLM_USE_RUST_MESSAGES", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +@dataclass(slots=True) +class _RustMessagesState: + enabled: bool + messages: RustMessages | None = None + amessages: RustAmessages | None = None + + +_STATE: Final[_RustMessagesState] = _RustMessagesState(enabled=_env_enables_rust_messages()) + + +def set_rust_messages( + enabled: bool | _Unset = _UNSET, + *, + messages: RustMessages | None | _Unset = _UNSET, + amessages: RustAmessages | None | _Unset = _UNSET, +) -> None: + if not isinstance(enabled, _Unset): + _STATE.enabled = enabled + if not isinstance(messages, _Unset): + _STATE.messages = messages + if not isinstance(amessages, _Unset): + _STATE.amessages = amessages + + +def rust_messages_enabled() -> bool: + return _STATE.enabled + + +def load_rust_messages() -> RustMessages | None: + if _STATE.messages is not None: + return _STATE.messages + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + if native_bridge is None: + return None + return cast(RustMessages, getattr(native_bridge, "messages", None)) + + +def load_rust_amessages() -> RustAmessages | None: + if _STATE.amessages is not None: + return _STATE.amessages + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + if native_bridge is None: + return None + return cast(RustAmessages, getattr(native_bridge, "amessages", None)) + + +def messages( + *, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_messages = load_rust_messages() + if rust_messages is None: + return None + return rust_messages( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ) + + +async def amessages( + *, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_amessages = load_rust_amessages() + if rust_amessages is None: + return None + return await rust_amessages( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 36a088b6b1a5..3654b9fb232b 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -3,10 +3,15 @@ from __future__ import annotations import os -from typing import Any, Awaitable, Final, Protocol, Union, cast +from typing import TYPE_CHECKING, Awaitable, Final, Protocol, Union, cast import httpx +from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds + +if TYPE_CHECKING: + from litellm.rust_bridge.messages import RustAmessages, RustMessages + class RustOcr(Protocol): def __call__( @@ -64,6 +69,8 @@ def use_litellm_rust( *, ocr: RustOcr | None | _Unset = _UNSET, aocr: RustAocr | None | _Unset = _UNSET, + messages: RustMessages | None | _Unset = _UNSET, + amessages: RustAmessages | None | _Unset = _UNSET, ) -> None: global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl _rust_ocr_enabled = enabled @@ -71,6 +78,16 @@ def use_litellm_rust( _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): _rust_aocr_impl = aocr + if isinstance(messages, _Unset) and isinstance(amessages, _Unset): + return + from litellm.rust_bridge.messages import set_rust_messages + + if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset): + set_rust_messages(enabled, messages=messages, amessages=amessages) + elif not isinstance(messages, _Unset): + set_rust_messages(enabled, messages=messages) + else: + set_rust_messages(enabled, amessages=amessages) def rust_ocr_enabled() -> bool: @@ -99,22 +116,14 @@ def load_rust_aocr() -> RustAocr | None: return cast(RustAocr, getattr(native_bridge, "aocr", None)) -def _timeout_to_seconds(timeout: Union[float, httpx.Timeout] | None) -> float | None: - if timeout is None: - return None - if isinstance(timeout, httpx.Timeout): - return timeout.read - return float(timeout) - - def ocr( *, model: str, - document: dict[str, Any], + document: dict[str, object], api_key: str | None, api_base: str | None, custom_llm_provider: str | None, - extra_headers: dict[str, Any] | None, + extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: Union[float, httpx.Timeout] | None, ) -> dict[str, object] | None: @@ -123,11 +132,11 @@ def ocr( return None return rust_ocr( model=model, - document=cast(dict[str, object], document), + document=document, api_key=api_key, api_base=api_base, custom_llm_provider=custom_llm_provider, - extra_headers=cast(dict[str, object] | None, extra_headers), + extra_headers=extra_headers, optional_params=optional_params, timeout_seconds=_timeout_to_seconds(timeout), ) @@ -136,11 +145,11 @@ def ocr( async def aocr( *, model: str, - document: dict[str, Any], + document: dict[str, object], api_key: str | None, api_base: str | None, custom_llm_provider: str | None, - extra_headers: dict[str, Any] | None, + extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: Union[float, httpx.Timeout] | None, ) -> dict[str, object] | None: @@ -149,11 +158,11 @@ async def aocr( return None return await rust_aocr( model=model, - document=cast(dict[str, object], document), + document=document, api_key=api_key, api_base=api_base, custom_llm_provider=custom_llm_provider, - extra_headers=cast(dict[str, object] | None, extra_headers), + extra_headers=extra_headers, optional_params=optional_params, timeout_seconds=_timeout_to_seconds(timeout), ) diff --git a/litellm/rust_bridge/timeouts.py b/litellm/rust_bridge/timeouts.py new file mode 100644 index 000000000000..4407986c3da2 --- /dev/null +++ b/litellm/rust_bridge/timeouts.py @@ -0,0 +1,15 @@ +"""Shared timeout conversion for the native Rust bridges.""" + +from __future__ import annotations + +from typing import Union + +import httpx + + +def timeout_to_seconds(timeout: Union[float, httpx.Timeout] | None) -> float | None: + if timeout is None: + return None + if isinstance(timeout, httpx.Timeout): + return timeout.read + return float(timeout) diff --git a/tests/test_litellm/rust_bridge/__init__.py b/tests/test_litellm/rust_bridge/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_litellm/rust_bridge/test_messages.py b/tests/test_litellm/rust_bridge/test_messages.py new file mode 100644 index 000000000000..6f0e1e190a19 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_messages.py @@ -0,0 +1,274 @@ +"""Tests for the optional Rust-backed Anthropic Messages path.""" + +import importlib + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.router import GenericLiteLLMParams + +rust_messages = importlib.import_module("litellm.rust_bridge.messages") +rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") + +FAKE_MESSAGES_RESPONSE: dict[str, object] = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "hello world"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 3}, +} + +REQUEST_BODY: dict[str, object] = { + "model": "claude-sonnet-4-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], +} + + +class RecordingMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append( + { + "model": model, + "body": body, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_MESSAGES_RESPONSE) + + +class RecordingAsyncMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append( + { + "model": model, + "body": body, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_MESSAGES_RESPONSE) + + +class ExplodingAsyncMessages: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> dict[str, object]: + self.calls += 1 + raise AssertionError("bridge must not be called") + + +@pytest.fixture(autouse=True) +def _reset_rust_flag(): + litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL + yield + litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL + + +def test_load_rust_messages_returns_injected_impl(): + bridge = RecordingMessages() + litellm.use_litellm_rust(True, messages=bridge) + assert rust_messages.load_rust_messages() is bridge + + +def test_load_rust_amessages_returns_injected_impl(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + assert rust_messages.load_rust_amessages() is bridge + + +def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) + litellm.use_litellm_rust(True) + assert rust_messages.load_rust_messages() is None + result = rust_messages.messages( + model="claude", + body=REQUEST_BODY, + api_key="k", + api_base="b", + custom_llm_provider="azure_ai", + extra_headers={}, + timeout=30.0, + ) + assert result is None + + +def test_messages_wrapper_forwards_args_and_converts_timeout(): + bridge = RecordingMessages() + litellm.use_litellm_rust(True, messages=bridge) + + response = rust_messages.messages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + custom_llm_provider="azure_ai", + extra_headers={"anthropic-beta": "token-efficient-tools-2025-02-19"}, + timeout=httpx.Timeout(600.0, read=42.0), + ) + + assert response == FAKE_MESSAGES_RESPONSE + assert bridge.calls[0] == { + "model": "claude-sonnet-4-5", + "body": REQUEST_BODY, + "api_key": "sk-azure", + "api_base": "https://resource.services.ai.azure.com/anthropic", + "custom_llm_provider": "azure_ai", + "extra_headers": {"anthropic-beta": "token-efficient-tools-2025-02-19"}, + "timeout_seconds": 42.0, + } + + +@pytest.mark.asyncio +async def test_amessages_wrapper_forwards_args(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await rust_messages.amessages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + custom_llm_provider="azure_ai", + extra_headers=None, + timeout=12.5, + ) + + assert response == FAKE_MESSAGES_RESPONSE + assert bridge.calls[0]["model"] == "claude-sonnet-4-5" + assert bridge.calls[0]["timeout_seconds"] == 12.5 + + +def _gate(**overrides): + kwargs = { + "custom_llm_provider": "azure_ai", + "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), + "stream": False, + "model": "claude-sonnet-4-5", + "api_key": "sk-azure", + "api_base": "https://resource.services.ai.azure.com/anthropic", + "headers": {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"}, + "request_body": dict(REQUEST_BODY), + "timeout": 30.0, + } + kwargs.update(overrides) + return BaseLLMHTTPHandler._maybe_rust_anthropic_messages(**kwargs) + + +@pytest.mark.asyncio +async def test_gate_invokes_rust_and_marks_response_header(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate() + + assert response is not None + assert response["id"] == "msg_123" + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + call = bridge.calls[0] + assert call["model"] == "claude-sonnet-4-5" + assert call["body"] == REQUEST_BODY + assert call["api_key"] == "sk-azure" + assert call["api_base"] == "https://resource.services.ai.azure.com/anthropic" + assert call["extra_headers"] == {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"} + assert call["timeout_seconds"] == 30.0 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_when_flag_absent(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_when_flag_false(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_for_non_azure_provider(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate(custom_llm_provider="anthropic") + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_when_streaming(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate(stream=True) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_gate_falls_back_when_bridge_unavailable(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) + litellm.use_litellm_rust(True) + + response = await _gate() + + assert response is None diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 87b4c96e3231..8fe7c588cf1b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1113 + "limit": 1112 }, "LIT007": { "limit": 0 From 194a950170f361d107d4c2cd39c5e78f2e287065 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:07:26 +0000 Subject: [PATCH 02/12] test(docs): exclude LITELLM_USE_RUST_MESSAGES rollout flag from env-doc check Mirrors the existing LITELLM_USE_RUST_OCR entry; the flag is an internal rollout toggle that is intentionally not in the public environment settings docs yet. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- tests/documentation_tests/test_env_keys.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 60fbd505d674..1086d67b2795 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -28,6 +28,7 @@ # environment settings docs until the feature is ready for broad use. EXCLUDED_ROLLOUT_FLAGS = { "LITELLM_USE_RUST_OCR", + "LITELLM_USE_RUST_MESSAGES", } EXCLUDED_TERMINAL_VARS = { From 95feb89aea00012379569e9365729654e215fb07 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:22:01 +0000 Subject: [PATCH 03/12] fix(rust_bridge): isolate OCR enable flag and drop dead messages global toggle use_litellm_rust only mutates the OCR enabled flag when configuring OCR (or called with no bridge kwargs, preserving the legacy contract), so configuring only the messages bridge no longer flips OCR state. Remove the vestigial global enabled/env state from the messages bridge. Routing is controlled per deployment by rust:true in the shared handler gate, so the messages module never consulted the global toggle; drop it rather than leave a no-op switch. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- basedpyright-code-budget.json | 2 +- litellm/rust_bridge/messages.py | 20 +----------------- litellm/rust_bridge/ocr.py | 13 +++++++----- tests/documentation_tests/test_env_keys.py | 1 - .../test_litellm/rust_bridge/test_messages.py | 21 +++++++++++++++++++ type-discipline-budget.json | 2 +- 6 files changed, 32 insertions(+), 27 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c05f98d22ae8..75d4d13eb714 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 10393 + "limit": 10389 }, "reportFunctionMemberAccess": { "limit": 11 diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index d51689f5f3b6..5abb21879d3e 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -2,7 +2,6 @@ from __future__ import annotations -import os from dataclasses import dataclass from typing import Awaitable, Final, Protocol, Union, cast @@ -46,43 +45,26 @@ class _Unset: _UNSET: Final[_Unset] = _Unset() -def _env_enables_rust_messages() -> bool: - return os.getenv("LITELLM_USE_RUST_MESSAGES", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - - @dataclass(slots=True) class _RustMessagesState: - enabled: bool messages: RustMessages | None = None amessages: RustAmessages | None = None -_STATE: Final[_RustMessagesState] = _RustMessagesState(enabled=_env_enables_rust_messages()) +_STATE: Final[_RustMessagesState] = _RustMessagesState() def set_rust_messages( - enabled: bool | _Unset = _UNSET, *, messages: RustMessages | None | _Unset = _UNSET, amessages: RustAmessages | None | _Unset = _UNSET, ) -> None: - if not isinstance(enabled, _Unset): - _STATE.enabled = enabled if not isinstance(messages, _Unset): _STATE.messages = messages if not isinstance(amessages, _Unset): _STATE.amessages = amessages -def rust_messages_enabled() -> bool: - return _STATE.enabled - - def load_rust_messages() -> RustMessages | None: if _STATE.messages is not None: return _STATE.messages diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 3654b9fb232b..35de2eb9727a 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -73,21 +73,24 @@ def use_litellm_rust( amessages: RustAmessages | None | _Unset = _UNSET, ) -> None: global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl - _rust_ocr_enabled = enabled + configuring_ocr = not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset) + configuring_messages = not isinstance(messages, _Unset) or not isinstance(amessages, _Unset) + if configuring_ocr or not configuring_messages: + _rust_ocr_enabled = enabled if not isinstance(ocr, _Unset): _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): _rust_aocr_impl = aocr - if isinstance(messages, _Unset) and isinstance(amessages, _Unset): + if not configuring_messages: return from litellm.rust_bridge.messages import set_rust_messages if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset): - set_rust_messages(enabled, messages=messages, amessages=amessages) + set_rust_messages(messages=messages, amessages=amessages) elif not isinstance(messages, _Unset): - set_rust_messages(enabled, messages=messages) + set_rust_messages(messages=messages) else: - set_rust_messages(enabled, amessages=amessages) + set_rust_messages(amessages=amessages) def rust_ocr_enabled() -> bool: diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 1086d67b2795..60fbd505d674 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -28,7 +28,6 @@ # environment settings docs until the feature is ready for broad use. EXCLUDED_ROLLOUT_FLAGS = { "LITELLM_USE_RUST_OCR", - "LITELLM_USE_RUST_MESSAGES", } EXCLUDED_TERMINAL_VARS = { diff --git a/tests/test_litellm/rust_bridge/test_messages.py b/tests/test_litellm/rust_bridge/test_messages.py index 6f0e1e190a19..f5be34ae0b67 100644 --- a/tests/test_litellm/rust_bridge/test_messages.py +++ b/tests/test_litellm/rust_bridge/test_messages.py @@ -109,6 +109,27 @@ def test_load_rust_messages_returns_injected_impl(): assert rust_messages.load_rust_messages() is bridge +def test_configuring_messages_does_not_enable_ocr(): + from litellm.rust_bridge.ocr import rust_ocr_enabled + + litellm.use_litellm_rust(False) + assert rust_ocr_enabled() is False + + litellm.use_litellm_rust(True, messages=RecordingMessages()) + + assert rust_ocr_enabled() is False + + +def test_bare_use_litellm_rust_still_toggles_ocr(): + from litellm.rust_bridge.ocr import rust_ocr_enabled + + litellm.use_litellm_rust(True) + assert rust_ocr_enabled() is True + + litellm.use_litellm_rust(False) + assert rust_ocr_enabled() is False + + def test_load_rust_amessages_returns_injected_impl(): bridge = RecordingAsyncMessages() litellm.use_litellm_rust(True, amessages=bridge) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 8fe7c588cf1b..0482f47e5bcd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1112 + "limit": 1111 }, "LIT007": { "limit": 0 From 403dd14580f6b96b48fd2d433019da5b2d82bcb2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:13:05 +0000 Subject: [PATCH 04/12] refactor(rust/messages): split Anthropic config into its own provider file and type the request/response contract Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .../crates/ai-gateway/src/messages/handler.rs | 11 +- .../crates/ai-gateway/src/messages/prepare.rs | 10 +- .../core/src/messages/transformation.rs | 33 +-- .../crates/core/src/messages/types.rs | 110 ++++++++- .../src/providers/anthropic/messages/mod.rs | 1 + .../anthropic/messages/transformation.rs | 142 +++++++++++ .../core/src/providers/anthropic/mod.rs | 1 + .../azure_ai/messages/transformation.rs | 221 ++++++++++-------- litellm-rust/crates/core/src/providers/mod.rs | 1 + 9 files changed, 397 insertions(+), 133 deletions(-) create mode 100644 litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs create mode 100644 litellm-rust/crates/core/src/providers/anthropic/mod.rs diff --git a/litellm-rust/crates/ai-gateway/src/messages/handler.rs b/litellm-rust/crates/ai-gateway/src/messages/handler.rs index ad58d74a87f6..dd4a2f22aa79 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/handler.rs @@ -35,12 +35,13 @@ pub(super) async fn execute_messages_provider_call( }); } - let response_json: Value = serde_json::from_str(&text).map_err(|err| { + let response = serde_json::from_str(&text).map_err(|err| { CoreError::InvalidResponse(format!("invalid messages response JSON: {err}")) })?; - - Ok(request + let transformed = request .config - .transform_response(&request.model, response_json)? - .into_json()) + .transform_response(&request.model, response)?; + serde_json::to_value(transformed).map_err(|err| { + CoreError::InvalidResponse(format!("failed to serialize messages response: {err}")) + }) } diff --git a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs index 72e321972679..47105b399541 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs @@ -51,7 +51,15 @@ pub(super) fn prepare_messages_call( } let url = config.complete_url(request.api_base, &model, &env_lookup)?; - let body = config.transform_request(request.body)?.body; + let typed_request = serde_json::from_value(request.body).map_err(|err| { + CoreError::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + })?; + let transformed = config.transform_request(typed_request)?; + let body = serde_json::to_value(transformed).map_err(|err| { + CoreError::InvalidRequest(format!( + "failed to serialize Anthropic messages request: {err}" + )) + })?; Ok(ProviderMessagesRequest { model, diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index 448f7dd6dc83..3a34a58de6f9 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -1,8 +1,6 @@ -use serde_json::Value; +use crate::error::CoreResult; -use crate::error::{json_type_name, CoreError, CoreResult}; - -use super::types::{MessagesRequestData, MessagesResponseData}; +use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MessagesAuthStrategy { @@ -44,29 +42,18 @@ pub trait AnthropicMessagesProviderConfig: Sync { ] } - fn transform_request(&self, body: Value) -> CoreResult { - if !body.is_object() { - return Err(CoreError::InvalidType { - expected: "object", - actual: json_type_name(&body), - }); - } - Ok(MessagesRequestData { body }) + fn transform_request( + &self, + request: AnthropicMessagesRequest, + ) -> CoreResult { + Ok(request) } fn transform_response( &self, _model: &str, - response_json: Value, - ) -> CoreResult { - if !response_json.is_object() { - return Err(CoreError::InvalidType { - expected: "object", - actual: json_type_name(&response_json), - }); - } - Ok(MessagesResponseData { - body: response_json, - }) + response: AnthropicMessagesResponse, + ) -> CoreResult { + Ok(response) } } diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index e56e924e9104..11fe17ea40f2 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -1,18 +1,110 @@ use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Map, Value}; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct MessagesRequestData { - pub body: Value, +#[serde(untagged)] +pub enum SystemPrompt { + Text(String), + Blocks(Vec), } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct MessagesResponseData { - pub body: Value, +#[serde(untagged)] +pub enum MessageContent { + Text(String), + Blocks(Vec), } -impl MessagesResponseData { - pub fn into_json(self) -> Value { - self.body - } +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct ContentBlock { + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_control: Option, + #[serde(flatten)] + pub extra: Map, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct CacheControl { + #[serde(rename = "type", skip_serializing_if = "Option::is_none")] + pub cache_type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub ttl: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub scope: Option, + #[serde(flatten)] + pub extra: Map, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct AnthropicMessage { + pub role: String, + pub content: MessageContent, + #[serde(flatten)] + pub extra: Map, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct AnthropicMessagesRequest { + pub model: String, + pub messages: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub system: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stop_sequences: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub stream: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_k: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_choice: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub thinking: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub service_tier: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub container: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub mcp_servers: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub context_management: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_config: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub speed: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub inference_geo: Option, + #[serde(flatten)] + pub extra: Map, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct AnthropicMessagesResponse { + pub id: String, + #[serde(rename = "type")] + pub message_type: String, + pub role: String, + pub model: String, + pub content: Vec, + // Anthropic always includes stop_reason / stop_sequence, null until the turn + // ends; serialize them even when None so callers see the same shape as Python. + pub stop_reason: Option, + pub stop_sequence: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub usage: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub container: Option, + #[serde(flatten)] + pub extra: Map, } diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs new file mode 100644 index 000000000000..f239b6921fac --- /dev/null +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs new file mode 100644 index 000000000000..829f2260d3c1 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -0,0 +1,142 @@ +use crate::error::{CoreError, CoreResult}; +use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; + +const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; +const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; +const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com"; +const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; + +pub struct AnthropicMessagesConfig; + +pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig; + +pub fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +pub fn resolve_anthropic_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + non_empty(api_key) + .map(str::to_string) + .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| { + CoreError::Auth( + "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \ + environment variable" + .to_string(), + ) + }) +} + +pub fn complete_anthropic_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + let api_base = non_empty(api_base) + .map(str::to_string) + .or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty())) + .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string()); + + let api_base = api_base.trim_end_matches('/'); + if api_base.ends_with(MESSAGES_PATH_SUFFIX) { + return api_base.to_string(); + } + format!("{api_base}{MESSAGES_PATH_SUFFIX}") +} + +impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + Ok(complete_anthropic_url(api_base, env_lookup)) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_anthropic_api_key(api_key, env_lookup) + } + + fn auth_strategy(&self) -> MessagesAuthStrategy { + MessagesAuthStrategy::Header("x-api-key") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn url_defaults_to_public_anthropic_endpoint() { + assert_eq!( + complete_anthropic_url(None, &|_| None), + "https://api.anthropic.com/v1/messages" + ); + } + + #[test] + fn url_appends_messages_suffix_to_custom_base() { + assert_eq!( + complete_anthropic_url(Some("https://proxy.internal"), &|_| None), + "https://proxy.internal/v1/messages" + ); + } + + #[test] + fn url_leaves_complete_messages_endpoint_untouched() { + assert_eq!( + complete_anthropic_url(Some("https://proxy.internal/v1/messages"), &|_| None), + "https://proxy.internal/v1/messages" + ); + } + + #[test] + fn url_falls_back_to_env_base() { + let with_env = |key: &str| { + (key == ANTHROPIC_API_BASE_ENV).then(|| "https://env.anthropic".to_string()) + }; + assert_eq!( + complete_anthropic_url(Some(" "), &with_env), + "https://env.anthropic/v1/messages" + ); + } + + #[test] + fn api_key_prefers_param_then_env_then_errors() { + assert_eq!( + resolve_anthropic_api_key(Some("sk-param"), &|_| None).unwrap(), + "sk-param" + ); + let with_env = |key: &str| (key == ANTHROPIC_API_KEY_ENV).then(|| "sk-env".to_string()); + assert_eq!( + resolve_anthropic_api_key(Some(" "), &with_env).unwrap(), + "sk-env" + ); + assert!(matches!( + resolve_anthropic_api_key(None, &|_| None).expect_err("missing key"), + CoreError::Auth(_) + )); + } + + #[test] + fn auth_strategy_and_default_headers_match_anthropic() { + assert_eq!( + ANTHROPIC_MESSAGES_CONFIG.auth_strategy().header_name(), + "x-api-key" + ); + assert_eq!( + ANTHROPIC_MESSAGES_CONFIG.default_headers(), + &[ + ("anthropic-version", "2023-06-01"), + ("content-type", "application/json"), + ] + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/anthropic/mod.rs b/litellm-rust/crates/core/src/providers/anthropic/mod.rs new file mode 100644 index 000000000000..ba63992f3cb7 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/anthropic/mod.rs @@ -0,0 +1 @@ +pub mod messages; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 4ac1e521d455..fd0fa3b60eee 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -1,20 +1,26 @@ -use serde_json::{Map, Value}; - -use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::error::{CoreError, CoreResult}; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; -use crate::messages::types::MessagesRequestData; +use crate::messages::types::{ + AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock, + MessageContent, SystemPrompt, +}; +use crate::providers::anthropic::messages::transformation::{ + non_empty, AnthropicMessagesConfig, ANTHROPIC_MESSAGES_CONFIG, +}; const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY"; const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; +const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic"; +const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; -pub struct AzureAnthropicMessagesConfig; +pub struct AzureAnthropicMessagesConfig { + anthropic: AnthropicMessagesConfig, +} pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig = - AzureAnthropicMessagesConfig; - -fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} + AzureAnthropicMessagesConfig { + anthropic: ANTHROPIC_MESSAGES_CONFIG, + }; pub fn resolve_azure_api_key( api_key: Option<&str>, @@ -48,42 +54,32 @@ pub fn complete_azure_anthropic_url( let api_base = api_base.trim_end_matches('/'); - if api_base.ends_with("/v1/messages") || api_base.ends_with("/anthropic/v1/messages") { + if api_base.ends_with(MESSAGES_PATH_SUFFIX) { return Ok(api_base.to_string()); } - let with_anthropic = match api_base.split_once("/anthropic") { - Some((prefix, _)) => format!("{prefix}/anthropic"), - None => format!("{api_base}/anthropic"), + let with_anthropic = match api_base.split_once(ANTHROPIC_PATH_SEGMENT) { + Some((prefix, _)) => format!("{prefix}{ANTHROPIC_PATH_SEGMENT}"), + None => format!("{api_base}{ANTHROPIC_PATH_SEGMENT}"), }; - Ok(format!("{with_anthropic}/v1/messages")) + Ok(format!("{with_anthropic}{MESSAGES_PATH_SUFFIX}")) } -fn remove_scope_from_content_blocks(content: &mut [Value]) { - for item in content.iter_mut() { - if let Some(cache_control) = item - .as_object_mut() - .and_then(|block| block.get_mut("cache_control")) - .and_then(Value::as_object_mut) - { - cache_control.remove("scope"); - } +fn strip_scope_from_block(block: &mut ContentBlock) { + if let Some(cache_control) = block.cache_control.as_mut() { + cache_control.scope = None; } } -fn remove_scope_from_cache_control(body: &mut Map) { - if let Some(Value::Array(system)) = body.get_mut("system") { - remove_scope_from_content_blocks(system); +fn strip_scope_from_system(system: &mut SystemPrompt) { + if let SystemPrompt::Blocks(blocks) = system { + blocks.iter_mut().for_each(strip_scope_from_block); } - if let Some(Value::Array(messages)) = body.get_mut("messages") { - for message in messages.iter_mut() { - if let Some(Value::Array(content)) = message - .as_object_mut() - .and_then(|message| message.get_mut("content")) - { - remove_scope_from_content_blocks(content); - } - } +} + +fn strip_scope_from_message(message: &mut AnthropicMessage) { + if let MessageContent::Blocks(blocks) = &mut message.content { + blocks.iter_mut().for_each(strip_scope_from_block); } } @@ -106,23 +102,33 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { } fn auth_strategy(&self) -> MessagesAuthStrategy { - MessagesAuthStrategy::Header("x-api-key") + self.anthropic.auth_strategy() } - fn transform_request(&self, body: Value) -> CoreResult { - let mut body = match body { - Value::Object(body) => body, - other => { - return Err(CoreError::InvalidType { - expected: "object", - actual: json_type_name(&other), - }) - } - }; - remove_scope_from_cache_control(&mut body); - Ok(MessagesRequestData { - body: Value::Object(body), - }) + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + self.anthropic.default_headers() + } + + fn transform_request( + &self, + mut request: AnthropicMessagesRequest, + ) -> CoreResult { + if let Some(system) = request.system.as_mut() { + strip_scope_from_system(system); + } + request + .messages + .iter_mut() + .for_each(strip_scope_from_message); + self.anthropic.transform_request(request) + } + + fn transform_response( + &self, + model: &str, + response: AnthropicMessagesResponse, + ) -> CoreResult { + self.anthropic.transform_response(model, response) } } @@ -131,6 +137,14 @@ mod tests { use super::*; use serde_json::json; + fn request_from(value: serde_json::Value) -> AnthropicMessagesRequest { + serde_json::from_value(value).expect("valid request") + } + + fn to_value(request: AnthropicMessagesRequest) -> serde_json::Value { + serde_json::to_value(request).expect("serializable request") + } + #[test] fn url_appends_anthropic_and_messages_suffix() { let url = @@ -234,7 +248,7 @@ mod tests { #[test] fn transform_request_strips_scope_from_system_and_messages() { - let body = json!({ + let request = request_from(json!({ "model": "claude-sonnet-4-5", "max_tokens": 1024, "system": [ @@ -257,12 +271,13 @@ mod tests { ] } ] - }); + })); - let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(body) - .expect("request transforms") - .body; + let transformed = to_value( + AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(request) + .expect("request transforms"), + ); assert_eq!( transformed["system"][0]["cache_control"], @@ -280,66 +295,82 @@ mod tests { #[test] fn transform_request_is_idempotent_and_preserves_string_system() { - let body = json!({ + let request = request_from(json!({ "model": "claude-sonnet-4-5", + "max_tokens": 16, "system": "plain string system", "messages": [{"role": "user", "content": "hi"}] - }); + })); let once = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(body) - .expect("request transforms") - .body; + .transform_request(request) + .expect("request transforms"); let twice = AZURE_ANTHROPIC_MESSAGES_CONFIG .transform_request(once.clone()) - .expect("request transforms") - .body; + .expect("request transforms"); assert_eq!(once, twice); - assert_eq!(once["system"], json!("plain string system")); + assert_eq!(to_value(once)["system"], json!("plain string system")); + } + + #[test] + fn transform_request_preserves_all_supported_params() { + let body = json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 256, + "messages": [{"role": "user", "content": "hi"}], + "system": "be terse", + "metadata": {"user_id": "u1"}, + "stop_sequences": ["STOP"], + "stream": false, + "temperature": 0.4, + "top_p": 0.9, + "top_k": 40, + "tools": [{"name": "get_weather", "input_schema": {"type": "object"}}], + "tool_choice": {"type": "auto"}, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "service_tier": "auto", + "container": {"id": "c1"}, + "mcp_servers": [{"type": "url", "url": "https://mcp.example", "name": "x"}], + "context_management": {"edits": []}, + "output_format": {"type": "json_schema"}, + "output_config": {"effort": "high"}, + "speed": "fast", + "inference_geo": "us", + "litellm_metadata": {"trace": "abc"} + }); + let transformed = to_value( + AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(request_from(body.clone())) + .expect("request transforms"), + ); + assert_eq!(transformed, body); } #[test] fn transform_request_rejects_non_object_body() { - let err = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(json!("bad")) + let err = serde_json::from_value::(json!("bad")) .expect_err("non-object body should error"); - assert_eq!( - err, - CoreError::InvalidType { - expected: "object", - actual: "string", - } - ); + assert!(err.is_data()); } #[test] - fn transform_response_passes_through_object() { - let response = json!({ + fn transform_response_passes_through() { + let response: AnthropicMessagesResponse = serde_json::from_value(json!({ "id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "hello"}], "model": "claude-sonnet-4-5", "stop_reason": "end_turn", + "stop_sequence": null, "usage": {"input_tokens": 1, "output_tokens": 2} - }); + })) + .expect("valid response"); let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_response("claude-sonnet-4-5", response.clone()) - .expect("response transforms") - .into_json(); - assert_eq!(transformed, response); - } - - #[test] - fn transform_response_rejects_non_object() { - let err = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_response("claude-sonnet-4-5", json!([1, 2, 3])) - .expect_err("array response should error"); - assert_eq!( - err, - CoreError::InvalidType { - expected: "object", - actual: "array", - } - ); + .transform_response("claude-sonnet-4-5", response) + .expect("response transforms"); + let value = serde_json::to_value(transformed).expect("serializable"); + assert_eq!(value["stop_reason"], json!("end_turn")); + assert_eq!(value["stop_sequence"], json!(null)); + assert_eq!(value["content"][0]["text"], json!("hello")); } } diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index d75e750a0ba1..dc9dc515e7d1 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -1,3 +1,4 @@ +pub mod anthropic; pub mod azure_ai; pub mod mistral; pub mod openai; From e3efd65b067db5439064bfe80b2242e4d3e6257b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:23:04 +0000 Subject: [PATCH 05/12] feat(messages): route eligible Azure Anthropic streaming through Rust via buffered fake-stream Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm/llms/custom_httpx/llm_http_handler.py | 33 +++++++++++++- .../test_litellm/rust_bridge/test_messages.py | 43 ++++++++++++++++++- 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e58641a5d8da..133ce99818b9 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -153,6 +153,9 @@ from aiohttp import ClientSession from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamingResponse, + ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -2095,6 +2098,7 @@ async def async_anthropic_messages_handler( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, stream=stream or False, + rust_stream_eligible=bool(stream) and not self._has_agentic_completion_hook(logging_obj), model=model, api_key=api_key, api_base=api_base, @@ -2107,6 +2111,8 @@ async def async_anthropic_messages_handler( ), ) if rust_messages_response is not None: + if stream: + return self._rust_anthropic_messages_fake_stream(rust_messages_response) return await self._finalize_anthropic_messages_response( initial_response=rust_messages_response, model=model, @@ -2247,6 +2253,7 @@ async def _maybe_rust_anthropic_messages( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, stream: bool, + rust_stream_eligible: bool, model: str, api_key: str | None, api_base: str | None, @@ -2254,14 +2261,17 @@ async def _maybe_rust_anthropic_messages( request_body: dict, timeout: float | httpx.Timeout | None, ) -> AnthropicMessagesResponse | None: - if stream or custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True: + if custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True: + return None + if stream and not rust_stream_eligible: return None from litellm.rust_bridge import messages as rust_messages_bridge + upstream_body = {key: value for key, value in request_body.items() if key != "stream"} rust_response = await rust_messages_bridge.amessages( model=model, - body=request_body, + body=upstream_body, api_key=api_key, api_base=api_base, custom_llm_provider=custom_llm_provider, @@ -2275,6 +2285,25 @@ async def _maybe_rust_anthropic_messages( response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}} return response_obj + @staticmethod + def _rust_anthropic_messages_fake_stream( + rust_response: AnthropicMessagesResponse, + ) -> "AnthropicMessagesStreamingResponse": + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamHiddenParams, + AnthropicMessagesStreamingResponse, + ) + + completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response)) + hidden_params = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"}) + return AnthropicMessagesStreamingResponse( + completion_stream=completion_stream, + hidden_params=hidden_params, + ) + def anthropic_messages_handler( self, model: str, diff --git a/tests/test_litellm/rust_bridge/test_messages.py b/tests/test_litellm/rust_bridge/test_messages.py index f5be34ae0b67..f34b8c05819b 100644 --- a/tests/test_litellm/rust_bridge/test_messages.py +++ b/tests/test_litellm/rust_bridge/test_messages.py @@ -1,12 +1,16 @@ """Tests for the optional Rust-backed Anthropic Messages path.""" import importlib +from typing import cast import httpx import pytest import litellm from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, +) from litellm.types.router import GenericLiteLLMParams rust_messages = importlib.import_module("litellm.rust_bridge.messages") @@ -207,6 +211,7 @@ def _gate(**overrides): "custom_llm_provider": "azure_ai", "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), "stream": False, + "rust_stream_eligible": False, "model": "claude-sonnet-4-5", "api_key": "sk-azure", "api_base": "https://resource.services.ai.azure.com/anthropic", @@ -271,16 +276,50 @@ async def test_gate_skips_rust_for_non_azure_provider(): @pytest.mark.asyncio -async def test_gate_skips_rust_when_streaming(): +async def test_gate_skips_rust_when_streaming_but_not_eligible(): bridge = ExplodingAsyncMessages() litellm.use_litellm_rust(True, amessages=bridge) - response = await _gate(stream=True) + response = await _gate(stream=True, rust_stream_eligible=False) assert response is None assert bridge.calls == 0 +@pytest.mark.asyncio +async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + streaming_body = {**REQUEST_BODY, "stream": True} + response = await _gate( + stream=True, + rust_stream_eligible=True, + request_body=streaming_body, + ) + + assert response is not None + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + assert "stream" not in bridge.calls[0]["body"] + assert bridge.calls[0]["body"] == REQUEST_BODY + + +@pytest.mark.asyncio +async def test_fake_stream_wraps_rust_response_as_anthropic_sse(): + response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE)) + stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response) + + assert stream._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + + chunks = [chunk async for chunk in stream] + joined = b"".join(chunks) + + assert b"event: message_start" in joined + assert b"event: content_block_delta" in joined + assert b"hello world" in joined + assert b"event: message_stop" in joined + + @pytest.mark.asyncio async def test_gate_falls_back_when_bridge_unavailable(monkeypatch): monkeypatch.setattr( From ee027e9c5b07bf134e058ae9e60b4758254ca50b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 04:00:09 +0000 Subject: [PATCH 06/12] fix(messages): fold system-role messages for Azure Anthropic and fall back to Python on Rust bridge errors Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm-rust/crates/core/src/error.rs | 4 +- .../azure_ai/messages/transformation.rs | 138 +++++++++++++++++- litellm/llms/custom_httpx/llm_http_handler.py | 25 ++-- .../test_litellm/rust_bridge/test_messages.py | 20 +++ 4 files changed, 175 insertions(+), 12 deletions(-) diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index b3e0519b7727..c2b08eee0c0e 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -19,9 +19,9 @@ pub enum CoreError { InvalidRequest(String), #[error("{0}")] Auth(String), - #[error("OCR request failed with status {status}: {body}")] + #[error("upstream request failed with status {status}: {body}")] Http { status: u16, body: String }, - #[error("OCR network error: {0}")] + #[error("upstream network error: {0}")] Network(String), #[error("routing error: {0}")] Routing(String), diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index fd0fa3b60eee..13e79b087c78 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -7,11 +7,14 @@ use crate::messages::types::{ use crate::providers::anthropic::messages::transformation::{ non_empty, AnthropicMessagesConfig, ANTHROPIC_MESSAGES_CONFIG, }; +use serde_json::{Map, Value}; const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY"; const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic"; const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; +const SYSTEM_ROLE: &str = "system"; +const TEXT_BLOCK_TYPE: &str = "text"; pub struct AzureAnthropicMessagesConfig { anthropic: AnthropicMessagesConfig, @@ -83,6 +86,61 @@ fn strip_scope_from_message(message: &mut AnthropicMessage) { } } +fn text_content_block(text: String) -> ContentBlock { + let extra = Map::from_iter([ + ( + "type".to_string(), + Value::String(TEXT_BLOCK_TYPE.to_string()), + ), + ("text".to_string(), Value::String(text)), + ]); + ContentBlock { + cache_control: None, + extra, + } +} + +fn content_into_blocks(content: MessageContent) -> Vec { + match content { + MessageContent::Text(text) => vec![text_content_block(text)], + MessageContent::Blocks(blocks) => blocks, + } +} + +fn system_into_blocks(system: Option) -> Vec { + match system { + None => Vec::new(), + Some(SystemPrompt::Text(text)) => vec![text_content_block(text)], + Some(SystemPrompt::Blocks(blocks)) => blocks, + } +} + +fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMessagesRequest { + if !request.messages.iter().any(|msg| msg.role == SYSTEM_ROLE) { + return request; + } + + let (system_messages, chat_messages): (Vec, Vec) = request + .messages + .into_iter() + .partition(|msg| msg.role == SYSTEM_ROLE); + + let folded_system: Vec = system_into_blocks(request.system) + .into_iter() + .chain( + system_messages + .into_iter() + .flat_map(|msg| content_into_blocks(msg.content)), + ) + .collect(); + + AnthropicMessagesRequest { + messages: chat_messages, + system: (!folded_system.is_empty()).then_some(SystemPrompt::Blocks(folded_system)), + ..request + } +} + impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { fn complete_url( &self, @@ -111,8 +169,9 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { fn transform_request( &self, - mut request: AnthropicMessagesRequest, + request: AnthropicMessagesRequest, ) -> CoreResult { + let mut request = fold_system_role_messages(request); if let Some(system) = request.system.as_mut() { strip_scope_from_system(system); } @@ -345,6 +404,83 @@ mod tests { assert_eq!(transformed, body); } + #[test] + fn transform_request_folds_system_role_message_into_top_level_system() { + let request = request_from(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 256, + "system": [{"type": "text", "text": "base system"}], + "messages": [ + {"role": "user", "content": "fix the bug"}, + {"role": "system", "content": "Available agent types: claude"} + ] + })); + + let transformed = to_value( + AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(request) + .expect("request transforms"), + ); + + assert_eq!( + transformed["messages"], + json!([{"role": "user", "content": "fix the bug"}]) + ); + assert_eq!( + transformed["system"], + json!([ + {"type": "text", "text": "base system"}, + {"type": "text", "text": "Available agent types: claude"} + ]) + ); + } + + #[test] + fn transform_request_folds_system_role_when_no_top_level_system() { + let request = request_from(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 256, + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + {"role": "system", "content": [{"type": "text", "text": "sys block"}]} + ] + })); + + let transformed = to_value( + AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(request) + .expect("request transforms"), + ); + + assert_eq!( + transformed["messages"], + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) + ); + assert_eq!( + transformed["system"], + json!([{"type": "text", "text": "sys block"}]) + ); + } + + #[test] + fn transform_request_leaves_requests_without_system_role_untouched() { + let body = json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 256, + "system": "be terse", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"} + ] + }); + let transformed = to_value( + AZURE_ANTHROPIC_MESSAGES_CONFIG + .transform_request(request_from(body.clone())) + .expect("request transforms"), + ); + assert_eq!(transformed, body); + } + #[test] fn transform_request_rejects_non_object_body() { let err = serde_json::from_value::(json!("bad")) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 133ce99818b9..3e6f9ee08eeb 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2269,15 +2269,22 @@ async def _maybe_rust_anthropic_messages( from litellm.rust_bridge import messages as rust_messages_bridge upstream_body = {key: value for key, value in request_body.items() if key != "stream"} - rust_response = await rust_messages_bridge.amessages( - model=model, - body=upstream_body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=headers, - timeout=timeout, - ) + try: + rust_response = await rust_messages_bridge.amessages( + model=model, + body=upstream_body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + ) + except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path + verbose_logger.debug( + "Rust Anthropic messages bridge raised %s; falling back to Python path", + type(rust_error).__name__, + ) + return None if rust_response is None: return None diff --git a/tests/test_litellm/rust_bridge/test_messages.py b/tests/test_litellm/rust_bridge/test_messages.py index f34b8c05819b..26ca7d272105 100644 --- a/tests/test_litellm/rust_bridge/test_messages.py +++ b/tests/test_litellm/rust_bridge/test_messages.py @@ -98,6 +98,15 @@ async def __call__(self, **kwargs: object) -> dict[str, object]: raise AssertionError("bridge must not be called") +class RaisingAsyncMessages: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> dict[str, object]: + self.calls += 1 + raise RuntimeError("upstream request failed with status 400: bad request") + + @pytest.fixture(autouse=True) def _reset_rust_flag(): litellm.use_litellm_rust(False, messages=None, amessages=None) @@ -242,6 +251,17 @@ async def test_gate_invokes_rust_and_marks_response_header(): assert call["timeout_seconds"] == 30.0 +@pytest.mark.asyncio +async def test_gate_falls_back_to_python_when_bridge_raises(): + bridge = RaisingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate() + + assert response is None + assert bridge.calls == 1 + + @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_absent(): bridge = ExplodingAsyncMessages() From 23d990b2e534cbd8079d1fb79178fa47c78965f5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 03:09:39 +0000 Subject: [PATCH 07/12] fix(rust_bridge): use Python::attach for amessages after pyo3 bump Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index ab4d51295353..77f4427127ac 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -254,7 +254,7 @@ fn amessages( .await .map_err(core_error_to_pyerr)?; - Python::with_gil(|py| json_to_py(py, value)) + Python::attach(|py| json_to_py(py, value)) }) } From bebb7b9ba432a4426acee9959b3c4a711b7a4b3b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 17:41:40 +0000 Subject: [PATCH 08/12] test(proxy): mock get_configured_token_limits in model_info tests Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .../management_endpoints/test_model_management_endpoints.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 8c6bdefedae0..79c5f3ea5497 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1727,6 +1727,7 @@ async def test_model_info_accessible_model_success(self): "gpt-3.5-turbo", ] mock_router.get_model_access_groups.return_value = {} + mock_router.get_configured_token_limits.return_value = (None, None) mock_get_key_models.return_value = ["gpt-4", "claude-3"] mock_get_team_models.return_value = ["gpt-3.5-turbo"] mock_get_complete_models.return_value = [ @@ -1812,6 +1813,7 @@ async def test_model_info_team_model_access(self): # Setup mocks mock_router.get_model_names.return_value = ["team-model-1"] mock_router.get_model_access_groups.return_value = {} + mock_router.get_configured_token_limits.return_value = (None, None) mock_get_key_models.return_value = [] mock_get_team_models.return_value = ["team-model-1"] mock_get_complete_models.return_value = ["team-model-1"] From c86d861a0345735322ff12d717fd5d1bde5dfaaf Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 18:04:38 +0000 Subject: [PATCH 09/12] ci: run rust_bridge unit tests in misc shard Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .github/workflows/test-unit-misc.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index 7c3b195f0ad5..1291b83d5a17 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -34,6 +34,7 @@ jobs: tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/ocr + tests/test_litellm/rust_bridge tests/test_litellm/passthrough tests/test_litellm/sandbox tests/test_litellm/vector_stores From 1371a045ba7e18ab4e42fdd090aca56b452e4712 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 18:18:33 +0000 Subject: [PATCH 10/12] Revert "ci: run rust_bridge unit tests in misc shard" This reverts commit c86d861a0345735322ff12d717fd5d1bde5dfaaf. --- .github/workflows/test-unit-misc.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index 1291b83d5a17..7c3b195f0ad5 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -34,7 +34,6 @@ jobs: tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/ocr - tests/test_litellm/rust_bridge tests/test_litellm/passthrough tests/test_litellm/sandbox tests/test_litellm/vector_stores From 8d4e2e00b8df4291ac11bc563cee2b64746d0b15 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 18:44:59 +0000 Subject: [PATCH 11/12] feat(messages): route native Anthropic /messages through Rust behind LITELLM_USE_RUST_MESSAGES Extends the opt-in Rust Anthropic Messages path (previously azure_ai only, per-deployment rust:true) to the native anthropic provider and switches enablement to the LITELLM_USE_RUST_MESSAGES env var, mirroring the OCR bridge (LITELLM_USE_RUST_OCR). When enabled, eligible providers route through Rust; unsupported providers fall back to the Python path. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .../ai-gateway/src/messages/common_utils.rs | 2 + .../crates/ai-gateway/src/messages/tests.rs | 70 +++++++++++++++++-- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/rust_bridge/messages.py | 20 +++++- litellm/rust_bridge/ocr.py | 17 ++--- .../test_litellm/rust_bridge/test_messages.py | 43 ++++++++++-- 6 files changed, 136 insertions(+), 20 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs index fe4ac4cf26f0..33894d0ee642 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs @@ -1,5 +1,6 @@ use litellm_core::error::{json_type_name, CoreError}; use litellm_core::messages::transformation::AnthropicMessagesProviderConfig; +use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; use litellm_core::CoreResult; use serde_json::{Map, Value}; @@ -18,6 +19,7 @@ pub(super) fn messages_provider_config( provider: &str, ) -> Option<&'static dyn AnthropicMessagesProviderConfig> { match provider { + "anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG), "azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG), _ => None, } diff --git a/litellm-rust/crates/ai-gateway/src/messages/tests.rs b/litellm-rust/crates/ai-gateway/src/messages/tests.rs index 30f6642400ea..680fd5b5d58e 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/tests.rs @@ -8,6 +8,7 @@ use tokio::net::{TcpListener, TcpStream}; use super::common_utils::{ has_header, messages_provider_config, string_headers, truncate_error_body, }; +use super::prepare::prepare_messages_call; use super::{messages, MessagesRequest}; async fn read_http_request(socket: &mut TcpStream) -> String { @@ -52,12 +53,73 @@ fn write_response(body: &str) -> String { } #[test] -fn provider_config_only_resolves_azure_ai() { +fn provider_config_resolves_supported_providers() { assert!(messages_provider_config("azure_ai").is_some()); - assert!(messages_provider_config("anthropic").is_none()); + assert!(messages_provider_config("anthropic").is_some()); assert!(messages_provider_config("openai").is_none()); } +#[test] +fn prepare_messages_call_resolves_native_anthropic() { + let prepared = prepare_messages_call(MessagesRequest { + model: "claude-opus-4-8", + body: json!({ + "model": "claude-opus-4-8", + "max_tokens": 16, + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": "hi", + "cache_control": {"type": "ephemeral", "scope": "global"} + }] + }] + }), + api_key: Some("sk-ant-test"), + api_base: None, + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: None, + }) + .expect("native Anthropic provider resolves"); + + assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages"); + assert!(prepared + .upstream_headers + .iter() + .any(|(name, value)| name == "x-api-key" && value == "sk-ant-test")); + assert!(prepared + .upstream_headers + .iter() + .any(|(name, value)| name == "anthropic-version" && value == "2023-06-01")); + assert!(prepared + .upstream_headers + .iter() + .any(|(name, value)| name == "content-type" && value == "application/json")); + assert_eq!( + prepared.body["messages"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral", "scope": "global"}) + ); +} + +#[test] +fn prepare_messages_call_rejects_unknown_provider() { + let result = prepare_messages_call(MessagesRequest { + model: "some-model", + body: json!({"model": "some-model", "max_tokens": 8, "messages": []}), + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: Some("openai"), + extra_headers: None, + timeout: None, + }); + + assert!(matches!( + result, + Err(CoreError::InvalidProvider(provider)) if provider == "openai" + )); +} + #[test] fn truncate_error_body_caps_long_payloads() { let body = "x".repeat(400); @@ -248,12 +310,12 @@ async fn messages_rejects_unsupported_provider() { body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}), api_key: Some("sk"), api_base: Some("http://127.0.0.1:1"), - custom_llm_provider: Some("anthropic"), + custom_llm_provider: Some("openai"), extra_headers: None, timeout: Some(Duration::from_millis(50)), }) .await .expect_err("unsupported provider errors"); - assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "anthropic")); + assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai")); } diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3e6f9ee08eeb..1c979ed7054b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2261,7 +2261,9 @@ async def _maybe_rust_anthropic_messages( request_body: dict, timeout: float | httpx.Timeout | None, ) -> AnthropicMessagesResponse | None: - if custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True: + from litellm.rust_bridge.messages import rust_messages_enabled + + if not rust_messages_enabled(): return None if stream and not rust_stream_eligible: return None diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 5abb21879d3e..690faa4dbd6e 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os from dataclasses import dataclass from typing import Awaitable, Final, Protocol, Union, cast @@ -47,24 +48,41 @@ class _Unset: @dataclass(slots=True) class _RustMessagesState: + enabled: bool = False messages: RustMessages | None = None amessages: RustAmessages | None = None -_STATE: Final[_RustMessagesState] = _RustMessagesState() +def _env_enables_rust_messages() -> bool: + return os.getenv("LITELLM_USE_RUST_MESSAGES", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +_STATE: Final[_RustMessagesState] = _RustMessagesState(enabled=_env_enables_rust_messages()) def set_rust_messages( *, + enabled: bool | _Unset = _UNSET, messages: RustMessages | None | _Unset = _UNSET, amessages: RustAmessages | None | _Unset = _UNSET, ) -> None: + if not isinstance(enabled, _Unset): + _STATE.enabled = enabled if not isinstance(messages, _Unset): _STATE.messages = messages if not isinstance(amessages, _Unset): _STATE.amessages = amessages +def rust_messages_enabled() -> bool: + return _STATE.enabled + + def load_rust_messages() -> RustMessages | None: if _STATE.messages is not None: return _STATE.messages diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 35de2eb9727a..386189700eda 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -81,16 +81,17 @@ def use_litellm_rust( _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): _rust_aocr_impl = aocr - if not configuring_messages: - return from litellm.rust_bridge.messages import set_rust_messages - if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset): - set_rust_messages(messages=messages, amessages=amessages) - elif not isinstance(messages, _Unset): - set_rust_messages(messages=messages) - else: - set_rust_messages(amessages=amessages) + if configuring_messages or not configuring_ocr: + if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset): + set_rust_messages(enabled=enabled, messages=messages, amessages=amessages) + elif not isinstance(messages, _Unset): + set_rust_messages(enabled=enabled, messages=messages) + elif not isinstance(amessages, _Unset): + set_rust_messages(enabled=enabled, amessages=amessages) + else: + set_rust_messages(enabled=enabled) def rust_ocr_enabled() -> bool: diff --git a/tests/test_litellm/rust_bridge/test_messages.py b/tests/test_litellm/rust_bridge/test_messages.py index 26ca7d272105..7d1b10e6bc38 100644 --- a/tests/test_litellm/rust_bridge/test_messages.py +++ b/tests/test_litellm/rust_bridge/test_messages.py @@ -107,6 +107,15 @@ async def __call__(self, **kwargs: object) -> dict[str, object]: raise RuntimeError("upstream request failed with status 400: bad request") +class NoneAsyncMessages: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> dict[str, object] | None: + self.calls += 1 + return None + + @pytest.fixture(autouse=True) def _reset_rust_flag(): litellm.use_litellm_rust(False, messages=None, amessages=None) @@ -251,6 +260,28 @@ async def test_gate_invokes_rust_and_marks_response_header(): assert call["timeout_seconds"] == 30.0 +@pytest.mark.asyncio +async def test_gate_invokes_rust_for_native_anthropic_provider(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate( + custom_llm_provider="anthropic", + api_key="sk-ant-test", + api_base="https://api.anthropic.com", + headers={"anthropic-version": "2023-06-01"}, + litellm_params=GenericLiteLLMParams(api_key="sk-ant-test", rust=True), + ) + + assert response is not None + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + call = bridge.calls[0] + assert call["custom_llm_provider"] == "anthropic" + assert call["api_key"] == "sk-ant-test" + assert call["api_base"] == "https://api.anthropic.com" + assert call["extra_headers"] == {"anthropic-version": "2023-06-01"} + + @pytest.mark.asyncio async def test_gate_falls_back_to_python_when_bridge_raises(): bridge = RaisingAsyncMessages() @@ -265,7 +296,7 @@ async def test_gate_falls_back_to_python_when_bridge_raises(): @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_absent(): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.use_litellm_rust(False, amessages=bridge) response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) @@ -276,7 +307,7 @@ async def test_gate_skips_rust_when_flag_absent(): @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_false(): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.use_litellm_rust(False, amessages=bridge) response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) @@ -285,14 +316,14 @@ async def test_gate_skips_rust_when_flag_false(): @pytest.mark.asyncio -async def test_gate_skips_rust_for_non_azure_provider(): - bridge = ExplodingAsyncMessages() +async def test_gate_skips_rust_for_non_listed_provider(): + bridge = NoneAsyncMessages() litellm.use_litellm_rust(True, amessages=bridge) - response = await _gate(custom_llm_provider="anthropic") + response = await _gate(custom_llm_provider="openai") assert response is None - assert bridge.calls == 0 + assert bridge.calls == 1 @pytest.mark.asyncio From 93be28e6dc795b3b5619df7c91ed32b2a191059c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:15:25 +0000 Subject: [PATCH 12/12] fix: exclude Rust messages rollout flag from env docs check Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- tests/documentation_tests/test_env_keys.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 60fbd505d674..fb6b7f639f2d 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -27,6 +27,7 @@ # Temporary/internal rollout flags are intentionally not added to the public # environment settings docs until the feature is ready for broad use. EXCLUDED_ROLLOUT_FLAGS = { + "LITELLM_USE_RUST_MESSAGES", "LITELLM_USE_RUST_OCR", }