diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index edfb3536ad3..c05f98d22ae 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 109b648f5db..557fe5d53d4 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 00000000000..b784d2b62a1 --- /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 3b566027646..7bc9642d192 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 d8ef7bb5ba1..db4c8211a5a 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 00000000000..6281270b964 --- /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 00000000000..fe4ac4cf26f --- /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 00000000000..ad58d74a87f --- /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 00000000000..7ed81474c47 --- /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 00000000000..72e32197267 --- /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 00000000000..30f6642400e --- /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 00000000000..6840ff57cc4 --- /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 555a04ce853..27154f5a08b 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 00000000000..ec2fbb969a6 --- /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 00000000000..448f7dd6dc8 --- /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 00000000000..e56e924e910 --- /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 00000000000..f239b6921fa --- /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 00000000000..4ac1e521d45 --- /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 3621ff6a2fd..5d13fa93e00 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 946a99f990c..80662b66d70 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 b47fc50e196..e58641a5d8d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2091,6 +2091,34 @@ class BaseLLMHTTPHandler: }, ) + 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 @@ class BaseLLMHTTPHandler: 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 @@ class BaseLLMHTTPHandler: "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 00000000000..d51689f5f3b --- /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 36a088b6b1a..3654b9fb232 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 00000000000..4407986c3da --- /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 00000000000..e69de29bb2d 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 00000000000..6f0e1e190a1 --- /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 87b4c96e323..8fe7c588cf1 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