diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index fbba9969c28..3d6fffe7304 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -113,6 +113,7 @@ jobs: tests/test_litellm/rag tests/test_litellm/realtime_api tests/test_litellm/rerank_api + tests/test_litellm/rust_bridge tests/test_litellm/sandbox tests/test_litellm/test_router tests/test_litellm/vector_stores diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 8b6896f3846..0c9faeda6e7 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -295,6 +295,8 @@ fn core_error_kind(error: &CoreError) -> &'static str { CoreError::Http { .. } => "HttpError", CoreError::InvalidResponse(_) => "InvalidResponse", CoreError::Network(_) => "NetworkError", + CoreError::Connect(_) => "ConnectError", CoreError::Routing(_) => "RoutingError", + CoreError::Unsupported(_) => "UnsupportedRequest", } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index ffe2e0122c0..95df566dc53 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -324,6 +324,8 @@ fn core_error_kind(error: &CoreError) -> &'static str { CoreError::Http { .. } => "HttpError", CoreError::InvalidResponse(_) => "InvalidResponse", CoreError::Network(_) => "NetworkError", + CoreError::Connect(_) => "ConnectError", CoreError::Routing(_) => "RoutingError", + CoreError::Unsupported(_) => "UnsupportedRequest", } } diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index a34b2edd7b8..7e38d10c6ff 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -105,12 +105,20 @@ impl IntoResponse for MessagesRouteError { ), CoreError::Http { .. } | CoreError::Network(_) + | CoreError::Connect(_) | CoreError::InvalidResponse(_) | CoreError::InvalidType { .. } | CoreError::MissingField(_) => ( StatusCode::BAD_GATEWAY, "messages provider request failed".to_string(), ), + // The gateway has no Python implementation to decline to, so a + // request the core cannot serve is reported to the caller. The + // reason is a fixed internal string, never provider content. + CoreError::Unsupported(reason) => ( + StatusCode::BAD_REQUEST, + format!("messages request is not supported: {reason}"), + ), }; ( status, diff --git a/litellm-rust/crates/core/src/chat_completions/client.rs b/litellm-rust/crates/core/src/chat_completions/client.rs new file mode 100644 index 00000000000..f2ef73ed030 --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/client.rs @@ -0,0 +1,15 @@ +use std::sync::OnceLock; +use std::time::Duration; + +use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_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(CHAT_COMPLETIONS_TIMEOUT_SECS)) + .connect_timeout(Duration::from_secs(CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS)) + .build() + .unwrap_or_else(|_| reqwest::Client::new()) + }) +} diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs new file mode 100644 index 00000000000..36eaf242a5a --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -0,0 +1,28 @@ +use serde_json::{Map, Value}; + +use crate::error::CoreResult; +use crate::http_utils::string_headers as shared_string_headers; +use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; + +use super::transformation::ChatCompletionsProviderConfig; + +const HEADER_CONTEXT: &str = "chat completions"; + +pub(super) fn chat_completions_provider_config( + provider: &str, +) -> Option<&'static dyn ChatCompletionsProviderConfig> { + match provider { + "anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG), + #[cfg(feature = "bedrock-auth")] + "bedrock" => Some( + &crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, + ), + _ => None, + } +} + +pub(super) fn string_headers( + extra_headers: Option>, +) -> CoreResult> { + shared_string_headers(HEADER_CONTEXT, extra_headers) +} diff --git a/litellm-rust/crates/core/src/chat_completions/conversation.rs b/litellm-rust/crates/core/src/chat_completions/conversation.rs new file mode 100644 index 00000000000..f7bdc60af37 --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/conversation.rs @@ -0,0 +1,254 @@ +//! Provider-neutral conversation shape. +//! +//! Both Anthropic Messages and Bedrock Converse want the same thing out of an +//! OpenAI message list: the system prompt lifted out, consecutive same-role +//! turns merged, and text blocks that are never empty. That normalization is +//! shared here so a provider config only renders the result into its own wire +//! shape. +//! +//! Mirrors Python's `anthropic_messages_pt` / +//! `_bedrock_converse_messages_pt` for the text-only surface this route +//! accepts; anything richer is declined upstream by the capability gate. + +use crate::constants::EMPTY_TEXT_PLACEHOLDER; + +use super::types::{ChatMessage, ChatMessageContent}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TurnRole { + User, + Assistant, +} + +impl TurnRole { + pub fn as_str(self) -> &'static str { + match self { + Self::User => "user", + Self::Assistant => "assistant", + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Turn { + pub role: TurnRole, + pub texts: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Conversation { + pub system: Vec, + pub turns: Vec, +} + +/// True when the conversation can be sent as-is. +/// +/// Python inserts a placeholder first user turn only under +/// `litellm.modify_params`, which the core cannot see, so a conversation that +/// does not open on a user turn is declined rather than guessed at. +impl Conversation { + pub fn opens_on_user_turn(&self) -> bool { + self.turns + .first() + .is_some_and(|turn| turn.role == TurnRole::User) + } +} + +fn message_texts(content: &ChatMessageContent) -> Vec { + match content { + ChatMessageContent::Text(text) => vec![text.clone()], + ChatMessageContent::Parts(parts) => parts + .iter() + .filter_map(|part| part.get("text").and_then(|text| text.as_str())) + .map(str::to_string) + .collect(), + } +} + +/// Python rewrites empty or whitespace-only text rather than dropping it, so an +/// entirely empty content list never reaches a provider that rejects one. +fn sanitize(text: String) -> String { + if text.trim().is_empty() { + return EMPTY_TEXT_PLACEHOLDER.to_string(); + } + text +} + +pub fn build_conversation(messages: &[ChatMessage]) -> Conversation { + let system = messages + .iter() + .filter(|message| message.role == "system") + .filter_map(|message| message.content.as_ref()) + .flat_map(message_texts) + .filter(|text| !text.is_empty()) + .collect(); + + let turns = messages + .iter() + .filter(|message| message.role != "system") + .fold(Vec::::new(), |mut turns, message| { + let role = if message.role == "assistant" { + TurnRole::Assistant + } else { + TurnRole::User + }; + let texts = message + .content + .as_ref() + .map(message_texts) + .unwrap_or_default() + .into_iter() + .map(sanitize); + match turns.last_mut() { + Some(last) if last.role == role => last.texts.extend(texts), + _ => turns.push(Turn { + role, + texts: texts.collect(), + }), + } + turns + }); + + // Anthropic and Bedrock both reject trailing whitespace on the final + // assistant turn, so Python right-strips it there; mirror that exactly. + let turns = match turns.split_last() { + Some((last, rest)) if last.role == TurnRole::Assistant => rest + .iter() + .cloned() + .chain([Turn { + role: last.role, + texts: last + .texts + .iter() + .map(|text| text.trim_end().to_string()) + .collect(), + }]) + .collect(), + _ => turns, + }; + + Conversation { system, turns } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn messages(value: serde_json::Value) -> Vec { + serde_json::from_value(value).expect("valid messages") + } + + #[test] + fn lifts_system_messages_out_of_the_turn_list() { + let conversation = build_conversation(&messages(json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]))); + assert_eq!(conversation.system, vec!["be terse".to_string()]); + assert_eq!( + conversation.turns, + vec![Turn { + role: TurnRole::User, + texts: vec!["hi".to_string()] + }] + ); + } + + #[test] + fn merges_consecutive_same_role_turns() { + let conversation = build_conversation(&messages(json!([ + {"role": "user", "content": "one"}, + {"role": "user", "content": "two"}, + {"role": "assistant", "content": "ack"}, + {"role": "user", "content": "three"} + ]))); + assert_eq!( + conversation.turns, + vec![ + Turn { + role: TurnRole::User, + texts: vec!["one".to_string(), "two".to_string()] + }, + Turn { + role: TurnRole::Assistant, + texts: vec!["ack".to_string()] + }, + Turn { + role: TurnRole::User, + texts: vec!["three".to_string()] + }, + ] + ); + } + + #[test] + fn flattens_text_parts_in_order() { + let conversation = build_conversation(&messages(json!([ + {"role": "user", "content": [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "second"} + ]} + ]))); + assert_eq!( + conversation.turns[0].texts, + vec!["first".to_string(), "second".to_string()] + ); + } + + #[test] + fn rewrites_empty_and_whitespace_only_text_to_the_python_placeholder() { + let conversation = build_conversation(&messages(json!([ + {"role": "user", "content": ""}, + {"role": "assistant", "content": " "}, + {"role": "user", "content": "real"} + ]))); + assert_eq!(conversation.turns[0].texts, vec![EMPTY_TEXT_PLACEHOLDER]); + assert_eq!(conversation.turns[1].texts, vec![EMPTY_TEXT_PLACEHOLDER]); + } + + #[test] + fn right_strips_only_the_final_assistant_turn() { + let conversation = build_conversation(&messages(json!([ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "kept "}, + {"role": "user", "content": "more"}, + {"role": "assistant", "content": "stripped "} + ]))); + assert_eq!(conversation.turns[1].texts, vec!["kept ".to_string()]); + assert_eq!(conversation.turns[3].texts, vec!["stripped".to_string()]); + } + + #[test] + fn does_not_strip_when_the_last_turn_is_a_user_turn() { + let conversation = build_conversation(&messages(json!([ + {"role": "assistant", "content": "kept "}, + {"role": "user", "content": "hi "} + ]))); + assert_eq!(conversation.turns[0].texts, vec!["kept ".to_string()]); + assert_eq!(conversation.turns[1].texts, vec!["hi ".to_string()]); + } + + #[test] + fn reports_whether_the_conversation_opens_on_a_user_turn() { + assert!( + build_conversation(&messages(json!([{"role": "user", "content": "hi"}]))) + .opens_on_user_turn() + ); + assert!( + !build_conversation(&messages(json!([{"role": "assistant", "content": "hi"}]))) + .opens_on_user_turn() + ); + assert!(!Conversation::default().opens_on_user_turn()); + } + + #[test] + fn drops_empty_system_text_the_way_python_skips_empty_system_blocks() { + let conversation = build_conversation(&messages(json!([ + {"role": "system", "content": ""}, + {"role": "system", "content": "kept"}, + {"role": "user", "content": "hi"} + ]))); + assert_eq!(conversation.system, vec!["kept".to_string()]); + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs new file mode 100644 index 00000000000..afc4529fd26 --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -0,0 +1,147 @@ +use serde_json::Value; + +use crate::error::{CoreError, CoreResult}; +use crate::http_utils::truncate_error_body; + +use super::client::http_client; +use super::transformation::ChatCompletionsAuth; +use super::types::{ + ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData, +}; + +pub(super) async fn execute_chat_completions_provider_call( + request: ProviderChatCompletionsRequest, +) -> CoreResult { + let body = serde_json::to_vec(&request.body).map_err(|err| { + CoreError::InvalidRequest(format!( + "failed to serialize chat completions request: {err}" + )) + })?; + let headers = signed_headers(&request, &body).await?; + + let mut request_builder = http_client().post(&request.url).body(body); + for (key, value) in &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| { + // Failing to establish the connection means the request never went out, + // so the host can still serve it. Everything else here, a timeout + // above all, may have reached the provider and been answered. + if err.is_connect() || err.is_builder() { + CoreError::Connect(err.to_string()) + } else { + 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 body: Value = serde_json::from_str(&text).map_err(|err| { + CoreError::InvalidResponse(format!("invalid chat completions response JSON: {err}")) + })?; + request + .config + .transform_response(&request.model, ProviderChatResponseData { body }) + .map_err(as_response_error) +} + +/// Re-tag an error raised while normalizing a response the provider already +/// returned. +/// +/// A config reports the same variants on either side of the call: a missing +/// field or an unsupported block can mean "this request cannot be translated" +/// during prepare and "this response cannot be normalized" here. Only the +/// second kind has already been billed, and a host that keeps a reference +/// implementation must not retry those, so collapse them to one variant that +/// can only mean the provider was already called. +pub(super) fn as_response_error(err: CoreError) -> CoreError { + match err { + already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already, + other => CoreError::InvalidResponse(other.to_string()), + } +} + +#[cfg(feature = "bedrock-auth")] +pub(super) async fn signed_headers( + request: &ProviderChatCompletionsRequest, + body: &[u8], +) -> CoreResult> { + use std::collections::BTreeMap; + use std::time::SystemTime; + + use crate::providers::bedrock::aws_base::{ + aws_auth_config, aws_signature_headers, host_supplied_credentials, + is_sigv4_computed_header, resolve_credentials, sign_bedrock_post, + }; + + let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else { + return Ok(request.upstream_headers.clone()); + }; + // Reattaching a header the signer also emits would put both copies on the + // wire, and Bedrock rejects that pair. Python instead drops the caller's + // copy and prefers a forwarded Authorization over the signature, so leave + // the request to Python rather than serving it a different way here. + if request + .upstream_headers + .iter() + .any(|(name, _)| is_sigv4_computed_header(name)) + { + return Err(CoreError::Unsupported( + "request forwards a header AWS SigV4 computes", + )); + } + let env_lookup = |key: &str| std::env::var(key).ok(); + let unsigned: BTreeMap = request.upstream_headers.iter().cloned().collect(); + // A host with its own resolution chain hands the result down; only fall + // back to deriving credentials here when it supplied none. + let credentials = match host_supplied_credentials(&request.optional_params) { + Some(credentials) => credentials, + None => { + resolve_credentials( + aws_auth_config(&request.optional_params, &env_lookup), + &env_lookup, + ) + .await? + } + }; + let signature = sign_bedrock_post( + &request.url, + body, + &aws_signature_headers(&unsigned), + region, + &credentials, + SystemTime::now(), + )?; + // Every original header goes back on the wire alongside the computed ones, + // as Python reattaches them. The guard above already rejected the names + // that would collide, so no name appears twice. + Ok(unsigned.into_iter().chain(signature).collect()) +} + +#[cfg(not(feature = "bedrock-auth"))] +pub(super) async fn signed_headers( + request: &ProviderChatCompletionsRequest, + _body: &[u8], +) -> CoreResult> { + match &request.auth { + ChatCompletionsAuth::AwsSigV4 { .. } => Err(CoreError::Unsupported( + "AWS SigV4 requires the bedrock-auth feature", + )), + _ => Ok(request.upstream_headers.clone()), + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs new file mode 100644 index 00000000000..f30ac1a24bf --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -0,0 +1,59 @@ +//! The `/chat/completions` call, the Rust equivalent of Python's +//! `litellm.completion()`. +//! +//! [`chat_completions`] is the top-level entrypoint: give it a model, the +//! OpenAI-shaped message list, the provider-mapped optional params, and +//! credentials, and it resolves the provider, translates the conversation, +//! calls the provider, and returns a typed OpenAI-shaped response. + +mod client; +mod common_utils; +pub mod conversation; +pub(crate) mod handler; +mod prepare; +pub mod response_utils; +pub mod transformation; +pub mod types; + +use serde_json::{Map, Value}; + +use crate::error::CoreResult; + +use handler::execute_chat_completions_provider_call; +use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config}; +use types::{ChatCompletionsRequest, ChatCompletionsResponse}; + +pub async fn chat_completions( + request: ChatCompletionsRequest<'_>, +) -> CoreResult { + execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await +} + +/// Whether the core would accept this request, without resolving credentials or +/// touching the network. +/// +/// A host that keeps the Python implementation asks this first so it can emit +/// its pre-call logging exactly once, on whichever path is about to run. +/// Returns the decline reason, or `None` when the request is accepted. +pub fn chat_completions_decline_reason( + model: &str, + custom_llm_provider: Option<&str>, + messages: Value, + optional_params: &Map, +) -> Option<&'static str> { + let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else { + return Some("provider is not on the rust chat completions path"); + }; + let Ok(messages) = parse_messages(messages) else { + return Some("unreadable message list"); + }; + if messages.is_empty() { + return Some("empty message list"); + } + config + .unsupported_reason(&messages, optional_params) + .map(|reason| reason.0) +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs new file mode 100644 index 00000000000..1e1c8d1bafd --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -0,0 +1,118 @@ +use serde_json::Value; + +use crate::error::{CoreError, CoreResult}; +use crate::http_utils::has_header; +use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; + +use super::common_utils::{chat_completions_provider_config, string_headers}; +use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig}; +use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest}; + +pub(super) fn resolve_provider_config<'a>( + model: &'a str, + custom_llm_provider: Option<&'a str>, +) -> CoreResult<(String, &'static dyn ChatCompletionsProviderConfig)> { + let provider_info = get_custom_llm_provider(model, custom_llm_provider) + .or_else(|| { + custom_llm_provider.map(|provider| CustomLlmProvider { + model, + custom_llm_provider: provider, + }) + }) + .ok_or_else(|| { + CoreError::InvalidProvider( + "unable to resolve custom_llm_provider for chat completions request".to_string(), + ) + })?; + let config = chat_completions_provider_config(provider_info.custom_llm_provider) + .ok_or_else(|| CoreError::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; + Ok((provider_info.model.to_string(), config)) +} + +pub(super) fn parse_messages(messages: Value) -> CoreResult> { + serde_json::from_value(messages).map_err(|err| { + CoreError::InvalidRequest(format!("invalid chat completions messages: {err}")) + }) +} + +pub(super) fn prepare_chat_completions_call( + request: ChatCompletionsRequest<'_>, +) -> CoreResult { + let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?; + let env_lookup = |key: &str| std::env::var(key).ok(); + + let messages = parse_messages(request.messages)?; + if messages.is_empty() { + return Err(CoreError::InvalidRequest( + "chat completions requires at least one message".to_string(), + )); + } + if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) { + return Err(CoreError::Unsupported(reason.0)); + } + + let mut headers = string_headers(request.extra_headers)?; + let auth = config.auth( + request.api_key, + &model, + &request.optional_params, + &env_lookup, + )?; + match &auth { + ChatCompletionsAuth::Header { name, value } => { + // The deployment's credential replaces whatever the caller forwarded + // under the same name, mirroring Python's + // `{**headers, **anthropic_headers}`: letting a request header win + // would let its sender choose the principal the call bills to. + // + // The exception is a scheme the provider hands off to entirely, such + // as an Anthropic OAuth bearer, where Python drops `x-api-key` + // instead of resolving one. Re-adding it there would put the + // credential into a header the host removed on purpose. + if !config.defers_to_forwarded_auth(&headers) { + headers.retain(|(header, _)| !header.eq_ignore_ascii_case(name)); + headers.push(((*name).to_string(), value.clone())); + } + } + ChatCompletionsAuth::Bearer { token } => { + // Bedrock's `get_request_headers` assigns `headers["Authorization"]` + // unconditionally once a bearer token resolves, so the deployment's + // identity outranks whatever the caller forwarded. Keeping the + // caller's would bill and authorize the call as a different + // principal than the same deployment uses on Python. + // + // The `Header` arm below keeps the opposite precedence on purpose: + // Anthropic's transform honours a forwarded OAuth bearer. + headers.retain(|(name, _)| !name.eq_ignore_ascii_case("authorization")); + headers.push(("authorization".to_string(), format!("Bearer {token}"))); + } + // SigV4 signs the serialized body, so the handler adds its headers. + ChatCompletionsAuth::AwsSigV4 { .. } => {} + } + + 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, + &request.optional_params, + &env_lookup, + )?; + let transformed = + config.transform_request(&model, messages, request.optional_params.clone())?; + + Ok(ProviderChatCompletionsRequest { + model, + config, + url, + body: transformed.body, + upstream_headers: headers, + auth, + optional_params: request.optional_params, + timeout: request.timeout, + }) +} diff --git a/litellm-rust/crates/core/src/chat_completions/response_utils.rs b/litellm-rust/crates/core/src/chat_completions/response_utils.rs new file mode 100644 index 00000000000..1ada5d43980 --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/response_utils.rs @@ -0,0 +1,101 @@ +//! Response normalization shared by every chat completions provider config. + +use std::time::{SystemTime, UNIX_EPOCH}; + +use super::types::{ChatCompletionsUsage, PromptTokensDetails}; + +/// OpenAI finish reasons, mirroring Python's `_FINISH_REASON_MAP` for the +/// reasons the providers on this route can emit. Python warns and falls back to +/// `stop` for anything unmapped, so do the same. +const FINISH_REASONS: &[(&str, &str)] = &[ + ("end_turn", "stop"), + ("stop_sequence", "stop"), + ("max_tokens", "length"), + ("refusal", "content_filter"), + ("compaction", "length"), + ("guardrail_intervened", "content_filter"), + ("content_filtered", "content_filter"), + ("content_filter", "content_filter"), + ("stop", "stop"), + ("length", "length"), +]; + +pub fn finish_reason_for(provider_reason: &str) -> &'static str { + FINISH_REASONS + .iter() + .find(|(reason, _)| *reason == provider_reason) + .map_or("stop", |(_, mapped)| *mapped) +} + +/// Python folds cache tokens into `prompt_tokens` and reports the split under +/// `prompt_tokens_details`; mirror that so cost tracking agrees on both paths. +pub fn usage_from_parts( + input_tokens: u64, + output_tokens: u64, + cache_read_tokens: u64, + cache_creation_tokens: u64, +) -> ChatCompletionsUsage { + let prompt_tokens = input_tokens + cache_read_tokens + cache_creation_tokens; + ChatCompletionsUsage { + prompt_tokens, + completion_tokens: output_tokens, + total_tokens: prompt_tokens + output_tokens, + prompt_tokens_details: PromptTokensDetails { + cached_tokens: cache_read_tokens, + cache_creation_tokens, + text_tokens: input_tokens, + }, + } +} + +pub fn unix_now() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |elapsed| elapsed.as_secs()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn maps_every_reason_the_route_can_observe() { + assert_eq!(finish_reason_for("end_turn"), "stop"); + assert_eq!(finish_reason_for("stop_sequence"), "stop"); + assert_eq!(finish_reason_for("max_tokens"), "length"); + assert_eq!(finish_reason_for("refusal"), "content_filter"); + assert_eq!(finish_reason_for("guardrail_intervened"), "content_filter"); + // Converse emits these two, and folding them into `stop` would report a + // filtered completion as a normal one. + assert_eq!(finish_reason_for("content_filtered"), "content_filter"); + assert_eq!(finish_reason_for("content_filter"), "content_filter"); + } + + #[test] + fn defaults_an_unmapped_reason_to_stop_like_python() { + // Python warns and falls back to `stop` for a reason its own map does + // not carry, so only a reason absent from `_FINISH_REASON_MAP` belongs + // here. + assert_eq!(finish_reason_for("something_new"), "stop"); + assert_eq!(finish_reason_for(""), "stop"); + } + + #[test] + fn folds_cache_tokens_into_prompt_tokens() { + let usage = usage_from_parts(10, 4, 7, 3); + assert_eq!(usage.prompt_tokens, 20); + assert_eq!(usage.completion_tokens, 4); + assert_eq!(usage.total_tokens, 24); + assert_eq!(usage.prompt_tokens_details.cached_tokens, 7); + assert_eq!(usage.prompt_tokens_details.cache_creation_tokens, 3); + assert_eq!(usage.prompt_tokens_details.text_tokens, 10); + } + + #[test] + fn reports_raw_input_tokens_when_no_cache_is_involved() { + let usage = usage_from_parts(12, 5, 0, 0); + assert_eq!(usage.prompt_tokens, 12); + assert_eq!(usage.total_tokens, 17); + assert_eq!(usage.prompt_tokens_details.text_tokens, 12); + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs new file mode 100644 index 00000000000..e2383723cb0 --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -0,0 +1,820 @@ +use serde_json::{Map, Value, json}; + +use crate::error::CoreError; + +use super::prepare::prepare_chat_completions_call; +use super::transformation::ChatCompletionsAuth; +use super::types::ChatCompletionsRequest; + +fn request<'a>( + model: &'a str, + provider: Option<&'a str>, + messages: Value, + optional_params: Value, +) -> ChatCompletionsRequest<'a> { + ChatCompletionsRequest { + model, + messages, + optional_params: match optional_params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }, + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: provider, + extra_headers: None, + timeout: None, + } +} + +/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers +/// carry resolved credentials), so unwrap the failure case by hand. +fn decline(request: ChatCompletionsRequest<'_>) -> CoreError { + match prepare_chat_completions_call(request) { + Err(error) => error, + Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), + } +} + +#[test] +fn resolves_the_provider_from_the_model_prefix() { + let prepared = prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .expect("prepares"); + assert_eq!(prepared.model, "claude-sonnet-4-5"); + assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages"); + assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5")); +} + +#[test] +fn strips_an_explicit_provider_prefix_from_the_model() { + let prepared = prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .expect("prepares"); + assert_eq!(prepared.model, "claude-sonnet-4-5"); +} + +#[test] +fn adds_the_auth_and_default_headers() { + let prepared = prepare_chat_completions_call(request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .expect("prepares"); + assert!( + prepared + .upstream_headers + .contains(&("x-api-key".to_string(), "sk-test".to_string())) + ); + assert!( + prepared + .upstream_headers + .contains(&("anthropic-version".to_string(), "2023-06-01".to_string())) + ); + assert!(matches!( + prepared.auth, + ChatCompletionsAuth::Header { + name: "x-api-key", + .. + } + )); +} + +#[test] +fn the_deployment_credential_replaces_a_caller_supplied_auth_header() { + // Python builds `{**headers, **anthropic_headers}`, so the deployment's key + // overwrites a forwarded one. Honouring the caller's would let whoever sends + // the request choose the Anthropic principal it bills to. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([( + "X-Api-Key".to_string(), + json!("sk-caller"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .collect(); + assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys[0].1, "sk-test"); +} + +#[test] +fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() { + // Anthropic's `validate_environment` pops `x-api-key` and sets `authorization` + // for an OAuth token, so re-adding the key here would put the credential into + // a header the host removed on purpose. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([ + ( + "Authorization".to_string(), + json!("Bearer sk-ant-oat01-token"), + ), + ("X-Api-Key".to_string(), json!("sk-caller")), + ])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + assert!( + !prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"), + "the resolved key must not be applied over an OAuth bearer, got {:?}", + prepared.upstream_headers + ); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-ant-oat01-token") + ); +} + +#[test] +fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() { + // Only an OAuth bearer replaces the credential. Python sends the deployment's + // `x-api-key` alongside any other forwarded `authorization`, so deferring on + // the mere presence of that header would drop the deployment's auth. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([ + ("Authorization".to_string(), json!("Bearer unrelated")), + ("X-Api-Key".to_string(), json!("sk-caller")), + ])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .collect(); + assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys[0].1, "sk-test"); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer unrelated"), + "the unrelated authorization must survive, got {:?}", + prepared.upstream_headers + ); +} + +#[test] +fn declines_an_unsupported_request_before_resolving_credentials() { + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}), + ); + call.api_key = None; + // No api_key is set and no env is consulted: the gate must run first, so the + // error is the decline rather than a missing-credential error. + assert_eq!(decline(call), CoreError::Unsupported("streaming")); +} + +#[test] +fn rejects_an_unknown_provider() { + assert_eq!( + decline(request( + "openai/gpt-4o", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )), + CoreError::InvalidProvider("openai".to_string()) + ); +} + +#[test] +fn rejects_a_model_with_no_resolvable_provider() { + assert!(matches!( + decline(request( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )), + CoreError::InvalidProvider(_) + )); +} + +#[test] +fn rejects_an_empty_or_malformed_message_list() { + assert_eq!( + decline(request( + "anthropic/claude-sonnet-4-5", + None, + json!([]), + json!({}), + )), + CoreError::InvalidRequest("chat completions requires at least one message".to_string()) + ); + assert!(matches!( + decline(request( + "anthropic/claude-sonnet-4-5", + None, + json!("not a list"), + json!({}), + )), + CoreError::InvalidRequest(_) + )); +} + +#[test] +fn rejects_non_string_extra_headers() { + let mut call = request( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); + assert_eq!( + decline(call), + CoreError::InvalidRequest( + "chat completions extra_headers.x-trace must be a string, got number".to_string() + ) + ); +} + +#[cfg(feature = "bedrock-auth")] +#[test] +fn prepares_a_bedrock_call_without_resolving_credentials() { + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + ); + call.api_key = None; + let prepared = prepare_chat_completions_call(call).expect("prepares"); + assert_eq!( + prepared.url, + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" + ); + assert_eq!( + prepared.auth, + ChatCompletionsAuth::AwsSigV4 { + region: "us-east-1".to_string() + } + ); + // SigV4 signs the serialized body, so prepare must not have added an + // Authorization header; the handler does it. + assert!( + !prepared + .upstream_headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")) + ); + assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16})); +} + +#[cfg(feature = "bedrock-auth")] +#[tokio::test] +async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { + // Python signs only the AWS header set and reattaches the rest, so a header + // the caller forwarded rides along without joining the canonical request. + // Signing it makes Converse 403 on a deployment that works on Python. + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" + }), + ); + // A key would resolve to a bearer token and never reach the signer. + call.api_key = None; + call.extra_headers = Some(Map::from_iter([( + "x-request-id".to_string(), + json!("abc-123"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#) + .await + .expect("signs"); + + let authorization = signed + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("authorization")) + .map(|(_, value)| value.clone()) + .expect("carries an authorization header"); + assert!( + authorization.starts_with("AWS4-HMAC-SHA256"), + "expected a SigV4 signature, got {authorization}" + ); + assert!( + !authorization.contains("x-request-id"), + "forwarded header reached SignedHeaders: {authorization}" + ); + // It still goes on the wire, it is just not part of the signature. + assert!( + signed + .iter() + .any(|(name, value)| name == "x-request-id" && value == "abc-123"), + "forwarded header was dropped instead of reattached" + ); +} + +#[cfg(feature = "bedrock-auth")] +#[tokio::test] +async fn a_forwarded_header_the_signer_computes_declines_to_python() { + // Reattaching the caller's copy next to the computed one puts the name on + // the wire twice and Bedrock rejects the pair, so a request carrying one + // has to go to Python instead of being signed here. + for forwarded in [ + "Authorization", + "x-amz-date", + "x-amz-security-token", + "Date", + ] { + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" + }), + ); + call.api_key = None; + call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#) + .await + .expect_err("{forwarded} should decline instead of being signed"); + assert!( + matches!(error, CoreError::Unsupported(_)), + "{forwarded} declined as {error:?}, which the host would not fall back on" + ); + } +} + +#[cfg(feature = "bedrock-auth")] +#[test] +fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() { + // `get_request_headers` assigns `headers["Authorization"]` unconditionally + // once a bearer token resolves, so the deployment's identity wins on + // Python. Keeping the caller's would authorize and bill the call as a + // different principal, and only when the deployment carries `rust: true`. + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + ); + call.extra_headers = Some(Map::from_iter([( + "Authorization".to_string(), + json!("Bearer caller-supplied"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let authorizations: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("authorization")) + .map(|(_, value)| value.as_str()) + .collect(); + assert_eq!( + authorizations, + vec!["Bearer sk-test"], + "the deployment token must be the only authorization on the wire" + ); +} + +#[test] +fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() { + // The opposite precedence, and deliberate: Anthropic's own transform + // honours a forwarded OAuth bearer, so the Bedrock fix above must not be + // generalized into a rule that the configured key always wins. + // + // An OAuth bearer is the whole of that exception. This forwarded a plain + // `x-api-key` until round 17, which read as the same claim and was not: + // Python overwrites a forwarded `x-api-key` with the deployment's. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([( + "authorization".to_string(), + json!("Bearer sk-ant-oat01-forwarded"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .map(|(_, value)| value.as_str()) + .collect(); + assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-ant-oat01-forwarded") + ); +} + +#[cfg(feature = "bedrock-auth")] +#[test] +fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { + // The configured bearer identity has its own account and quota boundary, + // so a request carrying one must not be signed as whatever principal the + // host's AWS credentials resolve to. + let prepared = prepare_chat_completions_call(request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + )) + .expect("prepares"); + assert_eq!( + prepared.auth, + ChatCompletionsAuth::Bearer { + token: "sk-test".to_string() + } + ); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-test"), + "prepare did not carry the bearer token" + ); +} + +fn decline_reason( + model: &str, + provider: Option<&str>, + messages: Value, + params: Value, +) -> Option<&'static str> { + let params = match params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }; + super::chat_completions_decline_reason(model, provider, messages, ¶ms) +} + +#[test] +fn the_gate_accepts_what_prepare_accepts() { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + ), + None + ); +} + +#[test] +fn the_gate_declines_without_resolving_credentials_or_calling_out() { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}), + ), + Some("streaming") + ); + assert_eq!( + decline_reason( + "openai/gpt-4o", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ), + Some("provider is not on the rust chat completions path") + ); + assert_eq!( + decline_reason( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ), + Some("provider is not on the rust chat completions path") + ); + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!("nope"), + json!({}) + ), + Some("unreadable message list") + ); + assert_eq!( + decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})), + Some("empty message list") + ); +} + +#[test] +fn the_gate_agrees_with_prepare_on_every_case_it_accepts() { + // A gate that accepts what prepare then declines would make the host emit + // its pre-call logging on a path that falls back, so pin the agreement. + for (messages, params) in [ + ( + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 8}), + ), + ( + json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]), + json!({"temperature": 0.1}), + ), + ( + json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]), + json!({}), + ), + ] { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + messages.clone(), + params.clone() + ), + None, + "gate declined {messages}" + ); + prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + None, + messages.clone(), + params, + )) + .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); + } +} + +mod round_trip { + use super::*; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + + use crate::chat_completions::chat_completions; + + 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 http_response(status: &str, body: &str) -> String { + format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ) + } + + /// Serve one request from a stub upstream and hand back what it received. + async fn serve_once( + status: &'static str, + body: &'static str, + ) -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let port = listener.local_addr().expect("addr").port(); + let handle = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts"); + let received = read_http_request(&mut socket).await; + socket + .write_all(http_response(status, body).as_bytes()) + .await + .expect("writes response"); + socket.flush().await.expect("flushes"); + received + }); + (format!("http://127.0.0.1:{port}/v1/messages"), handle) + } + + fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { + ChatCompletionsRequest { + model: "anthropic/claude-sonnet-4-5", + messages, + optional_params: match params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }, + api_key: Some("sk-test"), + api_base: Some(api_base), + custom_llm_provider: None, + extra_headers: None, + timeout: Some(std::time::Duration::from_secs(10)), + } + } + + const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; + + #[tokio::test] + async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { + let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; + let response = chat_completions(call( + &api_base, + json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]), + json!({"max_tokens": 16}), + )) + .await + .expect("call succeeds"); + + let received = handle.await.expect("server task"); + let sent: Value = serde_json::from_str( + received + .split_once("\r\n\r\n") + .expect("request has a body") + .1, + ) + .expect("body is json"); + assert_eq!( + sent["messages"], + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) + ); + assert_eq!( + sent["system"], + json!([{"type": "text", "text": "be terse"}]) + ); + assert_eq!(sent["max_tokens"], json!(16)); + assert!(received.to_lowercase().contains("x-api-key: sk-test")); + + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello") + ); + assert_eq!(response.usage.total_tokens, 15); + } + + #[tokio::test] + async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { + // The provider was called and billed, so the host must not retry this + // on its own path. `MissingField` here would read as a pre-send + // decline and be retried; `InvalidResponse` cannot. + const NO_USAGE: &str = + r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; + let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("response cannot be normalized"); + handle.await.expect("server task"); + assert!( + matches!(err, CoreError::InvalidResponse(_)), + "expected a post-send error, got {err:?}" + ); + } + + #[tokio::test] + async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { + const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; + let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("response cannot be normalized"); + handle.await.expect("server task"); + assert!( + matches!(err, CoreError::InvalidResponse(_)), + "expected a post-send error, got {err:?}" + ); + } + + #[tokio::test] + async fn an_upstream_error_status_keeps_its_code() { + let (api_base, handle) = + serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("upstream rejects"); + handle.await.expect("server task"); + assert!( + matches!(err, CoreError::Http { status: 429, .. }), + "expected a 429, got {err:?}" + ); + } + + #[tokio::test] + async fn a_connection_that_is_never_established_declines_instead_of_failing() { + // Nothing was sent, so nothing was billed and the host can still serve + // the request. Classing this with the post-send failures would turn a + // recoverable fallback into a user-facing error on exactly the + // deployments whose transport is configured only on the Python client. + let port = { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + listener.local_addr().expect("has an address").port() + // Dropped here, so the port is closed and the connect is refused. + }; + let err = chat_completions(call( + &format!("http://127.0.0.1:{port}/v1/messages"), + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("nothing is listening"); + assert!( + matches!(err, CoreError::Connect(_)), + "expected a pre-send connect failure, got {err:?}" + ); + } + + #[test] + fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { + use crate::chat_completions::handler::as_response_error; + + for original in [ + CoreError::MissingField("usage"), + CoreError::Unsupported("non-text response content block"), + CoreError::InvalidRequest("whatever".to_string()), + CoreError::Auth("whatever".to_string()), + ] { + let label = format!("{original:?}"); + assert!( + matches!(as_response_error(original), CoreError::InvalidResponse(_)), + "{label} must not stay retryable once the provider has answered" + ); + } + // An upstream status is already unambiguous, so it survives intact. + assert!(matches!( + as_response_error(CoreError::Http { + status: 500, + body: "boom".to_string() + }), + CoreError::Http { status: 500, .. } + )); + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs new file mode 100644 index 00000000000..a30ce9dc77c --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -0,0 +1,155 @@ +use serde_json::{Map, Value}; + +use crate::error::CoreResult; + +use super::types::{ + ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData, + ProviderChatResponseData, +}; + +/// How the upstream call is authenticated. API-key strategies are resolved in +/// `prepare`; SigV4 needs the serialized body, so the handler signs it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ChatCompletionsAuth { + Header { name: &'static str, value: String }, + Bearer { token: String }, + AwsSigV4 { region: String }, +} + +/// Why a request cannot be served by the Rust path. +/// +/// The core declines rather than guessing: the host turns this into a +/// transparent fallback to the Python implementation, which covers the full +/// surface. Acceptance is an allowlist, so a parameter or message shape the +/// core has never seen declines by construction instead of being translated +/// wrong. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Unsupported(pub &'static str); + +pub const STREAM_PARAM: &str = "stream"; + +/// Message fields that carry no meaning for the upstream body, so their +/// presence does not make a request untranslatable. +const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; + +pub trait ChatCompletionsProviderConfig: Sync { + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn auth( + &self, + api_key: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + &[("content-type", "application/json")] + } + + /// Whether an auth header the caller already supplied is the credential this + /// request should authenticate with, so the resolved one is not applied. + /// + /// Defaults to false: the deployment's credential outranks anything + /// forwarded, which is what every provider wants for its own auth header. + /// A provider overrides this only for a scheme it hands off to entirely. + fn defers_to_forwarded_auth(&self, _headers: &[(String, String)]) -> bool { + false + } + + /// Provider parameter names (post-mapping) the Rust path knows how to place + /// in the upstream body. Anything outside this set declines the request. + fn supported_params(&self) -> &'static [&'static str]; + + /// Parameters consumed as call configuration (credentials, endpoints) + /// rather than placed in the body. Accepted, never serialized. + fn config_params(&self) -> &'static [&'static str] { + &[] + } + + fn unsupported_reason( + &self, + messages: &[ChatMessage], + optional_params: &Map, + ) -> Option { + unsupported_param( + self.supported_params(), + self.config_params(), + optional_params, + ) + .or_else(|| messages.iter().find_map(unsupported_message)) + } + + fn transform_request( + &self, + model: &str, + messages: Vec, + optional_params: Map, + ) -> CoreResult; + + fn transform_response( + &self, + model: &str, + response: ProviderChatResponseData, + ) -> CoreResult; +} + +pub fn unsupported_param( + supported: &'static [&'static str], + config: &'static [&'static str], + optional_params: &Map, +) -> Option { + if optional_params + .get(STREAM_PARAM) + .and_then(Value::as_bool) + .unwrap_or(false) + { + return Some(Unsupported("streaming")); + } + optional_params + .keys() + .any(|key| { + key != STREAM_PARAM + && !supported.contains(&key.as_str()) + && !config.contains(&key.as_str()) + }) + .then_some(Unsupported("unrecognized request parameter")) +} + +/// Message shapes the core can translate faithfully: text content, either a +/// plain string or a non-empty list of parts that are all +/// `{"type": "text", "text": ...}`. Tool calls, tool results, and multimodal +/// parts decline so Python's fuller translation handles them. +pub fn unsupported_message(message: &ChatMessage) -> Option { + if message + .extra + .keys() + .any(|key| !IGNORABLE_MESSAGE_FIELDS.contains(&key.as_str())) + { + return Some(Unsupported("unrecognized message field")); + } + if !matches!(message.role.as_str(), "system" | "user" | "assistant") { + return Some(Unsupported("unrecognized message role")); + } + match &message.content { + None => Some(Unsupported("message without content")), + Some(ChatMessageContent::Text(_)) => None, + Some(ChatMessageContent::Parts(parts)) if parts.is_empty() => { + Some(Unsupported("message without content")) + } + Some(ChatMessageContent::Parts(parts)) => parts + .iter() + .any(|part| { + part.get("type").and_then(Value::as_str) != Some("text") + || part.get("text").and_then(Value::as_str).is_none() + || part.as_object().is_some_and(|object| object.len() != 2) + }) + .then_some(Unsupported("non-text message content")), + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs new file mode 100644 index 00000000000..35dd543a986 --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -0,0 +1,112 @@ +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig}; + +/// A `/chat/completions` call as it crosses into the core. +/// +/// `optional_params` arrives already mapped to the provider's own parameter +/// names by the host, exactly as the messages route receives an already +/// Anthropic-shaped body. The core owns the conversation translation, the +/// provider call, and the response normalization. +pub struct ChatCompletionsRequest<'a> { + pub model: &'a str, + pub messages: Value, + pub optional_params: Map, + 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(super) struct ProviderChatCompletionsRequest { + pub(super) model: String, + pub(super) config: &'static dyn ChatCompletionsProviderConfig, + pub(super) url: String, + pub(super) body: Value, + pub(super) upstream_headers: Vec<(String, String)>, + pub(super) auth: ChatCompletionsAuth, + #[cfg_attr(not(feature = "bedrock-auth"), allow(dead_code))] + pub(super) optional_params: Map, + pub(super) timeout: Option, +} + +/// The provider-shaped request body a config produces. Named rather than a bare +/// `Value` so the transform contract stays a typed one, mirroring +/// [`crate::audio_transcription::types::AudioTranscriptionRequestData`]. +pub struct ProviderChatRequestData { + pub body: Value, +} + +/// The raw provider response body handed back to a config for normalization. +pub struct ProviderChatResponseData { + pub body: Value, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum ChatMessageContent { + Text(String), + Parts(Vec), +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ChatMessage { + pub role: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(flatten)] + pub extra: Map, +} + +/// OpenAI `usage`, including the `prompt_tokens_details` split LiteLLM's Python +/// path reports so cost tracking sees the same numbers on either path. +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct PromptTokensDetails { + pub cached_tokens: u64, + pub cache_creation_tokens: u64, + pub text_tokens: u64, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct ChatCompletionsUsage { + pub prompt_tokens: u64, + pub completion_tokens: u64, + pub total_tokens: u64, + pub prompt_tokens_details: PromptTokensDetails, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ChatCompletionsChoiceMessage { + pub role: String, + // Whether an empty turn is `None` or `""` is the provider's choice, not a + // shared invariant: Anthropic's transform ends on `merged_text or None` + // while Converse assigns the joined string unconditionally. Each config + // mirrors its own, so keep this optional and serialize it even when None. + pub content: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ChatCompletionsChoice { + pub index: u64, + pub message: ChatCompletionsChoiceMessage, + pub finish_reason: String, +} + +/// The normalized response handed back to the host. +/// +/// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the +/// `ModelResponse` it already created, and echoing the provider's own id here +/// would change it. Pinned by `response_carries_no_id` in `tests.rs`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ChatCompletionsResponse { + pub created: u64, + pub model: String, + pub choices: Vec, + pub usage: ChatCompletionsUsage, +} diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index caada1d98b0..e1ac0a4fc8f 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -12,8 +12,30 @@ pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10; /// Max characters of an upstream error body echoed across the call boundary /// before truncation, so provider bodies are bounded and data-minimized. -pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256; +pub(crate) const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256; /// Provider name used for Anthropic Messages when a deployment's provider model /// does not carry an explicit provider prefix. pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic"; + +/// Prefix identifying an Anthropic OAuth token. Mirrors Python's +/// `ANTHROPIC_OAUTH_TOKEN_PREFIX`, which is what makes `validate_environment` +/// authenticate with `authorization` and drop `x-api-key` entirely. +pub(crate) const ANTHROPIC_OAUTH_TOKEN_PREFIX: &str = "sk-ant-oat"; + +/// Full-request timeout ceiling for chat completions provider calls, in +/// seconds. Mirrors the Python chat completions default. +pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600; + +/// Connect timeout for chat completions provider calls, in seconds. +pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10; + +/// `object` field every non-streaming chat completion response carries. +pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion"; + +/// Placeholder Python substitutes for empty or whitespace-only message text, +/// which Anthropic and Bedrock both reject. Must match +/// `_EMPTY_TEXT_PLACEHOLDER` in +/// `litellm/litellm_core_utils/prompt_templates/factory.py`. +pub const EMPTY_TEXT_PLACEHOLDER: &str = + "[System: Empty message content sanitised to satisfy protocol]"; diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index c2b08eee0c0..739532f8cb5 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -23,8 +23,19 @@ pub enum CoreError { Http { status: u16, body: String }, #[error("upstream network error: {0}")] Network(String), + /// The provider was never reached: DNS, TCP, TLS or proxy setup failed + /// before any byte of the request went out. Nothing was billed, so a host + /// that keeps a reference implementation can serve the request itself. + /// A timeout is deliberately not this, since the provider may have received + /// and answered the request already. + #[error("could not reach the provider: {0}")] + Connect(String), #[error("routing error: {0}")] Routing(String), + /// The request is outside the surface this route covers in Rust. Hosts that + /// keep a reference implementation treat this as "fall back", not "fail". + #[error("unsupported by the rust path: {0}")] + Unsupported(&'static str), } pub fn json_type_name(value: &serde_json::Value) -> &'static str { diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs new file mode 100644 index 00000000000..c541f50275b --- /dev/null +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -0,0 +1,112 @@ +//! Header and upstream-body helpers shared by every route module. + +use serde_json::{Map, Value}; + +use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; +use crate::error::{CoreError, CoreResult, json_type_name}; + +/// Bound an upstream error body before it crosses a host boundary, so provider +/// bodies stay data-minimized. +pub fn truncate_error_body(body: &str) -> String { + if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS { + return body.to_string(); + } + let truncated: String = body.chars().take(UPSTREAM_ERROR_BODY_MAX_CHARS).collect(); + format!("{truncated}... (truncated)") +} + +pub fn string_headers( + context: &'static str, + 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!( + "{context} extra_headers.{key} must be a string, got {}", + json_type_name(&value) + )) + }) + }) + .collect() +} + +pub fn has_header(headers: &[(String, String)], name: &str) -> bool { + headers + .iter() + .any(|(key, _)| key.eq_ignore_ascii_case(name)) +} + +pub fn has_bearer_auth(headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, value)| { + if !name.eq_ignore_ascii_case("authorization") { + return false; + } + let value = value.trim(); + value.len() > 7 + && value[..7].eq_ignore_ascii_case("bearer ") + && !value[7..].trim().is_empty() + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn truncate_leaves_short_bodies_untouched() { + assert_eq!(truncate_error_body("short"), "short"); + } + + #[test] + fn truncate_bounds_long_bodies_by_characters() { + let body = "\u{00e9}".repeat(UPSTREAM_ERROR_BODY_MAX_CHARS + 10); + let truncated = truncate_error_body(&body); + assert!(truncated.ends_with("... (truncated)")); + assert_eq!( + truncated.chars().count(), + UPSTREAM_ERROR_BODY_MAX_CHARS + "... (truncated)".chars().count() + ); + } + + #[test] + fn string_headers_rejects_non_string_values() { + let headers = Map::from_iter([("x-trace".to_string(), json!(7))]); + let err = string_headers("chat completions", Some(headers)).expect_err("non-string value"); + assert_eq!( + err, + CoreError::InvalidRequest( + "chat completions extra_headers.x-trace must be a string, got number".to_string() + ) + ); + } + + #[test] + fn header_lookup_is_case_insensitive() { + let headers = vec![("X-Api-Key".to_string(), "k".to_string())]; + assert!(has_header(&headers, "x-api-key")); + assert!(!has_header(&headers, "authorization")); + } + + #[test] + fn bearer_detection_requires_a_non_empty_token() { + assert!(has_bearer_auth(&[( + "Authorization".to_string(), + "Bearer abc".to_string() + )])); + assert!(!has_bearer_auth(&[( + "Authorization".to_string(), + "Bearer ".to_string() + )])); + assert!(!has_bearer_auth(&[( + "Authorization".to_string(), + "Basic abc".to_string() + )])); + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 51ea19750ea..dce4a425ea0 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,8 +1,10 @@ pub mod audio_transcription; pub mod caching; pub mod call_lifecycle; +pub mod chat_completions; pub mod constants; pub mod error; +pub mod http_utils; pub mod messages; pub mod ocr; pub mod providers; diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 9dcfcaa71e3..a14dffbc1fe 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,19 +1,15 @@ use serde_json::{Map, Value}; -use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS; -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::CoreResult; +use crate::http_utils::string_headers as shared_string_headers; use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; use super::transformation::AnthropicMessagesProviderConfig; -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) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body}; + +const HEADER_CONTEXT: &str = "messages"; pub(super) fn messages_provider_config( provider: &str, @@ -28,37 +24,5 @@ pub(super) fn messages_provider_config( 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)) -} - -pub(super) fn has_bearer_auth(headers: &[(String, String)]) -> bool { - headers.iter().any(|(name, value)| { - if !name.eq_ignore_ascii_case("authorization") { - return false; - } - let value = value.trim(); - value.len() > 7 - && value[..7].eq_ignore_ascii_case("bearer ") - && !value[7..].trim().is_empty() - }) + shared_string_headers(HEADER_CONTEXT, extra_headers) } diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/mod.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs new file mode 100644 index 00000000000..4534ac0182c --- /dev/null +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs @@ -0,0 +1,444 @@ +use super::*; +use serde_json::json; + +fn messages(value: Value) -> Vec { + serde_json::from_value(value).expect("valid messages") +} + +fn params(value: Value) -> Map { + match value { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + } +} + +fn transform(model: &str, msgs: Value, opts: Value) -> Value { + ANTHROPIC_CHAT_COMPLETIONS_CONFIG + .transform_request(model, messages(msgs), params(opts)) + .expect("request transforms") + .body +} + +fn transform_response(body: Value) -> CoreResult { + ANTHROPIC_CHAT_COMPLETIONS_CONFIG + .transform_response("claude-sonnet-4-5", ProviderChatResponseData { body }) +} + +fn reason(msgs: Value, opts: Value) -> Option { + ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) +} + +#[test] +fn builds_the_messages_body_python_builds() { + let body = transform( + "claude-sonnet-4-5", + json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]), + json!({"max_tokens": 128, "temperature": 0.2}), + ); + assert_eq!( + body, + json!({ + "model": "claude-sonnet-4-5", + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]} + ], + "system": [{"type": "text", "text": "be terse"}], + "max_tokens": 128, + "temperature": 0.2 + }) + ); +} + +#[test] +fn omits_system_when_no_system_message_is_present() { + let body = transform( + "claude-sonnet-4-5", + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + ); + assert!(body.get("system").is_none()); +} + +#[test] +fn merges_consecutive_turns_and_wraps_every_text_in_a_block() { + let body = transform( + "claude-sonnet-4-5", + json!([ + {"role": "user", "content": "one"}, + {"role": "user", "content": [{"type": "text", "text": "two"}]}, + {"role": "assistant", "content": "ack"} + ]), + json!({"max_tokens": 16}), + ); + assert_eq!( + body["messages"], + json!([ + {"role": "user", "content": [ + {"type": "text", "text": "one"}, + {"type": "text", "text": "two"} + ]}, + {"role": "assistant", "content": [{"type": "text", "text": "ack"}]} + ]) + ); +} + +#[test] +fn right_strips_a_trailing_assistant_prefill_like_python() { + let body = transform( + "claude-sonnet-4-5", + json!([ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "Argentina "} + ]), + json!({"max_tokens": 16}), + ); + assert_eq!( + body["messages"][1]["content"][0]["text"], + json!("Argentina") + ); +} + +#[test] +fn passes_every_supported_param_through_untouched() { + let body = transform( + "claude-sonnet-4-5", + json!([{"role": "user", "content": "hi"}]), + json!({ + "max_tokens": 64, + "temperature": 0.1, + "top_p": 0.9, + "stop_sequences": ["STOP"] + }), + ); + assert_eq!(body["max_tokens"], json!(64)); + assert_eq!(body["temperature"], json!(0.1)); + assert_eq!(body["top_p"], json!(0.9)); + assert_eq!(body["stop_sequences"], json!(["STOP"])); +} + +#[test] +fn declines_top_k_because_python_gates_it_by_model_below_this_point() { + // `temperature` and `top_p` arrive already resolved, because + // `map_openai_params` applies `_apply_sampling_param` to them before the + // gate runs. `top_k` bypasses that and is gated inside `transform_request`, + // the function this route replaces, so forwarding it would send `top_k` to + // a model that removed sampling params and take a 400 after the call, where + // Python drops it and succeeds. + assert_eq!( + reason( + json!([{"role": "user", "content": "hi"}]), + json!({"top_k": 40}) + ), + Some(Unsupported("unrecognized request parameter")) + ); +} + +#[test] +fn declines_streaming_before_anything_else() { + assert_eq!( + reason( + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true, "max_tokens": 16}) + ), + Some(Unsupported("streaming")) + ); +} + +#[test] +fn accepts_an_explicit_stream_false() { + assert_eq!( + reason( + json!([{"role": "user", "content": "hi"}]), + json!({"stream": false, "max_tokens": 16}) + ), + None + ); +} + +#[test] +fn declines_any_param_outside_the_allowlist() { + for param in [ + json!({"tools": []}), + json!({"tool_choice": {"type": "auto"}}), + json!({"thinking": {"type": "enabled"}}), + json!({"system": "injected"}), + json!({"metadata": {"user_id": "u1"}}), + json!({"output_config": {"effort": "high"}}), + ] { + assert_eq!( + reason(json!([{"role": "user", "content": "hi"}]), param.clone()), + Some(Unsupported("unrecognized request parameter")), + "expected {param} to decline" + ); + } +} + +#[test] +fn declines_tool_calls_tool_results_and_multimodal_content() { + assert_eq!( + reason( + json!([ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": null, "tool_calls": [ + {"id": "c1", "type": "function", + "function": {"name": "f", "arguments": "{}"}} + ]} + ]), + json!({}) + ), + Some(Unsupported("unrecognized message field")) + ); + assert_eq!( + reason( + json!([ + {"role": "user", "content": "hi"}, + {"role": "tool", "tool_call_id": "c1", "content": "ok"} + ]), + json!({}) + ), + Some(Unsupported("unrecognized message field")) + ); + assert_eq!( + reason( + json!([{"role": "user", "content": [ + {"type": "image_url", "image_url": {"url": "https://x/y.png"}} + ]}]), + json!({}) + ), + Some(Unsupported("non-text message content")) + ); + assert_eq!( + reason( + json!([{"role": "user", "content": [ + {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}} + ]}]), + json!({}) + ), + Some(Unsupported("non-text message content")) + ); +} + +#[test] +fn declines_a_message_whose_content_list_is_empty() { + // An empty list passes every per-part check, so without this it would reach + // the provider as an empty `content` array and fail after the call rather + // than declining to Python before it. + assert_eq!( + reason(json!([{"role": "user", "content": []}]), json!({})), + Some(Unsupported("message without content")) + ); + assert_eq!( + reason( + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]), + json!({}) + ), + None + ); +} + +#[test] +fn declines_a_conversation_that_does_not_open_on_a_user_turn() { + assert_eq!( + reason( + json!([ + {"role": "system", "content": "be terse"}, + {"role": "assistant", "content": "prefill"} + ]), + json!({}) + ), + Some(Unsupported("conversation does not open on a user turn")) + ); +} + +#[test] +fn accepts_a_plain_text_conversation() { + assert_eq!( + reason( + json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": [{"type": "text", "text": "again"}]} + ]), + json!({"max_tokens": 16, "temperature": 0.5}) + ), + None + ); +} + +#[test] +fn normalizes_a_text_response_into_openai_shape() { + let response = transform_response(json!({ + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20260101", + "content": [{"type": "text", "text": "hello"}, {"type": "text", "text": " there"}], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": {"input_tokens": 11, "output_tokens": 4} + })) + .expect("response transforms"); + + assert_eq!(response.model, "claude-sonnet-4-5-20260101"); + assert_eq!(response.choices.len(), 1); + assert_eq!(response.choices[0].index, 0); + assert_eq!(response.choices[0].message.role, "assistant"); + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello there") + ); + assert_eq!(response.choices[0].finish_reason, "stop"); + assert_eq!(response.usage.prompt_tokens, 11); + assert_eq!(response.usage.completion_tokens, 4); + assert_eq!(response.usage.total_tokens, 15); +} + +#[test] +fn folds_cache_tokens_into_prompt_tokens_like_python() { + let response = transform_response(json!({ + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 10, + "output_tokens": 2, + "cache_read_input_tokens": 5, + "cache_creation_input_tokens": 3 + } + })) + .expect("response transforms"); + assert_eq!(response.usage.prompt_tokens, 18); + assert_eq!(response.usage.total_tokens, 20); + assert_eq!(response.usage.prompt_tokens_details.cached_tokens, 5); + assert_eq!( + response.usage.prompt_tokens_details.cache_creation_tokens, + 3 + ); + assert_eq!(response.usage.prompt_tokens_details.text_tokens, 10); +} + +#[test] +fn maps_max_tokens_stop_reason_to_length() { + let response = transform_response(json!({ + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "max_tokens", + "usage": {"input_tokens": 1, "output_tokens": 1} + })) + .expect("response transforms"); + assert_eq!(response.choices[0].finish_reason, "length"); +} + +#[test] +fn a_refusal_returns_the_completion_python_returns() { + // `refusal` is a stop_reason, not a content block type, so the content is + // ordinary text and this normalizes rather than declining. Python maps it + // to content_filter in _FINISH_REASON_MAP and returns the completion. + let response = transform_response(json!({ + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "I can't help with that."}], + "stop_reason": "refusal", + "usage": {"input_tokens": 9, "output_tokens": 6} + })) + .expect("a refusal still transforms"); + assert_eq!(response.choices[0].finish_reason, "content_filter"); + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("I can't help with that.") + ); +} + +#[test] +fn reports_no_content_rather_than_an_empty_string() { + let response = transform_response(json!({ + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 0} + })) + .expect("response transforms"); + assert_eq!(response.choices[0].message.content, None); +} + +#[test] +fn response_carries_no_id_so_python_keeps_its_chatcmpl_id() { + let response = transform_response(json!({ + "id": "msg_should_not_leak", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + })) + .expect("response transforms"); + let value = serde_json::to_value(response).expect("serializable"); + assert!( + value.get("id").is_none(), + "the rust response must not carry an id, got {value}" + ); +} + +#[test] +fn declines_a_response_carrying_a_non_text_block() { + let err = transform_response(json!({ + "model": "claude-sonnet-4-5", + "content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}], + "stop_reason": "tool_use", + "usage": {"input_tokens": 1, "output_tokens": 1} + })) + .expect_err("non-text block"); + assert_eq!( + err, + CoreError::Unsupported("non-text response content block") + ); +} + +#[test] +fn errors_on_a_response_missing_required_fields() { + assert_eq!( + transform_response(json!("nope")).expect_err("not an object"), + CoreError::InvalidResponse("messages response is not an object".to_string()) + ); + assert_eq!( + transform_response(json!({"model": "m", "usage": {}})).expect_err("no content"), + CoreError::MissingField("content") + ); + assert_eq!( + transform_response(json!({"model": "m", "content": []})).expect_err("no usage"), + CoreError::MissingField("usage") + ); + assert_eq!( + transform_response(json!({"content": [], "usage": {}})).expect_err("no model"), + CoreError::MissingField("model") + ); +} + +#[test] +fn resolves_the_messages_url_and_x_api_key_auth() { + let config = &ANTHROPIC_CHAT_COMPLETIONS_CONFIG; + assert_eq!( + config + .complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None) + .expect("url builds"), + "https://api.anthropic.com/v1/messages" + ); + assert_eq!( + config + .auth(Some("sk-x"), "claude-sonnet-4-5", &Map::new(), &|_| None) + .expect("auth resolves"), + ChatCompletionsAuth::Header { + name: "x-api-key", + value: "sk-x".to_string() + } + ); + assert_eq!( + config.default_headers(), + &[ + ("anthropic-version", "2023-06-01"), + ("content-type", "application/json"), + ] + ); +} diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs new file mode 100644 index 00000000000..3658642b539 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -0,0 +1,211 @@ +use serde_json::{Map, Value, json}; + +use crate::chat_completions::conversation::{Conversation, build_conversation}; +use crate::chat_completions::transformation::{ + ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message, + unsupported_param, +}; +use crate::chat_completions::types::{ + ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, ChatMessage, + ProviderChatRequestData, ProviderChatResponseData, +}; +use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX; +use crate::error::{CoreError, CoreResult}; +use crate::providers::anthropic::messages::transformation::{ + complete_anthropic_url, resolve_anthropic_api_key, +}; + +use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts}; + +/// Anthropic parameter names, post `map_openai_params`, that the Rust path can +/// place verbatim in the Messages body. +/// +/// `top_k` is deliberately absent even though the Messages API takes it. +/// `temperature` and `top_p` reach this gate already resolved, because +/// `map_openai_params` runs first and applies `_apply_sampling_param` to them. +/// `top_k` bypasses `map_openai_params` entirely, so Python applies that same +/// per-model gate inside `transform_request`, the function this route replaces. +/// Forwarding it would send `top_k` to a model that removed sampling params and +/// take a 400 after the call, where Python drops it and succeeds. +const SUPPORTED_PARAMS: &[&str] = &["max_tokens", "temperature", "top_p", "stop_sequences"]; + +pub struct AnthropicChatCompletionsConfig; + +pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicChatCompletionsConfig = + AnthropicChatCompletionsConfig; + +fn text_block(text: &str) -> Value { + json!({"type": "text", "text": text}) +} + +fn anthropic_body(model: &str, conversation: &Conversation, params: Map) -> Value { + let messages: Vec = conversation + .turns + .iter() + .map(|turn| { + json!({ + "role": turn.role.as_str(), + "content": turn.texts.iter().map(|text| text_block(text)).collect::>(), + }) + }) + .collect(); + + let system: Vec = conversation.system.iter().map(|s| text_block(s)).collect(); + + let body = Map::from_iter( + [ + ("model".to_string(), json!(model)), + ("messages".to_string(), json!(messages)), + ] + .into_iter() + // Python builds `{"model", "messages", **optional_params}` with + // `system` already folded into optional_params, so a caller-supplied + // key of the same name wins here too. + .chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))) + .chain(params), + ); + Value::Object(body) +} + +impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + Ok(complete_anthropic_url(api_base, env_lookup)) + } + + fn auth( + &self, + api_key: Option<&str>, + _model: &str, + _optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + Ok(ChatCompletionsAuth::Header { + name: "x-api-key", + value: resolve_anthropic_api_key(api_key, env_lookup)?, + }) + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + &[ + ("anthropic-version", "2023-06-01"), + ("content-type", "application/json"), + ] + } + + /// An OAuth bearer is the whole credential: Python's `validate_environment` + /// authenticates with it and drops `x-api-key` rather than resolving one, so + /// the resolved key must not be applied over the top. Any other forwarded + /// `authorization` is unrelated to this header and does not defer, which is + /// also what Python does: it sends the deployment's `x-api-key` alongside. + fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case("authorization") + && value + .strip_prefix("Bearer ") + .is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) + }) + } + + fn supported_params(&self) -> &'static [&'static str] { + SUPPORTED_PARAMS + } + + fn unsupported_reason( + &self, + messages: &[ChatMessage], + optional_params: &Map, + ) -> Option { + unsupported_param(SUPPORTED_PARAMS, &[], optional_params) + .or_else(|| messages.iter().find_map(unsupported_message)) + // Anthropic rejects a request whose first turn is not a user turn. + // Python only repairs that under `litellm.modify_params`, which the + // core cannot observe, so decline instead of guessing. + .or_else(|| { + (!build_conversation(messages).opens_on_user_turn()) + .then_some(Unsupported("conversation does not open on a user turn")) + }) + } + + fn transform_request( + &self, + model: &str, + messages: Vec, + optional_params: Map, + ) -> CoreResult { + Ok(ProviderChatRequestData { + body: anthropic_body(model, &build_conversation(&messages), optional_params), + }) + } + + fn transform_response( + &self, + _model: &str, + response: ProviderChatResponseData, + ) -> CoreResult { + let body = response.body.as_object().ok_or_else(|| { + CoreError::InvalidResponse("messages response is not an object".into()) + })?; + + let content = body + .get("content") + .and_then(Value::as_array) + .ok_or(CoreError::MissingField("content"))?; + // The route declines tool and thinking requests, so a non-text block + // means the response carries something this path never asked for. + // Decline rather than silently dropping it; the host falls back. + if content + .iter() + .any(|block| block.get("type").and_then(Value::as_str) != Some("text")) + { + return Err(CoreError::Unsupported("non-text response content block")); + } + let text: String = content + .iter() + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .collect(); + + let usage = body + .get("usage") + .and_then(Value::as_object) + .ok_or(CoreError::MissingField("usage"))?; + let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0); + + Ok(ChatCompletionsResponse { + created: unix_now(), + model: body + .get("model") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("model"))? + .to_string(), + choices: vec![ChatCompletionsChoice { + index: 0, + message: ChatCompletionsChoiceMessage { + role: "assistant".to_string(), + content: (!text.is_empty()).then_some(text), + }, + finish_reason: finish_reason_for( + body.get("stop_reason") + .and_then(Value::as_str) + .unwrap_or(""), + ) + .to_string(), + }], + usage: usage_from_parts( + field("input_tokens"), + field("output_tokens"), + field("cache_read_input_tokens"), + field("cache_creation_input_tokens"), + ), + }) + } +} + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; diff --git a/litellm-rust/crates/core/src/providers/anthropic/mod.rs b/litellm-rust/crates/core/src/providers/anthropic/mod.rs index ba63992f3cb..0bb20991ff7 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/mod.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/mod.rs @@ -1 +1,2 @@ +pub mod chat_completions; pub mod messages; diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index 86eb589e2c0..5e885734182 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -8,11 +8,8 @@ use crate::audio_transcription::types::{ }; use crate::error::{CoreError, CoreResult, json_type_name}; -use super::aws_base::AwsAuthConfig; -use super::constants::{ - AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE, - DEFAULT_BEDROCK_REGION, -}; +pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region}; +use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}; const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"]; @@ -21,64 +18,6 @@ pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig = pub struct BedrockAudioTranscriptionConfig; -pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { - let mut stripped = model; - for prefix in ["bedrock/converse/", "bedrock/", "converse/"] { - if let Some(value) = stripped.strip_prefix(prefix) { - stripped = value; - break; - } - } - let mut region = None; - if let Some((candidate, remainder)) = stripped.split_once('/') - && is_bedrock_region(candidate) - { - region = Some(candidate.to_string()); - stripped = remainder; - } - for prefix in ["nova-2/", "nova/"] { - if let Some(value) = stripped.strip_prefix(prefix) { - stripped = value; - break; - } - } - if region.is_none() { - region = stripped - .strip_prefix("arn:") - .and_then(|value| value.split(':').nth(3)) - .filter(|value| !value.is_empty()) - .map(str::to_string); - } - (stripped.to_string(), region) -} - -fn is_bedrock_region(value: &str) -> bool { - value.len() > 3 - && value.contains('-') - && value - .chars() - .all(|char| char.is_ascii_alphanumeric() || char == '-') -} - -pub fn resolve_bedrock_region( - model_region: Option<&str>, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> String { - if let Some(region) = optional_params - .get("aws_region_name") - .and_then(Value::as_str) - { - return region.to_string(); - } - if let Some(region) = model_region { - return region.to_string(); - } - env_lookup(AWS_REGION_NAME) - .or_else(|| env_lookup(AWS_REGION)) - .unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string()) -} - fn audio_fields(audio: Value) -> CoreResult<(String, String)> { let object = audio.as_object().ok_or_else(|| CoreError::InvalidType { expected: "object", @@ -203,32 +142,6 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { } } -pub fn aws_auth_config( - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> AwsAuthConfig { - let value = |key: &str| { - optional_params - .get(key) - .and_then(Value::as_str) - .map(str::to_string) - }; - let env = |key: &str| env_lookup(key); - AwsAuthConfig { - access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")), - secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")), - session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")), - region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)), - session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")), - profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")), - role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")), - web_identity_token: value("aws_web_identity_token") - .or_else(|| env("AWS_WEB_IDENTITY_TOKEN")), - sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")), - external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")), - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index dc036a3cf21..b11639aa09b 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -12,13 +12,15 @@ use aws_sigv4::http_request::{ }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; +use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; use super::constants::{ - AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN, - AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_STS_ENDPOINT, - AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, BEDROCK_SERVICE, - DEFAULT_SESSION_NAME_PREFIX, + AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME, + AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, + AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, + BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, + SIGV4_COMPUTED_HEADER_NAMES, }; const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); @@ -401,6 +403,33 @@ fn default_session_name() -> String { format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}") } +/// The subset of `headers` SigV4 should cover. +/// +/// Python signs only these and reattaches the rest afterwards, so a forwarded +/// client header cannot change the canonical request and invalidate the +/// signature. Signing everything instead makes the request 403 on a header the +/// caller supplied, on a deployment that works on the Python path. +pub fn aws_signature_headers(headers: &BTreeMap) -> BTreeMap { + headers + .iter() + .filter(|(name, _)| { + let name = name.to_ascii_lowercase(); + AWS_SIGNED_HEADER_NAMES.contains(&name.as_str()) + || name.starts_with("x-amz-") + || name.starts_with("x-amzn-") + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +/// Whether the signer produces `name` itself. +/// +/// Python's reattach loop skips these, so a caller-supplied copy never reaches +/// the wire next to the computed one. +pub fn is_sigv4_computed_header(name: &str) -> bool { + SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str()) +} + pub fn sign_bedrock_post( url: &str, body: &[u8], @@ -441,6 +470,121 @@ pub fn sign_bedrock_post( .collect()) } +/// Model-id and region parsing shared by every Bedrock route. +pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { + let mut stripped = model; + for prefix in ["bedrock/converse/", "bedrock/", "converse/"] { + if let Some(value) = stripped.strip_prefix(prefix) { + stripped = value; + break; + } + } + let mut region = None; + if let Some((candidate, remainder)) = stripped.split_once('/') + && is_bedrock_region(candidate) + { + region = Some(candidate.to_string()); + stripped = remainder; + } + for prefix in ["nova-2/", "nova/"] { + if let Some(value) = stripped.strip_prefix(prefix) { + stripped = value; + break; + } + } + if region.is_none() { + // Python splits the whole ARN and takes field 3, the region. Stripping + // `arn:` first shifts every field down one, so the region is field 2 + // here; field 3 is the account id. + region = stripped + .strip_prefix("arn:") + .and_then(|value| value.split(':').nth(2)) + .filter(|value| !value.is_empty()) + .map(str::to_string); + } + (stripped.to_string(), region) +} + +fn is_bedrock_region(value: &str) -> bool { + value.len() > 3 + && value.contains('-') + && value + .chars() + .all(|char| char.is_ascii_alphanumeric() || char == '-') +} + +pub fn resolve_bedrock_region( + model_region: Option<&str>, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + if let Some(region) = optional_params + .get("aws_region_name") + .and_then(Value::as_str) + { + return region.to_string(); + } + if let Some(region) = model_region { + return region.to_string(); + } + env_lookup(AWS_REGION_NAME) + .or_else(|| env_lookup(AWS_REGION)) + .unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string()) +} + +pub fn aws_auth_config( + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> AwsAuthConfig { + let value = |key: &str| { + optional_params + .get(key) + .and_then(Value::as_str) + .map(str::to_string) + }; + let env = |key: &str| env_lookup(key); + AwsAuthConfig { + access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")), + secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")), + session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")), + region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)), + session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")), + profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")), + role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")), + web_identity_token: value("aws_web_identity_token") + .or_else(|| env("AWS_WEB_IDENTITY_TOKEN")), + sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")), + external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")), + } +} + +/// Credentials a host resolved through its own chain and handed down verbatim. +/// +/// A host with its own resolution (LiteLLM's Python `BaseAWSLLM`, which reads +/// profiles, STS and boto sessions) passes the result here so the core signs +/// with exactly those. Without this the core would re-derive from ambient +/// state, where an unrelated `AWS_ROLE_NAME` or `AWS_PROFILE_NAME` in the +/// environment outranks explicit keys in [`classify_auth`] and the two sides +/// would sign as different principals. +pub fn host_supplied_credentials(optional_params: &Map) -> Option { + let value = |key: &str| { + optional_params + .get(key) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + }; + let access_key_id = value("aws_access_key_id")?; + let secret_access_key = value("aws_secret_access_key")?; + Some(Credentials::new( + access_key_id, + secret_access_key, + value("aws_session_token").map(str::to_string), + None, + "litellm-host-supplied", + )) +} + #[cfg(test)] mod tests { use super::*; @@ -458,6 +602,18 @@ mod tests { ) } + #[test] + fn reads_the_region_field_of_a_model_arn_not_the_account_id() { + // Python's `_get_aws_region_from_model_arn` splits the whole ARN and + // takes field 3. Stripping `arn:` first shifts every field down one, so + // the region is field 2 here. Taking field 3 after the strip returns + // the account id, which is not a region at all. + let (_, region) = bedrock_model_id_and_region( + "bedrock/arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2", + ); + assert_eq!(region.as_deref(), Some("us-west-2")); + } + #[test] fn classification_preserves_python_precedence() { let config = AwsAuthConfig { @@ -610,6 +766,52 @@ mod tests { )); } + #[test] + fn a_forwarded_client_header_is_not_folded_into_the_signature() { + // Python signs only the AWS header set, so a header a caller forwarded + // cannot change the canonical request. Signing it instead makes the + // request 403 the moment anything on the wire rewrites or drops it. + let (url, body, mut headers) = parity_inputs(); + headers.insert("x-request-id".to_string(), "abc-123".to_string()); + headers.insert("Accept-Encoding".to_string(), "gzip".to_string()); + headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string()); + let signable = aws_signature_headers(&headers); + + assert!(!signable.contains_key("x-request-id")); + assert!(!signable.contains_key("Accept-Encoding")); + // The AWS-prefixed one is genuinely part of the signature. + assert!(signable.contains_key("x-amzn-trace-id")); + assert!(signable.contains_key("Content-Type")); + + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + None, + None, + "test", + ); + let signed = sign_bedrock_post( + &url, + &body, + &signable, + "us-east-1", + &credentials, + SystemTime::UNIX_EPOCH, + ) + .expect("signs"); + let authorization = signed + .get("Authorization") + .expect("carries an authorization header"); + assert!( + !authorization.contains("x-request-id"), + "forwarded header reached SignedHeaders: {authorization}" + ); + assert!( + !authorization.contains("accept-encoding"), + "forwarded header reached SignedHeaders: {authorization}" + ); + } + #[test] fn signing_matches_botocore_golden_vector() { let (url, body, headers) = parity_inputs(); diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/mod.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs new file mode 100644 index 00000000000..4b75dcb8e9d --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs @@ -0,0 +1,580 @@ +use super::*; +use serde_json::json; + +fn messages(value: Value) -> Vec { + serde_json::from_value(value).expect("valid messages") +} + +fn params(value: Value) -> Map { + match value { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + } +} + +fn transform(msgs: Value, opts: Value) -> Value { + BEDROCK_CHAT_COMPLETIONS_CONFIG + .transform_request( + "anthropic.claude-sonnet-4-5-v1:0", + messages(msgs), + params(opts), + ) + .expect("request transforms") + .body +} + +fn transform_response(body: Value) -> CoreResult { + BEDROCK_CHAT_COMPLETIONS_CONFIG.transform_response( + "anthropic.claude-sonnet-4-5-v1:0", + ProviderChatResponseData { body }, + ) +} + +fn reason(msgs: Value, opts: Value) -> Option { + BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) +} + +#[test] +fn builds_the_converse_body_python_builds() { + let body = transform( + json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]), + json!({"maxTokens": 128, "temperature": 0.2}), + ); + assert_eq!( + body, + json!({ + "inferenceConfig": {"maxTokens": 128, "temperature": 0.2}, + "messages": [{"role": "user", "content": [{"text": "hi"}]}], + "system": [{"text": "be terse"}] + }) + ); +} + +#[test] +fn always_emits_inference_config_even_when_empty() { + let body = transform(json!([{"role": "user", "content": "hi"}]), json!({})); + assert_eq!(body["inferenceConfig"], json!({})); + assert!(body.get("system").is_none()); +} + +#[test] +fn places_only_inference_params_in_inference_config() { + let body = transform( + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 64, + "temperature": 0.1, + "topP": 0.9, + "stopSequences": ["STOP"] + }), + ); + assert_eq!( + body["inferenceConfig"], + json!({"maxTokens": 64, "temperature": 0.1, "topP": 0.9, "stopSequences": ["STOP"]}) + ); + assert!(body.get("additionalModelRequestFields").is_none()); +} + +#[test] +fn merges_consecutive_user_turns_into_one_message() { + let body = transform( + json!([ + {"role": "user", "content": "one"}, + {"role": "user", "content": [{"type": "text", "text": "two"}]}, + {"role": "assistant", "content": "ack"}, + {"role": "user", "content": "three"} + ]), + json!({}), + ); + assert_eq!( + body["messages"], + json!([ + {"role": "user", "content": [{"text": "one"}, {"text": "two"}]}, + {"role": "assistant", "content": [{"text": "ack"}]}, + {"role": "user", "content": [{"text": "three"}]} + ]) + ); +} + +#[test] +fn declines_streaming() { + assert_eq!( + reason( + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}) + ), + Some(Unsupported("streaming")) + ); +} + +#[test] +fn declines_top_k_because_python_routes_it_by_base_model() { + assert_eq!( + reason( + json!([{"role": "user", "content": "hi"}]), + json!({"topK": 40}) + ), + Some(Unsupported("unrecognized request parameter")) + ); +} + +#[test] +fn declines_tools_and_other_params_outside_the_allowlist() { + for param in [ + json!({"tools": []}), + json!({"tool_choice": {"auto": {}}}), + json!({"thinking": {"type": "enabled"}}), + json!({"requestMetadata": {"k": "v"}}), + json!({"outputConfig": {}}), + json!({"_parallel_tool_use_config": {}}), + ] { + assert_eq!( + reason(json!([{"role": "user", "content": "hi"}]), param.clone()), + Some(Unsupported("unrecognized request parameter")), + "expected {param} to decline" + ); + } +} + +#[test] +fn declines_blank_text_rather_than_substituting_the_anthropic_placeholder() { + for content in [ + json!(""), + json!(" "), + json!([{"type": "text", "text": " "}]), + ] { + assert_eq!( + reason( + json!([{"role": "user", "content": content}, {"role": "user", "content": "hi"}]), + json!({}) + ), + Some(Unsupported("blank message text")), + "expected blank content {content} to decline" + ); + } +} + +#[test] +fn declines_a_message_whose_content_list_is_empty() { + // The blank-text check scans parts, so an empty list clears it; Converse + // rejects an empty `content` array, which is a decline the core owes the + // host before the call rather than an error after it. + assert_eq!( + reason(json!([{"role": "user", "content": []}]), json!({})), + Some(Unsupported("message without content")) + ); + assert_eq!( + reason( + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]), + json!({}) + ), + None + ); +} + +#[test] +fn declines_a_conversation_that_opens_or_closes_on_an_assistant_turn() { + assert_eq!( + reason( + json!([ + {"role": "assistant", "content": "prefill"}, + {"role": "user", "content": "hi"} + ]), + json!({}) + ), + Some(Unsupported( + "conversation does not run user turn to user turn" + )) + ); + assert_eq!( + reason( + json!([ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "prefill"} + ]), + json!({}) + ), + Some(Unsupported( + "conversation does not run user turn to user turn" + )) + ); +} + +#[test] +fn accepts_a_user_to_user_text_conversation() { + assert_eq!( + reason( + json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "again"} + ]), + json!({"maxTokens": 16}) + ), + None + ); +} + +#[test] +fn builds_the_converse_url_from_the_region_in_the_model_id() { + let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG; + assert_eq!( + config + .complete_url(None, "us-east-1/anthropic.claude-v2", &Map::new(), &|_| { + None + }) + .expect("url builds"), + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" + ); +} + +#[test] +fn falls_back_to_the_region_env_then_the_default_region() { + let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG; + let with_env = |key: &str| (key == "AWS_REGION_NAME").then(|| "eu-west-1".to_string()); + assert_eq!( + config + .complete_url(None, "anthropic.claude-v2", &Map::new(), &with_env) + .expect("url builds"), + "https://bedrock-runtime.eu-west-1.amazonaws.com/model/anthropic.claude-v2/converse" + ); + assert_eq!( + config + .complete_url(None, "anthropic.claude-v2", &Map::new(), &|_| None) + .expect("url builds"), + "https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-v2/converse" + ); +} + +#[test] +fn prefers_an_explicit_runtime_endpoint_over_the_api_base() { + let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG; + let overrides = params(json!({"aws_bedrock_runtime_endpoint": "https://vpce.internal/"})); + assert_eq!( + config + .complete_url( + Some("https://ignored.example"), + "anthropic.claude-v2", + &overrides, + &|_| None + ) + .expect("url builds"), + "https://vpce.internal/model/anthropic.claude-v2/converse" + ); +} + +#[test] +fn signs_with_sigv4_in_the_resolved_region() { + let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG; + assert_eq!( + config + .auth( + None, + "eu-central-1/anthropic.claude-v2", + &Map::new(), + &|_| None + ) + .expect("auth resolves"), + ChatCompletionsAuth::AwsSigV4 { + region: "eu-central-1".to_string() + } + ); +} + +#[test] +fn a_bearer_token_outranks_sigv4_the_way_python_resolves_it() { + // Python's get_request_headers reads `api_key` as the Bedrock bearer token + // and only falls back to the env when the caller passed none, so each case + // pins one of its precedence rules. Signing as the host principal when a + // bearer identity is configured would cross an account and quota boundary. + let bedrock_env = + |key: &str| (key == "AWS_BEARER_TOKEN_BEDROCK").then(|| "from-env".to_string()); + let no_env = |_: &str| None; + let resolve = |api_key, env: &dyn Fn(&str) -> Option| { + BEDROCK_CHAT_COMPLETIONS_CONFIG + .auth( + api_key, + "eu-central-1/anthropic.claude-v2", + &Map::new(), + env, + ) + .expect("auth resolves") + }; + let bearer = |token: &str| ChatCompletionsAuth::Bearer { + token: token.to_string(), + }; + let sigv4 = ChatCompletionsAuth::AwsSigV4 { + region: "eu-central-1".to_string(), + }; + + // A caller-supplied key is the bearer token, and outranks the env. + assert_eq!( + resolve(Some("bedrock-api-key"), &bedrock_env), + bearer("bedrock-api-key") + ); + // No key, so the env supplies it. + assert_eq!(resolve(None, &bedrock_env), bearer("from-env")); + // An empty key is not a bearer token, and deliberately does NOT reach for + // the env, which is what Python's `is not None` check does. + assert_eq!(resolve(Some(""), &bedrock_env), sigv4); + // Whitespace is truthy in Python, so it stays a bearer token rather than + // silently becoming a host-credentialed SigV4 request. + assert_eq!(resolve(Some(" "), &no_env), bearer(" ")); + // Neither present, so SigV4 as before. + assert_eq!(resolve(None, &no_env), sigv4); +} + +#[test] +fn normalizes_a_converse_response_into_openai_shape() { + let response = transform_response(json!({ + "output": {"message": {"role": "assistant", "content": [ + {"text": "hello"}, {"text": " there"} + ]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15} + })) + .expect("response transforms"); + + assert_eq!(response.model, "anthropic.claude-sonnet-4-5-v1:0"); + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello there") + ); + assert_eq!(response.choices[0].finish_reason, "stop"); + assert_eq!(response.usage.prompt_tokens, 11); + assert_eq!(response.usage.completion_tokens, 4); + assert_eq!(response.usage.total_tokens, 15); +} + +#[test] +fn maps_converse_stop_reasons_python_maps() { + for (provider_reason, expected) in [ + ("end_turn", "stop"), + ("stop_sequence", "stop"), + ("max_tokens", "length"), + ("guardrail_intervened", "content_filter"), + // Converse emits this one, and Python's `_FINISH_REASON_MAP` carries + // it. Folding it into `stop` reports a filtered completion as a normal + // one to anything keying on the finish reason. + ("content_filtered", "content_filter"), + ("content_filter", "content_filter"), + ] { + let response = transform_response(json!({ + "output": {"message": {"content": [{"text": "x"}]}}, + "stopReason": provider_reason, + "usage": {"inputTokens": 1, "outputTokens": 1} + })) + .expect("response transforms"); + assert_eq!( + response.choices[0].finish_reason, expected, + "stopReason {provider_reason}" + ); + } +} + +#[test] +fn reports_an_empty_converse_answer_as_an_empty_string_not_null() { + // Converse assigns the joined text unconditionally + // (`chat_completion_message["content"] = content_str`), unlike Anthropic's + // `merged_text or None`, so an empty answer is `""` on both paths. A caller + // calling `.strip()` on it would break on the Rust path alone. Reachable + // through a filtered or guardrail-intervened response. + for content in [json!([]), json!([{"text": ""}])] { + let response = transform_response(json!({ + "output": {"message": {"content": content}}, + "stopReason": "content_filtered", + "usage": {"inputTokens": 1, "outputTokens": 0} + })) + .expect("response transforms"); + assert_eq!(response.choices[0].message.content, Some(String::new())); + } +} + +#[test] +fn reports_the_total_tokens_converse_sent_rather_than_recomputing_them() { + // Python reads `usage["totalTokens"]` straight through here, where Anthropic + // has no such field and adds the two counts instead. The two agree while the + // gate declines every cache_control request, so this is what keeps them + // agreeing if that ever widens. + let response = transform_response(json!({ + "output": {"message": {"content": [{"text": "x"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 4, "cacheReadInputTokens": 7, "totalTokens": 14} + })) + .expect("response transforms"); + assert_eq!( + response.usage.total_tokens, 14, + "provider total was recomputed" + ); + assert_eq!(response.usage.prompt_tokens, 17); + assert_eq!(response.usage.completion_tokens, 4); +} + +#[test] +fn falls_back_to_the_computed_total_when_converse_omits_it() { + // Python raises a KeyError on a body with no `totalTokens`. Reporting a zero + // instead would be a worse divergence than the one above, so the computed + // total stands in. + let response = transform_response(json!({ + "output": {"message": {"content": [{"text": "x"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 4} + })) + .expect("response transforms"); + assert_eq!(response.usage.total_tokens, 14); +} + +#[test] +fn declines_a_cache_control_message_so_widening_the_gate_is_a_red_test() { + // Converse only reports cache token counts when the request carries a + // cachePoint block, which is why the provider total and the computed one + // cannot disagree today. This is the tripwire: whoever widens the gate to + // admit prompt caching has to come back and re-check the usage mapping + // rather than discovering a silent number change in production. + assert_eq!( + reason( + json!([{"role": "user", "content": [ + {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}} + ]}]), + json!({}) + ), + Some(Unsupported("non-text message content")) + ); +} + +#[test] +fn folds_converse_cache_tokens_into_prompt_tokens() { + let response = transform_response(json!({ + "output": {"message": {"content": [{"text": "x"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 2, + "cacheReadInputTokens": 5, + "cacheWriteInputTokens": 3 + } + })) + .expect("response transforms"); + assert_eq!(response.usage.prompt_tokens, 18); + assert_eq!(response.usage.prompt_tokens_details.cached_tokens, 5); + assert_eq!( + response.usage.prompt_tokens_details.cache_creation_tokens, + 3 + ); + assert_eq!(response.usage.prompt_tokens_details.text_tokens, 10); +} + +#[test] +fn declines_a_response_carrying_a_tool_use_block() { + let err = transform_response(json!({ + "output": {"message": {"content": [ + {"toolUse": {"toolUseId": "t1", "name": "f", "input": {}}} + ]}}, + "stopReason": "tool_use", + "usage": {"inputTokens": 1, "outputTokens": 1} + })) + .expect_err("tool use block"); + assert_eq!( + err, + CoreError::Unsupported("non-text response content block") + ); +} + +#[test] +fn errors_on_a_response_missing_required_fields() { + assert_eq!( + transform_response(json!("nope")).expect_err("not an object"), + CoreError::InvalidResponse("converse response is not an object".to_string()) + ); + assert_eq!( + transform_response(json!({"usage": {}})).expect_err("no output"), + CoreError::MissingField("output.message.content") + ); + assert_eq!( + transform_response(json!({"output": {"message": {"content": []}}})).expect_err("no usage"), + CoreError::MissingField("usage") + ); +} + +#[test] +fn accepts_aws_call_configuration_without_serializing_it() { + let call_config = json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIA", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + "aws_region_name": "us-east-1", + "aws_profile_name": "litellm-stage", + "aws_role_name": "role", + "aws_session_name": "session", + "aws_web_identity_token": "wit", + "aws_sts_endpoint": "https://sts.example", + "aws_external_id": "ext", + "aws_bedrock_runtime_endpoint": "https://vpce.internal" + }); + assert_eq!( + reason( + json!([{"role": "user", "content": "hi"}]), + call_config.clone() + ), + None + ); + let body = transform(json!([{"role": "user", "content": "hi"}]), call_config); + assert_eq!( + body, + json!({ + "inferenceConfig": {"maxTokens": 16}, + "messages": [{"role": "user", "content": [{"text": "hi"}]}] + }), + "aws call configuration must not reach the Converse body" + ); +} + +#[test] +fn leaves_a_complete_converse_url_untouched() { + let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG; + let already_built = + "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-v2%3A0/converse"; + assert_eq!( + config + .complete_url( + Some(already_built), + "anthropic.claude-v2", + &Map::new(), + &|_| None + ) + .expect("url builds"), + already_built, + "a host that encoded the model id itself must not have it re-derived" + ); +} + +#[test] +fn host_supplied_credentials_outrank_ambient_profile_and_role_state() { + use crate::providers::bedrock::aws_base::host_supplied_credentials; + + let supplied = params(json!({ + "aws_access_key_id": "AKIAHOST", + "aws_secret_access_key": "hostsecret", + "aws_session_token": "hosttoken" + })); + let credentials = host_supplied_credentials(&supplied).expect("host credentials"); + assert_eq!(credentials.access_key_id(), "AKIAHOST"); + assert_eq!(credentials.secret_access_key(), "hostsecret"); + assert_eq!(credentials.session_token(), Some("hosttoken")); + + // Without a full static pair there is nothing to honor, so the core falls + // back to deriving credentials itself. + assert!(host_supplied_credentials(¶ms(json!({"aws_access_key_id": "AKIA"}))).is_none()); + assert!( + host_supplied_credentials(¶ms( + json!({"aws_access_key_id": " ", "aws_secret_access_key": "s"}) + )) + .is_none() + ); + assert!(host_supplied_credentials(&Map::new()).is_none()); +} diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs new file mode 100644 index 00000000000..b107950748e --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -0,0 +1,297 @@ +use serde_json::{Map, Value, json}; + +use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation}; +use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts}; +use crate::chat_completions::transformation::{ + ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message, + unsupported_param, +}; +use crate::chat_completions::types::{ + ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, + ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData, + ProviderChatResponseData, +}; +use crate::error::{CoreError, CoreResult}; + +use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region}; +use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}; + +/// Converse parameter names, post `map_openai_params`, that the Rust path can +/// place verbatim in `inferenceConfig`. +/// +/// `topK` is deliberately absent: Python routes it to +/// `additionalModelRequestFields` for Anthropic base models and to +/// `inferenceConfig` otherwise, and that branch reads the model catalog the +/// core cannot see. +const SUPPORTED_PARAMS: &[&str] = &["maxTokens", "temperature", "topP", "stopSequences"]; + +/// Params that belong in `inferenceConfig`, in the order Python's +/// `AmazonConverseConfig` declares them, so bodies compare cleanly. +const INFERENCE_CONFIG_PARAMS: &[&str] = SUPPORTED_PARAMS; + +const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "aws_bedrock_runtime_endpoint"; + +/// AWS call configuration a host passes down: consumed for signing and endpoint +/// resolution, never serialized into the Converse body. +const CONFIG_PARAMS: &[&str] = &[ + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_region_name", + "aws_session_name", + "aws_profile_name", + "aws_role_name", + "aws_web_identity_token", + "aws_sts_endpoint", + "aws_external_id", + AWS_BEDROCK_RUNTIME_ENDPOINT, +]; + +const CONVERSE_PATH_SUFFIX: &str = "/converse"; + +pub struct BedrockChatCompletionsConfig; + +pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: BedrockChatCompletionsConfig = + BedrockChatCompletionsConfig; + +fn converse_body(conversation: &Conversation, params: &Map) -> Value { + let messages: Vec = conversation + .turns + .iter() + .map(|turn| { + json!({ + "role": turn.role.as_str(), + "content": turn.texts.iter().map(|text| json!({"text": text})).collect::>(), + }) + }) + .collect(); + + let inference_config = Map::from_iter(INFERENCE_CONFIG_PARAMS.iter().filter_map(|name| { + params + .get(*name) + .map(|value| ((*name).to_string(), value.clone())) + })); + + let system: Vec = conversation + .system + .iter() + .map(|text| json!({"text": text})) + .collect(); + + Value::Object(Map::from_iter( + [ + ( + "inferenceConfig".to_string(), + Value::Object(inference_config), + ), + ("messages".to_string(), json!(messages)), + ] + .into_iter() + .chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))), + )) +} + +fn has_blank_text(message: &ChatMessage) -> bool { + match &message.content { + None => false, + Some(ChatMessageContent::Text(text)) => text.trim().is_empty(), + Some(ChatMessageContent::Parts(parts)) => parts.iter().any(|part| { + part.get("text") + .and_then(Value::as_str) + .is_none_or(|text| text.trim().is_empty()) + }), + } +} + +impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let (model_id, model_region) = bedrock_model_id_and_region(model); + let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup); + let endpoint = optional_params + .get(AWS_BEDROCK_RUNTIME_ENDPOINT) + .and_then(Value::as_str) + .or(api_base) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", ®ion)); + let endpoint = endpoint.trim_end_matches('/'); + // A host that already built the full Converse URL (LiteLLM's Python + // path encodes the model id itself) passes it through untouched, the + // way the Anthropic config leaves a complete `/v1/messages` URL alone. + if endpoint.ends_with(CONVERSE_PATH_SUFFIX) { + return Ok(endpoint.to_string()); + } + Ok(format!("{endpoint}/model/{model_id}{CONVERSE_PATH_SUFFIX}")) + } + + fn auth( + &self, + api_key: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + // Python reads `api_key` as the Bedrock bearer token and consults the + // env only when the caller passed none, so a caller-supplied empty key + // falls through to SigV4 without reaching for the environment. An + // all-whitespace token stays a bearer token here because Python sends + // it too: treating it as absent would sign as the host principal + // instead, which is the identity swap this branch exists to prevent. + let bearer = match api_key { + Some(key) => Some(key.to_string()), + None => env_lookup(AWS_BEARER_TOKEN_BEDROCK), + } + .filter(|token| !token.is_empty()); + if let Some(token) = bearer { + return Ok(ChatCompletionsAuth::Bearer { token }); + } + let (_, model_region) = bedrock_model_id_and_region(model); + Ok(ChatCompletionsAuth::AwsSigV4 { + region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), + }) + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + &[("Content-Type", "application/json")] + } + + fn supported_params(&self) -> &'static [&'static str] { + SUPPORTED_PARAMS + } + + fn config_params(&self) -> &'static [&'static str] { + CONFIG_PARAMS + } + + fn unsupported_reason( + &self, + messages: &[ChatMessage], + optional_params: &Map, + ) -> Option { + unsupported_param(SUPPORTED_PARAMS, CONFIG_PARAMS, optional_params) + .or_else(|| messages.iter().find_map(unsupported_message)) + // Python's Converse translation drops blank text blocks instead of + // substituting the placeholder the shared conversation builder + // applies, so decline blank text rather than diverge. + .or_else(|| { + messages + .iter() + .any(has_blank_text) + .then_some(Unsupported("blank message text")) + }) + // Converse has no assistant prefill: Python inserts a continue turn + // when a conversation opens or closes on an assistant message, and + // only under `litellm.modify_params`, which the core cannot see. + // Declining both ends also keeps the shared builder's final + // assistant right-strip (an Anthropic rule) unreachable here. + .or_else(|| { + let conversation = build_conversation(messages); + let ends_on_assistant = conversation + .turns + .last() + .is_some_and(|turn| turn.role == TurnRole::Assistant); + (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( + "conversation does not run user turn to user turn", + )) + }) + } + + fn transform_request( + &self, + _model: &str, + messages: Vec, + optional_params: Map, + ) -> CoreResult { + Ok(ProviderChatRequestData { + body: converse_body(&build_conversation(&messages), &optional_params), + }) + } + + fn transform_response( + &self, + model: &str, + response: ProviderChatResponseData, + ) -> CoreResult { + let body = response.body.as_object().ok_or_else(|| { + CoreError::InvalidResponse("converse response is not an object".into()) + })?; + + let content = body + .get("output") + .and_then(|output| output.get("message")) + .and_then(|message| message.get("content")) + .and_then(Value::as_array) + .ok_or(CoreError::MissingField("output.message.content"))?; + // The route declines tool requests, so anything other than a text block + // is something this path never asked for. Decline; the host falls back. + if content.iter().any(|block| { + block + .as_object() + .is_none_or(|block| block.len() != 1 || !block.contains_key("text")) + }) { + return Err(CoreError::Unsupported("non-text response content block")); + } + let text: String = content + .iter() + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .collect(); + + let usage = body + .get("usage") + .and_then(Value::as_object) + .ok_or(CoreError::MissingField("usage"))?; + let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0); + let computed = usage_from_parts( + field("inputTokens"), + field("outputTokens"), + field("cacheReadInputTokens"), + field("cacheWriteInputTokens"), + ); + // Converse reports `totalTokens` and Python passes it straight through, + // where Anthropic has no such field and Python adds the two counts + // instead, so only this provider overrides the computed total. Python + // does a bare `usage["totalTokens"]` lookup, so a body without the key + // raises there rather than reporting a zero; fall back to the computed + // total, which is the closest thing to that without failing the call. + let usage = ChatCompletionsUsage { + total_tokens: usage + .get("totalTokens") + .and_then(Value::as_u64) + .unwrap_or(computed.total_tokens), + ..computed + }; + + Ok(ChatCompletionsResponse { + created: unix_now(), + // Converse echoes no model id, so Python reports the requested one. + model: model.to_string(), + choices: vec![ChatCompletionsChoice { + index: 0, + message: ChatCompletionsChoiceMessage { + role: "assistant".to_string(), + // Converse assigns the joined string unconditionally, so an + // empty response is `""` here and not `None` as it is on + // Anthropic. A caller calling `.strip()` on it would break + // on this path alone. + content: Some(text), + }, + finish_reason: finish_reason_for( + body.get("stopReason").and_then(Value::as_str).unwrap_or(""), + ) + .to_string(), + }], + usage, + }) + } +} + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; diff --git a/litellm-rust/crates/core/src/providers/bedrock/constants.rs b/litellm-rust/crates/core/src/providers/bedrock/constants.rs index 785295207e7..be215cc9016 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/constants.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/constants.rs @@ -11,6 +11,31 @@ pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN"; pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; +pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK"; + +/// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors +/// Python's `_filter_headers_for_aws_signature` allowlist. +pub const AWS_SIGNED_HEADER_NAMES: &[&str] = &[ + "host", + "content-type", + "date", + "x-amz-date", + "x-amz-security-token", + "x-amz-content-sha256", + "x-amz-algorithm", + "x-amz-credential", + "x-amz-signedheaders", + "x-amz-signature", +]; +/// Headers the signer emits itself. Mirrors Python's `SIGV4_COMPUTED_HEADERS`, +/// which the reattach loop skips so a caller's copy cannot ride alongside the +/// computed one. +pub const SIGV4_COMPUTED_HEADER_NAMES: &[&str] = &[ + "authorization", + "x-amz-date", + "x-amz-security-token", + "date", +]; pub const BEDROCK_SERVICE: &str = "bedrock"; pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session"; pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2"; diff --git a/litellm-rust/crates/core/src/providers/bedrock/mod.rs b/litellm-rust/crates/core/src/providers/bedrock/mod.rs index b09675ad7dd..d9cd3efcb74 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/mod.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/mod.rs @@ -5,4 +5,5 @@ #[cfg(feature = "bedrock-auth")] pub mod audio_transcription; pub mod aws_base; +pub mod chat_completions; mod constants; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f0cc26a0cca..c6f81cf6916 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -6,6 +6,10 @@ use litellm_ai_gateway::io::audio_transcription::{ }; use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; +use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; +use litellm_core::chat_completions::{ + chat_completions as run_chat_completions, chat_completions_decline_reason, +}; use litellm_core::error::CoreError; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; @@ -16,6 +20,20 @@ use serde_json::{Map, Value}; mod gil; +pyo3::create_exception!( + _native, + RustBridgeDeclined, + pyo3::exceptions::PyException, + "The route declined before calling the provider, so the host may retry on its own path." +); + +pyo3::create_exception!( + _native, + RustUpstreamError, + pyo3::exceptions::PyException, + "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." +); + type MarshaledOcrInputs = ( Value, Option>, @@ -45,6 +63,15 @@ fn messages_response_to_py( json_to_py(py, value) } +fn chat_completions_response_to_py( + py: Python<'_>, + response: ChatCompletionsResponse, +) -> PyResult> { + let value = + serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?; + json_to_py(py, value) +} + fn core_error_to_pyerr(err: CoreError) -> PyErr { match err { CoreError::Auth(message) => PyValueError::new_err(message), @@ -56,6 +83,33 @@ fn core_error_to_pyerr(err: CoreError) -> PyErr { } } +/// Map a core error for a route whose host keeps a Python implementation. +/// +/// The distinction the host needs is whether the provider was already called. +/// Everything raised before the request goes out is safe for the host to retry +/// on its own path; anything after it is not, because the provider has already +/// done the work and billed for it. +fn chat_completions_error_to_pyerr(err: CoreError) -> PyErr { + match err { + CoreError::Unsupported(_) + | CoreError::Auth(_) + | CoreError::InvalidProvider(_) + | CoreError::InvalidRequest(_) + | CoreError::InvalidType { .. } + | CoreError::MissingField(_) + | CoreError::Routing(_) + // Nothing reached the provider, so serving it on Python cannot double + // bill and is the only way the caller gets an answer at all. + | CoreError::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), + CoreError::Http { status, body } => { + RustUpstreamError::new_err((status, format!("{status}: {body}"))) + } + CoreError::Network(message) | CoreError::InvalidResponse(message) => { + RustUpstreamError::new_err((0u16, message)) + } + } +} + fn optional_object_to_map( py: Python<'_>, name: &'static str, @@ -430,6 +484,143 @@ fn amessages( }) } +type MarshaledChatCompletionsInputs = ( + Value, + Map, + Option>, + Option, +); + +fn marshal_chat_completions_inputs( + py: Python<'_>, + messages: Py, + optional_params: Option>, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult { + let messages = py_to_json(py, messages.bind(py))?; + if !messages.is_array() { + return Err(PyValueError::new_err("messages must be a list")); + } + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + Ok(( + messages, + optional_params, + extra_headers, + optional_timeout(timeout_seconds), + )) +} + +/// The decline reason for this request, or `None` when the Rust path accepts +/// it. Resolves no credentials and performs no I/O, so a host can ask before +/// committing to either path. +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] +fn chat_completions_decline( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + custom_llm_provider: Option, +) -> PyResult> { + let messages = py_to_json(py, messages.bind(py))?; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + Ok(chat_completions_decline_reason( + &model, + custom_llm_provider.as_deref(), + messages, + &optional_params, + ) + .map(str::to_string)) +} + +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn chat_completions( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( + py, + messages, + optional_params, + extra_headers, + timeout_seconds, + )?; + + let result = gil::release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions( + ChatCompletionsRequest { + model: &model, + messages, + optional_params, + 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(response) => chat_completions_response_to_py(py, response), + Err(err) => Err(chat_completions_error_to_pyerr(err)), + } +} + +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn achat_completions( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( + py, + messages, + optional_params, + extra_headers, + timeout_seconds, + )?; + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let response = run_chat_completions(ChatCompletionsRequest { + model: &model, + messages, + optional_params, + 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(chat_completions_error_to_pyerr)?; + + Python::attach(|py| chat_completions_response_to_py(py, response)) + }) +} + #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); @@ -439,12 +630,18 @@ fn gil_stats(py: Python<'_>) -> PyResult> { #[pymodule] fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { + let py = module.py(); module.add_function(wrap_pyfunction!(ocr, module)?)?; module.add_function(wrap_pyfunction!(aocr, module)?)?; module.add_function(wrap_pyfunction!(transcription, module)?)?; module.add_function(wrap_pyfunction!(atranscription, module)?)?; module.add_function(wrap_pyfunction!(messages, module)?)?; module.add_function(wrap_pyfunction!(amessages, module)?)?; + module.add("RustBridgeDeclined", py.get_type::())?; + module.add("RustUpstreamError", py.get_type::())?; + module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; + module.add_function(wrap_pyfunction!(chat_completions, module)?)?; + module.add_function(wrap_pyfunction!(achat_completions, module)?)?; module.add_class::()?; module.add_function(wrap_pyfunction!(gil_stats, module)?)?; Ok(()) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 3eb8c163d5c..b12c715c9f5 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -21,6 +21,14 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( } ) +# The per-deployment Rust opt-in. +RUST_KWARG_KEY: Final = "rust" + +# Keys `completion()` forwards from its own kwargs into `get_litellm_params`, +# which are otherwise invisible to it because that call site passes explicit +# named arguments rather than `**kwargs`. +FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY}) + # Pre-define optional kwargs keys as frozenset for O(1) lookups # These are extracted from kwargs only if present, avoiding unnecessary .get() calls OPTIONAL_KWARGS_KEYS: Final = ( @@ -47,6 +55,10 @@ OPTIONAL_KWARGS_KEYS: Final = ( "itpm", "otpm", "use_xai_oauth", + # The per-deployment Rust opt-in. `all_litellm_params` keeps it out + # of the provider body; this keeps it *in* litellm_params, which is + # where the chat completions handlers read it from. + RUST_KWARG_KEY, } ) | AWS_CREDENTIAL_KWARGS_KEYS diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 39d3947c07c..d9bb0d7abff 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -24,6 +24,8 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) +from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge +from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts from litellm.types.llms.anthropic import ( ContentBlockDelta, ContentBlockStart, @@ -361,30 +363,135 @@ class AnthropicChatCompletion(BaseLLM): if config is None: raise ValueError(f"Provider config not found for model: {model} and provider: {custom_llm_provider}") - data = config.transform_request( + def build_request() -> tuple[dict, dict]: # mutable-ok: rewritten in place downstream + """Translate the request the Python way, returning `(headers, data)`. + + The pair stays mutable because the streaming path rewrites it in + place (`data["stream"] = True`) before sending. + + Shared by the normal path and by the Rust path's fallback, which + builds it only when the Rust call did not serve the request. + """ + request_data: Final = config.transform_request( + model=model, + messages=messages, + optional_params={**optional_params, "is_vertex_request": is_vertex_request}, + litellm_params=litellm_params, + headers=headers, + ) + return update_request_with_filtered_beta( + headers=headers, + request_data=request_data, + provider=custom_llm_provider, + ) + + # The Rust core owns the whole call for the subset it accepts, so ask + # before transforming: whichever path runs emits pre_call exactly once. + # `get_config` merges the class-level defaults (Anthropic's required + # `max_tokens` among them) that `transform_request` would have applied. + rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy + **AnthropicConfig.get_config(model=model), + **optional_params, + } + serves_via_rust: Final = rust_chat_completions_accepts( model=model, messages=messages, - optional_params={**optional_params, "is_vertex_request": is_vertex_request}, + optional_params=rust_optional_params, + custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, - headers=headers, + stream=stream, ) - - headers, data = update_request_with_filtered_beta( - headers=headers, - request_data=data, - provider=custom_llm_provider, - ) - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key=api_key, - additional_args={ - "complete_input_dict": data, + if serves_via_rust: + rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict + "complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent + "model": model, + "messages": messages, + **rust_optional_params, + }, "api_base": api_base, "headers": headers, - }, - ) + } + logging_obj.pre_call(input=messages, api_key=api_key, additional_args=rust_logging_args) + log_rust_post_call: Final = rust_chat_completions_bridge.response_logger( + logging_obj=logging_obj, + messages=messages, + api_key=api_key, + additional_args=rust_logging_args, + ) + if acompletion is True: + + async def python_fallback() -> "ModelResponse | CustomStreamWrapper": + # pre_call already fired for this request above. The Rust + # path only declines before the provider is called, so this + # is the same attempt continuing, not a second one. + fallback_headers, fallback_data = build_request() + return await self.acompletion_function( + model=model, + messages=messages, + data=fallback_data, + api_base=api_base, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + encoding=encoding, + api_key=api_key, + provider_config=config, + logging_obj=logging_obj, + optional_params=optional_params, + stream=stream, + _is_function_call=_is_function_call, + litellm_params=litellm_params, + logger_fn=logger_fn, + headers=fallback_headers, + client=client, + json_mode=json_mode, + timeout=timeout, + ) + + return rust_chat_completions_bridge.achat_completions_or_fallback( + model=model, + messages=messages, + optional_params=rust_optional_params, + model_response=model_response, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + on_response=log_rust_post_call, + python_fallback=python_fallback, + ) + rust_response: Final = rust_chat_completions_bridge.chat_completions( + model=model, + messages=messages, + optional_params=rust_optional_params, + model_response=model_response, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + on_response=log_rust_post_call, + ) + if rust_response is not None: + return rust_response + + headers, data = build_request() + + ## LOGGING + # Reaching here with `serves_via_rust` set means the Rust attempt + # declined at call time, before the provider was called, and already + # logged this request. That is the same attempt continuing. + if not serves_via_rust: + logging_obj.pre_call( + input=messages, + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": headers, + }, + ) print_verbose(f"_is_function_call: {_is_function_call}") if acompletion is True: if ( diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 25e544f4521..3a093e1f939 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -14,6 +14,8 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) +from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge +from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper @@ -169,6 +171,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers: dict = {}, client: AsyncHTTPHandler | None = None, api_key: str | None = None, + skip_pre_call_logging: bool = False, ) -> ModelResponse | CustomStreamWrapper: request_data: Final = await litellm.AmazonConverseConfig()._async_transform_request( model=model, @@ -190,15 +193,19 @@ class BedrockConverseLLM(BaseAWSLLM): ) ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": prepped.headers, - }, - ) + # The Rust path already logged this request's pre_call before handing + # it here, and it only declines before the provider is called, so this + # is the same attempt continuing rather than a second one. + if not skip_pre_call_logging: + logging_obj.pre_call( + input=messages, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": prepped.headers, + }, + ) headers = dict(prepped.headers) if client is None or not isinstance(client, AsyncHTTPHandler): @@ -354,6 +361,94 @@ class BedrockConverseLLM(BaseAWSLLM): # Filter beta headers in HTTP headers before making the request headers = update_headers_with_filtered_beta(headers=headers, provider="bedrock_converse") + + # The Rust core owns the whole call for the subset it accepts. Ask + # before transforming so whichever path runs emits pre_call once, and + # hand down the credentials, region and endpoint this handler already + # resolved so both paths sign as the same principal. + rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy + **optional_params, + **{ # mutable-ok: merged into its mutable parent above + key: value + for key, value in ( + ("aws_access_key_id", credentials.access_key), + ("aws_secret_access_key", credentials.secret_key), + ("aws_session_token", credentials.token), + ("aws_region_name", aws_region_name), + ) + if value is not None + }, + } + serves_via_rust: Final = rust_chat_completions_accepts( + model=model, + messages=messages, + optional_params=rust_optional_params, + custom_llm_provider="bedrock", + litellm_params=litellm_params, + stream=stream, + ) + if serves_via_rust: + rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict + "complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent + "messages": messages, + **optional_params, + }, + "api_base": proxy_endpoint_url, + "headers": headers, + } + logging_obj.pre_call(input=messages, api_key="", additional_args=rust_logging_args) + log_rust_post_call: Final = rust_chat_completions_bridge.response_logger( + logging_obj=logging_obj, + messages=messages, + api_key="", + additional_args=rust_logging_args, + ) + if acompletion: + return rust_chat_completions_bridge.achat_completions_or_fallback( + model=model, + messages=messages, + optional_params=rust_optional_params, + model_response=model_response, + api_key=api_key, + api_base=proxy_endpoint_url, + custom_llm_provider="bedrock", + extra_headers=headers, + timeout=timeout, + on_response=log_rust_post_call, + python_fallback=lambda: self.async_completion( + model=model, + messages=messages, + api_base=proxy_endpoint_url, + model_response=model_response, + encoding=encoding, + logging_obj=logging_obj, + optional_params=optional_params, + stream=stream, + litellm_params=litellm_params, + logger_fn=logger_fn, + headers=headers, + timeout=timeout, + client=client, + credentials=credentials, + api_key=api_key, + skip_pre_call_logging=True, + ), + ) + rust_response: Final = rust_chat_completions_bridge.chat_completions( + model=model, + messages=messages, + optional_params=rust_optional_params, + model_response=model_response, + api_key=api_key, + api_base=proxy_endpoint_url, + custom_llm_provider="bedrock", + extra_headers=headers, + timeout=timeout, + on_response=log_rust_post_call, + ) + if rust_response is not None: + return rust_response + ### ROUTING (ASYNC, STREAMING, SYNC) if acompletion: if isinstance(client, HTTPHandler): @@ -420,15 +515,21 @@ class BedrockConverseLLM(BaseAWSLLM): ) ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": proxy_endpoint_url, - "headers": prepped.headers, - }, - ) + # Reaching here with `serves_via_rust` set means the synchronous Rust + # attempt declined at call time, before the provider was called, and + # already logged this request. That is the same attempt continuing. + # The asynchronous branch above returns before this point, and hands + # its own fallback `skip_pre_call_logging=True` for the same reason. + if not serves_via_rust: + logging_obj.pre_call( + input=messages, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": proxy_endpoint_url, + "headers": prepped.headers, + }, + ) if client is None or isinstance(client, AsyncHTTPHandler): _params: Final = {} if timeout is not None: diff --git a/litellm/main.py b/litellm/main.py index 98c220f94e0..f66767a8d42 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -75,7 +75,7 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import ( from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_litellm_params import ( - AWS_CREDENTIAL_KWARGS_KEYS, + FORWARDED_KWARGS_KEYS, OPTIONAL_KWARGS_KEYS, ) from litellm.litellm_core_utils.get_provider_specific_headers import ( @@ -5451,7 +5451,7 @@ def completion( tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), - **{key: kwargs[key] for key in AWS_CREDENTIAL_KWARGS_KEYS if key in kwargs}, + **{key: kwargs[key] for key in FORWARDED_KWARGS_KEYS if key in kwargs}, ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 912a0b0ebd0..ce662ee0374 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -311,6 +311,12 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # the request away from the admin's pinned configuration. "nvcf_function_id", "use_ssl", + # Per-deployment opt-in that hands the whole call to the Rust core. It is a + # deployment decision, not a request one: the Rust path uses its own client + # rather than the one the deployment configured, and reports no post_call, + # so a caller-supplied value picks a transport and a callback surface the + # admin did not choose. + "rust", # SDK-only field; also rejected outright in is_request_body_safe. "model_list", "vertex_ai_credentials", diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py new file mode 100644 index 00000000000..acda3086051 --- /dev/null +++ b/litellm/rust_bridge/chat_completions.py @@ -0,0 +1,453 @@ +"""Thin Python wrapper for the native Rust chat completions bridge. + +The Rust core owns the conversation translation, the provider call, and the +response normalization for the subset of `/chat/completions` requests it +accepts. This module only marshals inputs and hands the normalized result to +LiteLLM's existing `ModelResponse` builder. + +``None`` means the provider was never called, so the caller is free to serve the +request on the Python path. A failure after the call was issued raises instead: +retrying it there would bill the customer for the same work twice. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Protocol + +import httpx +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.exceptions import APIError +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_model_response_object, +) +from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned +from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge.timeouts import timeout_to_seconds +from litellm.types.utils import ModelResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +# Providers whose `/chat/completions` deployments the Rust core can serve. A +# provider outside this set never reaches the bridge. +RUST_CHAT_COMPLETIONS_PROVIDERS: Final = frozenset({"anthropic", "bedrock"}) + +# `litellm_params` values are `object`, so validate the one this module reads +# rather than narrowing an unparameterized `Mapping` and typing the result Any. +_LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + +RUST_RESPONSE_HEADER: Final = "x-litellm-rust" + +_TRUTHY_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) + + +class RustChatCompletions(Protocol): + def __call__( + self, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object] | None, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: Mapping[str, object] | None, + timeout_seconds: float | None, + ) -> Mapping[str, object]: + raise NotImplementedError + + +class RustAchatCompletions(Protocol): + def __call__( + self, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object] | None, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: Mapping[str, object] | None, + timeout_seconds: float | None, + ) -> Awaitable[Mapping[str, object]]: + raise NotImplementedError + + +class RustChatCompletionsDecline(Protocol): + def __call__( + self, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object] | None, + custom_llm_provider: str | None, + ) -> str | None: + raise NotImplementedError + + +class ResponseObserver(Protocol): + """Invoked with the payload the core returned, on success only. + + Lets the caller emit its own `post_call` on whichever path served the + request. Both entry points call it, so the synchronous and asynchronous + paths cannot drift apart the way the pre_call suppression once did. + """ + + def __call__(self, rust_response: Mapping[str, object], /) -> None: + raise NotImplementedError + + +def response_logger( + *, + logging_obj: LiteLLMLoggingObj, + messages: Sequence[object], + api_key: str, + additional_args: Mapping[str, object], +) -> ResponseObserver: + """A `ResponseObserver` that emits the caller's `post_call` for a Rust-served + request. + + The core owns the provider call, so the Python transform that normally + raises this event never runs; without it every `post_call` callback goes + silent on a Rust-served request and `original_response` stays unset. The + payload is the core's normalized response rather than the provider's wire + body, which is the closest thing that crosses the bridge. + """ + + def log(rust_response: Mapping[str, object], /) -> None: + logging_obj.post_call( + input=messages, + api_key=api_key, + original_response=json.dumps(rust_response), + additional_args=additional_args, + ) + + return log + + +class _Unset: + pass + + +_UNSET: Final[_Unset] = _Unset() + + +@dataclass(slots=True) +class _RustChatCompletionsState: + chat_completions: RustChatCompletions | None = None + achat_completions: RustAchatCompletions | None = None + decline: RustChatCompletionsDecline | None = None + + +_STATE: Final[_RustChatCompletionsState] = _RustChatCompletionsState() + + +def set_rust_chat_completions( + *, + chat_completions: RustChatCompletions | None | _Unset = _UNSET, + achat_completions: RustAchatCompletions | None | _Unset = _UNSET, + decline: RustChatCompletionsDecline | None | _Unset = _UNSET, +) -> None: + """Inject the native callables, so tests can supply a double instead of + patching module attributes.""" + if not isinstance(chat_completions, _Unset): + _STATE.chat_completions = chat_completions + if not isinstance(achat_completions, _Unset): + _STATE.achat_completions = achat_completions + if not isinstance(decline, _Unset): + _STATE.decline = decline + + +def load_rust_chat_completions() -> RustChatCompletions | None: + if _STATE.chat_completions is not None: + return _STATE.chat_completions + native_bridge: Final = get_native_bridge() + if native_bridge is None: + return None + loaded: RustChatCompletions | None = getattr(native_bridge, "chat_completions", None) + return loaded + + +def load_rust_achat_completions() -> RustAchatCompletions | None: + if _STATE.achat_completions is not None: + return _STATE.achat_completions + native_bridge: Final = get_native_bridge() + if native_bridge is None: + return None + loaded: RustAchatCompletions | None = getattr(native_bridge, "achat_completions", None) + return loaded + + +def _env_enables_rust() -> bool: + return os.getenv("LITELLM_RUST", "").strip().lower() in _TRUTHY_ENV_VALUES + + +def _load_rust_decline() -> RustChatCompletionsDecline | None: + if _STATE.decline is not None: + return _STATE.decline + native_bridge: Final = get_native_bridge() + if native_bridge is None: + return None + loaded: RustChatCompletionsDecline | None = getattr(native_bridge, "chat_completions_decline", None) + return loaded + + +def _anthropic_user_id_reaches_the_body(litellm_params: Mapping[str, object] | None) -> bool: + metadata: Final = litellm_params.get("metadata") if litellm_params is not None else None + try: + entries: Final = _LITELLM_METADATA_ADAPTER.validate_python(metadata) + except ValidationError: + return False + return entries.get("user_id") is not None + + +def _litellm_metadata_reaches_the_provider( + custom_llm_provider: str | None, litellm_params: Mapping[str, object] | None +) -> bool: + """Whether the Python transform would promote proxy-owned attribution into the + provider request, below this gate and inside the function the Rust route replaces. + + `AnthropicConfig.transform_request` promotes a valid `metadata["user_id"]` + into the Messages body, so the core never sees the key and would send the + request to Anthropic with the abuse-detection attribution missing. + + `AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the + Converse body whenever the operator armed `bedrock_request_metadata_fields`. + Owning that field also means evicting a caller-supplied one, which the core + cannot do either, so ownership alone is the condition rather than whether + anything resolved. + + Deliberately a superset of Python's condition in both cases: declining a + request Python would not have attributed anyway costs only the Rust path, + while missing one loses the attribution silently. + """ + match custom_llm_provider: + case "anthropic": + return _anthropic_user_id_reaches_the_body(litellm_params) + case "bedrock": + return bedrock_request_metadata_is_owned() + case _: + return False + + +def rust_chat_completions_accepts( + *, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object], + custom_llm_provider: str | None, + litellm_params: Mapping[str, object] | None, + stream: object, +) -> bool: + """Whether the Rust path will serve this request. + + Asked before the caller commits to either path, so pre-call logging is + emitted exactly once, on whichever path actually runs. The core's own + capability gate answers the second half; it resolves no credentials and + performs no I/O. + """ + if custom_llm_provider not in RUST_CHAT_COMPLETIONS_PROVIDERS: + return False + if stream: + return False + opted_in: Final = litellm_params is not None and litellm_params.get("rust") is True + if not opted_in and not _env_enables_rust(): + return False + if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): + verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") + return False + decline: Final = _load_rust_decline() + if decline is None: + return False + try: + reason: Final = decline( + model=model, + messages=messages, + optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + ) + 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 chat completions gate raised %s; staying on the Python path", + type(rust_error).__name__, + ) + return False + if reason is not None: + verbose_logger.debug("Rust chat completions declined (%s); using the Python path", reason) + return False + return True + + +def _rust_bridge_exceptions() -> tuple[type[BaseException], type[BaseException]] | None: + """`(declined, upstream_failed)` from the native module, or None when absent.""" + native_bridge: Final = get_native_bridge() + if native_bridge is None: + return None + declined: Final = getattr(native_bridge, "RustBridgeDeclined", None) + upstream: Final = getattr(native_bridge, "RustUpstreamError", None) + if declined is None or upstream is None: + return None + return declined, upstream + + +def _reraise_or_decline( + rust_error: BaseException, + *, + model: str, + custom_llm_provider: str | None, +) -> None: + """Re-raise a failure the provider already saw, or return so the caller declines. + + A request that never reached the provider is safe to serve on the Python + path. One that did is not: the provider has already done the work, so a + second attempt bills for it twice. Those surface as an `APIError` carrying + the upstream status, which LiteLLM's exception mapping already understands. + """ + exceptions: Final = _rust_bridge_exceptions() + if exceptions is None: + verbose_logger.debug( + "Rust chat completions bridge raised %s; falling back to Python path", + type(rust_error).__name__, + ) + return + declined, upstream_failed = exceptions + if isinstance(rust_error, upstream_failed): + args: Final = rust_error.args + status: Final = args[0] if args else 0 + message: Final = args[1] if len(args) > 1 else "" + raise APIError( + status_code=int(status) or 500, + message=f"litellm rust chat completions: {message}", + llm_provider=custom_llm_provider or "", + model=model, + ) + if not isinstance(rust_error, declined): + raise rust_error + verbose_logger.debug( + "Rust chat completions declined before calling the provider (%s); using the Python path", + rust_error, + ) + + +def _build_model_response( + rust_response: Mapping[str, object], + model_response: ModelResponse, +) -> ModelResponse: + built: Final = convert_to_model_response_object( + response_object=dict(rust_response), # mutable-ok: the converter takes a real dict and rewrites it + model_response_object=model_response, + hidden_params={"additional_headers": {RUST_RESPONSE_HEADER: "true"}}, # mutable-ok: rewritten by the converter + ) + if not isinstance(built, ModelResponse): + raise TypeError(f"expected a ModelResponse from the rust path, got {type(built).__name__}") + return built + + +def chat_completions( + *, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object], + model_response: ModelResponse, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: Mapping[str, object] | None, + timeout: float | httpx.Timeout | None, + on_response: ResponseObserver, +) -> ModelResponse | None: + rust_chat_completions: Final = load_rust_chat_completions() + if rust_chat_completions is None: + return None + try: + rust_response: Final = rust_chat_completions( + model=model, + messages=messages, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ) + except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw + _reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider) + return None + on_response(rust_response) + return _build_model_response(rust_response, model_response) + + +async def achat_completions( + *, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object], + model_response: ModelResponse, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: Mapping[str, object] | None, + timeout: float | httpx.Timeout | None, + on_response: ResponseObserver, +) -> ModelResponse | None: + rust_achat_completions: Final = load_rust_achat_completions() + if rust_achat_completions is None: + return None + try: + rust_response: Final = await rust_achat_completions( + model=model, + messages=messages, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ) + except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw + _reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider) + return None + on_response(rust_response) + return _build_model_response(rust_response, model_response) + + +async def achat_completions_or_fallback( + *, + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object], + model_response: ModelResponse, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: Mapping[str, object] | None, + timeout: float | httpx.Timeout | None, + on_response: ResponseObserver, + python_fallback: Callable[[], Awaitable[object]], +) -> object: + """Await the Rust path, falling back to the caller's own Python path when + the bridge is unavailable or the call fails. + + The caller supplies the fallback, so the bridge stays free of provider + dispatch. This exists because a caller that dispatches asynchronously has + already returned a coroutine by the time a Rust failure surfaces, and so + cannot fall back on its own. + """ + response: Final = await achat_completions( + model=model, + messages=messages, + optional_params=optional_params, + model_response=model_response, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout=timeout, + on_response=on_response, + ) + if response is not None: + return response + return await python_fallback() diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index fb4cb494bee..956da571d43 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -215,3 +215,32 @@ class TestMetadataFallsBackToLitellmMetadata: assert result["metadata"] is not litellm_metadata result["metadata"].pop("trace_id") assert litellm_metadata == {"trace_id": "trace-1"} + + +class TestRustOptIn: + """`rust: true` is a litellm param, so it has to reach `litellm_params`. + + `all_litellm_params` keeps it out of the provider body; without it also + being carried into `litellm_params` the chat completions handlers cannot + see the opt-in and the Rust path is silently never taken. + """ + + def test_rust_is_an_optional_kwargs_key(self): + assert "rust" in _OPTIONAL_KWARGS_KEYS + + def test_rust_is_forwarded_from_completion_kwargs(self): + from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS + + assert "rust" in FORWARDED_KWARGS_KEYS + + def test_rust_survives_into_litellm_params(self): + params = get_litellm_params(rust=True) + assert params["rust"] is True + + def test_rust_is_absent_when_the_deployment_did_not_set_it(self): + assert "rust" not in get_litellm_params() + + def test_rust_stays_out_of_the_provider_body(self): + from litellm.types.utils import all_litellm_params + + assert "rust" in all_litellm_params diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index f934c7184f8..f6cd6ac6734 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -1,7 +1,7 @@ import json import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -2045,3 +2045,389 @@ def test_non_bash_tool_result_skipped(): assert ( len(code_results) == 0 ), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}" + + +class TestRustChatCompletionsHook: + """The `rust: true` opt-in on `/chat/completions` for the Anthropic provider. + + The native callables are dependency-injected, so these run without the + compiled extension. + """ + + RUST_RESPONSE = { + "created": 1_700_000_000, + "model": "claude-sonnet-4-5-20260101", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hello from rust"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 11, + "completion_tokens": 4, + "total_tokens": 15, + "prompt_tokens_details": { + "cached_tokens": 0, + "cache_creation_tokens": 0, + "text_tokens": 11, + }, + }, + } + + @pytest.fixture(autouse=True) + def _reset_bridge(self, monkeypatch): + from litellm.rust_bridge import chat_completions as bridge + + monkeypatch.delenv("LITELLM_RUST", raising=False) + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) + yield + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) + + @staticmethod + def _completion_kwargs(**overrides): + from litellm.types.utils import ModelResponse + + kwargs = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "api_base": "https://api.anthropic.com/v1/messages", + "custom_llm_provider": "anthropic", + "custom_prompt_dict": {}, + "model_response": ModelResponse(), + "print_verbose": lambda *_args, **_kwargs: None, + "encoding": None, + "api_key": "sk-ant-test", + "logging_obj": MagicMock(), + "optional_params": {"max_tokens": 16}, + "timeout": 30.0, + "litellm_params": {"rust": True}, + "acompletion": False, + "headers": {}, + "client": None, + } + kwargs.update(overrides) + return kwargs + + @staticmethod + def _recording_logging_obj(): + """A logging object that keeps each hook's payload in a real list, so a + test can assert which path logged and what it carried.""" + calls = {"pre_call": [], "post_call": []} + logging_obj = MagicMock() + logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs) + logging_obj.post_call.side_effect = lambda **kwargs: calls["post_call"].append(kwargs) + return logging_obj, calls + + def _inject(self, *, decline_reason=None, sync_result=None, sync_error=None): + from litellm.rust_bridge import chat_completions as bridge + + seen = {"gate": [], "call": []} + + def gate(**kwargs): + seen["gate"].append(kwargs) + return decline_reason + + def native(**kwargs): + seen["call"].append(kwargs) + if sync_error is not None: + raise sync_error + return dict(sync_result if sync_result is not None else self.RUST_RESPONSE) + + bridge.set_rust_chat_completions(decline=gate, chat_completions=native) + return seen + + def test_rust_true_serves_the_call_and_stamps_the_header(self): + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + + seen = self._inject() + response = AnthropicChatCompletion().completion(**self._completion_kwargs()) + + assert response.choices[0].message.content == "hello from rust" + assert response._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + assert len(seen["call"]) == 1 + + def test_the_core_receives_the_untranslated_openai_messages(self): + """Rust owns the translation, so the handler must not pre-translate.""" + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + + seen = self._inject() + AnthropicChatCompletion().completion( + **self._completion_kwargs( + messages=[ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + ] + ) + ) + assert seen["call"][0]["messages"] == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + ] + + def test_the_anthropic_max_tokens_default_is_merged_in_before_the_gate(self): + """`transform_request` applies `AnthropicConfig.get_config`; the Rust + path skips it, so the handler has to merge it or Anthropic 400s on a + request that omits `max_tokens`.""" + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + + seen = self._inject() + AnthropicChatCompletion().completion(**self._completion_kwargs(optional_params={})) + assert "max_tokens" in seen["gate"][0]["optional_params"] + assert seen["call"][0]["optional_params"]["max_tokens"] > 0 + + def test_a_caller_supplied_max_tokens_outranks_the_default(self): + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + + seen = self._inject() + AnthropicChatCompletion().completion( + **self._completion_kwargs(optional_params={"max_tokens": 7}) + ) + assert seen["call"][0]["optional_params"]["max_tokens"] == 7 + + def test_without_the_opt_in_the_core_is_never_consulted(self): + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + seen = self._inject() + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ) as transform, patch.object( + AnthropicChatCompletion, "acompletion_function" + ): + try: + AnthropicChatCompletion().completion( + **self._completion_kwargs(litellm_params={}) + ) + except Exception: + # The Python path goes on to make an HTTP call; reaching it is + # the assertion, so the network failure below is expected. + pass + assert seen["gate"] == [] + assert seen["call"] == [] + assert transform.called + + def test_a_declined_request_never_reaches_the_native_call(self): + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + seen = self._inject(decline_reason="unrecognized request parameter") + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): + try: + AnthropicChatCompletion().completion(**self._completion_kwargs()) + except Exception: + pass + assert len(seen["gate"]) == 1 + assert seen["call"] == [] + + def test_streaming_stays_on_the_python_path(self): + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + seen = self._inject() + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): + try: + AnthropicChatCompletion().completion( + **self._completion_kwargs(optional_params={"max_tokens": 16, "stream": True}) + ) + except Exception: + pass + assert seen["gate"] == [] + + def test_pre_call_logging_fires_exactly_once_on_the_rust_path(self): + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + + seen = self._inject() + logging_obj = MagicMock() + AnthropicChatCompletion().completion( + **self._completion_kwargs(logging_obj=logging_obj) + ) + assert logging_obj.pre_call.call_count == 1 + assert len(seen["call"]) == 1 + + def test_post_call_logging_fires_on_the_rust_path(self): + """The Rust core owns the provider call, so the Python transform that + normally raises `post_call` never runs. Without the bridge hook every + post_call callback goes silent and `original_response` stays unset.""" + import json + + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + + self._inject() + logging_obj = MagicMock() + AnthropicChatCompletion().completion( + **self._completion_kwargs(logging_obj=logging_obj) + ) + + assert logging_obj.post_call.call_count == 1 + logged = logging_obj.post_call.call_args.kwargs["original_response"] + assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust" + + def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines(self, monkeypatch): + """A decline never reached the provider, so the Python path serves the + request and owns the only post_call. Firing the hook there too would + double every post_call callback for one request.""" + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.rust_bridge import chat_completions as bridge + + class _Declined(Exception): + pass + + class _FakeNative: + RustBridgeDeclined = _Declined + RustUpstreamError = type("_Upstream", (Exception,), {}) + + def declining_native(**_kwargs): + raise _Declined("blank message text") + + monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, chat_completions=declining_native + ) + + logging_obj, calls = self._recording_logging_obj() + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): + try: + AnthropicChatCompletion().completion( + **self._completion_kwargs(logging_obj=logging_obj) + ) + except Exception: + # The Python path goes on to make an HTTP call; the log count is + # the assertion, so a failure past this point is expected. + pass + + assert calls["post_call"] == [] + + @pytest.mark.asyncio + async def test_the_async_path_falls_back_when_the_core_declines(self, monkeypatch): + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + from litellm.rust_bridge import chat_completions as bridge + + class _Declined(Exception): + pass + + class _FakeNative: + RustBridgeDeclined = _Declined + RustUpstreamError = type("_Upstream", (Exception,), {}) + + monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + + async def declining_native(**_kwargs): + raise _Declined("blank message text") + + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=declining_native + ) + + sentinel = object() + + async def python_path(**_kwargs): + return sentinel + + with patch.object( + AnthropicChatCompletion, "acompletion_function", side_effect=python_path + ) as python_call: + result = await AnthropicChatCompletion().completion( + **self._completion_kwargs(acompletion=True) + ) + + assert result is sentinel + assert python_call.called, "a failing rust call must re-enter the python path" + + @pytest.mark.asyncio + async def test_the_async_path_serves_the_rust_response_without_the_fallback(self): + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + from litellm.rust_bridge import chat_completions as bridge + + async def native(**_kwargs): + return dict(self.RUST_RESPONSE) + + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=native + ) + + with patch.object(AnthropicChatCompletion, "acompletion_function") as python_call: + result = await AnthropicChatCompletion().completion( + **self._completion_kwargs(acompletion=True) + ) + + assert result.choices[0].message.content == "hello from rust" + assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + assert not python_call.called + + + def test_pre_call_logging_fires_once_when_the_sync_rust_call_declines(self, monkeypatch): + """One request, one pre_call, on the synchronous path too. Without the + suppression the Python path logs a second time for the same attempt.""" + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.rust_bridge import chat_completions as bridge + + class _Declined(Exception): + pass + + class _FakeNative: + RustBridgeDeclined = _Declined + RustUpstreamError = type("_Upstream", (Exception,), {}) + + monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + + def declining_native(**_kwargs): + raise _Declined("blank message text") + + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, chat_completions=declining_native + ) + + logging_obj, calls = self._recording_logging_obj() + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): + try: + AnthropicChatCompletion().completion( + **self._completion_kwargs(logging_obj=logging_obj) + ) + except Exception: + # The Python path goes on to make an HTTP call; the log count is + # the assertion, so a failure past this point is expected. + pass + + assert len(calls["pre_call"]) == 1 + assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == ( + "claude-sonnet-4-5" + ) + + def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch): + """The suppression must not swallow the log on the ordinary path.""" + from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + self._inject() + logging_obj, calls = self._recording_logging_obj() + with patch.object( + AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} + ): + try: + AnthropicChatCompletion().completion( + **self._completion_kwargs(litellm_params={}, logging_obj=logging_obj) + ) + except Exception: + pass + + assert len(calls["pre_call"]) == 1 + assert calls["pre_call"][0]["additional_args"]["complete_input_dict"] == { + "model": "m", + "messages": [], + } diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py new file mode 100644 index 00000000000..8e67a7e3438 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -0,0 +1,489 @@ +"""Tests for `BedrockConverseLLM.completion`'s Rust chat completions hook. + +The native callables are dependency-injected, so these run without the compiled +extension, and AWS credential resolution is stubbed so nothing reaches STS. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from botocore.credentials import Credentials +from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.rust_bridge import chat_completions as bridge +from litellm.types.utils import ModelResponse + +RUST_RESPONSE = { + "created": 1_700_000_000, + "model": "anthropic.claude-sonnet-4-5-v1:0", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hello from rust"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 11, + "completion_tokens": 4, + "total_tokens": 15, + "prompt_tokens_details": { + "cached_tokens": 0, + "cache_creation_tokens": 0, + "text_tokens": 11, + }, + }, +} + +RESOLVED_CREDENTIALS = Credentials( + access_key="AKIARESOLVED", + secret_key="resolved-secret", + token="resolved-token", +) + + +@pytest.fixture(autouse=True) +def reset_bridge(monkeypatch): + monkeypatch.delenv("LITELLM_RUST", raising=False) + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) + yield + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) + + +def _inject(*, decline_reason=None, error: Exception | None = None): + seen: dict[str, list[dict]] = {"gate": [], "call": []} + + def gate(**kwargs): + seen["gate"].append(kwargs) + return decline_reason + + def native(**kwargs): + seen["call"].append(kwargs) + if error is not None: + raise error + return dict(RUST_RESPONSE) + + bridge.set_rust_chat_completions(decline=gate, chat_completions=native) + return seen + + +def _completion_kwargs(**overrides): + kwargs = { + "model": "bedrock/us-east-1/anthropic.claude-sonnet-4-5-v1:0", + "messages": [{"role": "user", "content": "hi"}], + "api_base": None, + "custom_prompt_dict": {}, + "model_response": ModelResponse(), + "encoding": None, + "logging_obj": MagicMock(), + "optional_params": {"maxTokens": 16}, + "acompletion": False, + "timeout": 30.0, + "litellm_params": {"rust": True}, + "extra_headers": None, + "client": None, + "api_key": None, + } + kwargs.update(overrides) + return kwargs + + +def _run(**overrides): + with patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ): + return BedrockConverseLLM().completion(**_completion_kwargs(**overrides)) + + +def _recording_logging_obj(): + """A logging object that keeps each hook's payload in a real list, so a test + can assert which path logged and what it carried.""" + calls = {"pre_call": [], "post_call": []} + logging_obj = MagicMock() + logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs) + logging_obj.post_call.side_effect = lambda **kwargs: calls["post_call"].append(kwargs) + return logging_obj, calls + + +def test_rust_true_serves_the_call_and_stamps_the_header(): + seen = _inject() + response = _run() + + assert response.choices[0].message.content == "hello from rust" + assert response._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + assert len(seen["call"]) == 1 + + +def test_the_core_receives_the_credentials_this_handler_already_resolved(): + """Both paths must sign as the same principal, so the resolved credentials + are handed down rather than re-derived from ambient AWS state.""" + seen = _inject() + _run() + + params = seen["call"][0]["optional_params"] + assert params["aws_access_key_id"] == "AKIARESOLVED" + assert params["aws_secret_access_key"] == "resolved-secret" + assert params["aws_session_token"] == "resolved-token" + assert params["aws_region_name"] == "us-east-1" + + +def test_the_core_receives_the_converse_url_this_handler_already_built(): + seen = _inject() + _run() + + assert seen["call"][0]["api_base"].endswith( + "/model/anthropic.claude-sonnet-4-5-v1%3A0/converse" + ) + assert "bedrock-runtime.us-east-1.amazonaws.com" in seen["call"][0]["api_base"] + + +def test_the_core_receives_the_untranslated_openai_messages(): + seen = _inject() + _run( + messages=[ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + ] + ) + assert seen["call"][0]["messages"] == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + ] + + +def test_without_the_opt_in_the_core_is_never_consulted(): + seen = _inject() + try: + _run(litellm_params={}) + except Exception: + # The Python path goes on to make an HTTP call; not reaching the gate + # is the assertion, so a failure past this point is expected. + pass + assert seen["gate"] == [] + assert seen["call"] == [] + + +def test_streaming_stays_on_the_python_path(): + seen = _inject() + try: + _run(optional_params={"maxTokens": 16, "stream": True}) + except Exception: + pass + assert seen["gate"] == [] + + +def test_a_declined_request_never_reaches_the_native_call(): + seen = _inject(decline_reason="unrecognized request parameter") + try: + _run() + except Exception: + pass + assert len(seen["gate"]) == 1 + assert seen["call"] == [] + + +def test_pre_call_logging_fires_exactly_once_on_the_rust_path(): + _inject() + logging_obj = MagicMock() + _run(logging_obj=logging_obj) + assert logging_obj.pre_call.call_count == 1 + + +@pytest.mark.asyncio +async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch): + class _Declined(Exception): + pass + + class _FakeNative: + RustBridgeDeclined = _Declined + RustUpstreamError = type("_Upstream", (Exception,), {}) + + monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + + async def declining_native(**_kwargs): + raise _Declined("blank message text") + + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=declining_native + ) + + sentinel = object() + + async def python_path(**_kwargs): + return sentinel + + with ( + patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ), + patch.object( + BedrockConverseLLM, "async_completion", side_effect=python_path + ) as python_call, + ): + result = await BedrockConverseLLM().completion( + **_completion_kwargs(acompletion=True) + ) + + assert result is sentinel + assert python_call.called, "a failing rust call must re-enter the python path" + + +@pytest.mark.asyncio +async def test_the_async_path_serves_the_rust_response_without_the_fallback(): + async def native(**_kwargs): + return dict(RUST_RESPONSE) + + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=native + ) + + with ( + patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ), + patch.object(BedrockConverseLLM, "async_completion") as python_call, + ): + result = await BedrockConverseLLM().completion( + **_completion_kwargs(acompletion=True) + ) + + assert result.choices[0].message.content == "hello from rust" + assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + assert not python_call.called + + +@pytest.mark.asyncio +async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines(): + """One request, one pre_call. Without the suppression the Python fallback + logs a second one and non-idempotent callbacks run twice.""" + + class _Declined(Exception): + pass + + class _FakeNative: + RustBridgeDeclined = _Declined + RustUpstreamError = type("_Upstream", (Exception,), {}) + + async def declining_native(**_kwargs): + raise _Declined("blank message text") + + logging_obj = MagicMock() + served = [] + + async def python_path(**kwargs): + served.append(kwargs) + return ModelResponse() + + with ( + patch.object(bridge, "get_native_bridge", lambda: _FakeNative()), + patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ), + patch.object( + BedrockConverseLLM, "async_completion", side_effect=python_path + ), + ): + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=declining_native + ) + await BedrockConverseLLM().completion( + **_completion_kwargs(acompletion=True, logging_obj=logging_obj) + ) + + assert logging_obj.pre_call.call_count == 1 + assert served and served[0]["skip_pre_call_logging"] is True + + +CONVERSE_RESPONSE = { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 2, "totalTokens": 7}, +} + + +async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): + """Run the real `async_completion` with a stubbed transport.""" + import httpx as _httpx + + client = MagicMock() + + async def post(**_kwargs): + return _httpx.Response( + 200, + json=CONVERSE_RESPONSE, + request=_httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + + client.post = post + client.__class__ = AsyncHTTPHandler + + return await BedrockConverseLLM().async_completion( + model="anthropic.claude-sonnet-4-5-v1:0", + messages=[{"role": "user", "content": "hi"}], + api_base="https://bedrock-runtime.us-west-2.amazonaws.com/model/m/converse", + model_response=ModelResponse(), + timeout=30.0, + encoding=None, + logging_obj=logging_obj, + stream=None, + optional_params={"maxTokens": 16}, + litellm_params={"aws_region_name": "us-west-2"}, + credentials=RESOLVED_CREDENTIALS, + headers={}, + client=client, + skip_pre_call_logging=skip_pre_call_logging, + ) + + +@pytest.mark.asyncio +async def test_async_completion_honors_the_pre_call_suppression(): + logging_obj = MagicMock() + await _drive_async_completion(skip_pre_call_logging=True, logging_obj=logging_obj) + assert logging_obj.pre_call.call_count == 0 + + +@pytest.mark.asyncio +async def test_async_completion_logs_pre_call_by_default(): + """The suppression must be opt-in, so every existing caller keeps its log.""" + logging_obj = MagicMock() + await _drive_async_completion(skip_pre_call_logging=False, logging_obj=logging_obj) + assert logging_obj.pre_call.call_count == 1 + + +def _sync_client_returning_converse_response(): + client = MagicMock() + client.post = lambda **_kwargs: httpx.Response( + 200, + json=CONVERSE_RESPONSE, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + client.__class__ = HTTPHandler + return client + + +def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines(): + """One request, one pre_call, on the synchronous path too. + + The gate accepts and logs, then the native call declines before the + provider is reached, so execution continues into the Python path below. + That is the same attempt continuing; without the suppression it logs a + second pre_call and non-idempotent callbacks run twice for one request. + """ + + class _Declined(Exception): + pass + + class _FakeNative: + RustBridgeDeclined = _Declined + RustUpstreamError = type("_Upstream", (Exception,), {}) + + def declining_native(**_kwargs): + raise _Declined("blank message text") + + logging_obj = MagicMock() + + with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()): + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, chat_completions=declining_native + ) + response = _run( + logging_obj=logging_obj, + client=_sync_client_returning_converse_response(), + ) + + assert response.choices[0].message.content == "hi" + assert logging_obj.pre_call.call_count == 1 + + +def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(): + """The suppression must not swallow the log on a request the gate declined, + so a deployment with no `rust` flag keeps exactly the log it always had.""" + logging_obj = MagicMock() + response = _run( + logging_obj=logging_obj, + litellm_params={}, + client=_sync_client_returning_converse_response(), + ) + + assert response.choices[0].message.content == "hi" + assert logging_obj.pre_call.call_count == 1 + + +def test_post_call_logging_fires_on_the_sync_rust_path(): + """The Rust core owns the provider call, so the Converse transform that + normally raises `post_call` never runs. Without the bridge hook every + post_call callback goes silent and `original_response` stays unset.""" + import json + + _inject() + logging_obj = MagicMock() + _run(logging_obj=logging_obj) + + assert logging_obj.post_call.call_count == 1 + logged = logging_obj.post_call.call_args.kwargs["original_response"] + assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust" + + +@pytest.mark.asyncio +async def test_post_call_logging_fires_on_the_async_rust_path(): + """The asynchronous path runs through the same hook, so the two paths + cannot drift apart the way the pre_call suppression once did.""" + import json + + async def native(**_kwargs): + return dict(RUST_RESPONSE) + + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, achat_completions=native + ) + logging_obj = MagicMock() + + with patch.object( + BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS + ): + await BedrockConverseLLM().completion( + **_completion_kwargs(acompletion=True, logging_obj=logging_obj) + ) + + assert logging_obj.post_call.call_count == 1 + logged = logging_obj.post_call.call_args.kwargs["original_response"] + assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust" + + +def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines(): + """A decline never reached the provider, so the Python path serves the + request and owns the only post_call. Firing the hook there too would double + every post_call callback for one request.""" + + class _Declined(Exception): + pass + + class _FakeNative: + RustBridgeDeclined = _Declined + RustUpstreamError = type("_Upstream", (Exception,), {}) + + def declining_native(**_kwargs): + raise _Declined("blank message text") + + logging_obj, calls = _recording_logging_obj() + + with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()): + bridge.set_rust_chat_completions( + decline=lambda **_kwargs: None, chat_completions=declining_native + ) + response = _run( + logging_obj=logging_obj, + client=_sync_client_returning_converse_response(), + ) + + assert response.choices[0].message.content == "hi" + assert len(calls["post_call"]) == 1 + assert "hi" in calls["post_call"][0]["original_response"] diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 3102e69bf26..ecf7f89d487 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -2226,6 +2226,66 @@ class TestIsRequestBodySafeBlocksBedrockProjectOverride: ) +class TestIsRequestBodySafeBlocksRustOptIn: + """``rust`` hands the whole call to the Rust core, which signs and sends + with its own HTTP client rather than the one the deployment configured, and + reports no ``post_call``. The proxy splats the request body straight into + the router, and ``rust`` is a litellm param, so it lands in + ``litellm_params`` and the gate honours it: without this entry any + authenticated caller picks a transport and a callback surface the admin + never chose. It stays a deployment decision, liftable only by the same + admin opt-in as the rest of the list.""" + + def test_rust_in_request_body_is_rejected(self): + with pytest.raises(ValueError, match="rust"): + is_request_body_safe( + request_body={"model": "gpt-4", "rust": True}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_rust_under_extra_body_is_rejected(self): + with pytest.raises(ValueError, match="not allowed in request body"): + is_request_body_safe( + request_body={"model": "gpt-4", "extra_body": {"rust": True}}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_api_key_does_not_bypass_the_rust_block(self): + with pytest.raises(ValueError, match="rust"): + is_request_body_safe( + request_body={"model": "gpt-4", "api_key": "sk-anything", "rust": True}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_admin_opt_in_proxy_wide_allows_rust(self): + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", "rust": True}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + def test_body_without_rust_is_still_allowed(self): + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", "temperature": 0.7}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + class TestIsRequestBodySafeBlocksVertexCredentialAlias: @pytest.mark.parametrize("field", ["vertex_ai_credentials"]) def test_field_in_request_body_is_rejected(self, field): diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py new file mode 100644 index 00000000000..47cb66932b7 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -0,0 +1,420 @@ +"""Tests for the Rust chat completions bridge. + +The native callables are dependency-injected through +``set_rust_chat_completions`` rather than patched, so these run without the +compiled extension present. +""" + +from __future__ import annotations + +import pytest + +import litellm +from litellm.rust_bridge import chat_completions as bridge +from litellm.types.utils import ModelResponse + +RUST_RESPONSE = { + "created": 1_700_000_000, + "model": "claude-sonnet-4-5-20260101", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hello from rust"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 11, + "completion_tokens": 4, + "total_tokens": 15, + "prompt_tokens_details": { + "cached_tokens": 0, + "cache_creation_tokens": 0, + "text_tokens": 11, + }, + }, +} + +MESSAGES = [{"role": "user", "content": "hi"}] + + +class _FakeDeclined(Exception): + """Stands in for the native `RustBridgeDeclined`.""" + + +class _FakeUpstream(Exception): + """Stands in for the native `RustUpstreamError`; args are (status, message).""" + + +class _FakeNative: + RustBridgeDeclined = _FakeDeclined + RustUpstreamError = _FakeUpstream + + +def _fake_native_bridge(monkeypatch): + """Expose the bridge's exception classes without the compiled extension.""" + monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + + +def _hide_native_bridge(monkeypatch): + """Simulate a wheel built without the compiled extension. + + There is no injection seam for "the .so is absent", so the loader itself is + replaced; every other case here uses `set_rust_chat_completions`. + """ + monkeypatch.setattr(bridge, "get_native_bridge", lambda: None) + + +@pytest.fixture(autouse=True) +def reset_bridge(): + """Every test starts with no injected callables, and leaves none behind.""" + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) + yield + bridge.set_rust_chat_completions( + chat_completions=None, achat_completions=None, decline=None + ) + + +class _RecordingDecline: + """A stand-in for the native gate that records what it was asked.""" + + def __init__(self, reason: str | None = None): + self.reason = reason + self.calls: list[dict] = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + return self.reason + + +class _RecordingCall: + def __init__(self, result=None, error: Exception | None = None): + self.result = result if result is not None else dict(RUST_RESPONSE) + self.error = error + self.calls: list[dict] = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + if self.error is not None: + raise self.error + return self.result + + +class _RecordingAsyncCall(_RecordingCall): + async def __call__(self, **kwargs): + return _RecordingCall.__call__(self, **kwargs) + + +def _accepts(**overrides) -> bool: + kwargs = { + "model": "claude-sonnet-4-5", + "messages": MESSAGES, + "optional_params": {"max_tokens": 16}, + "custom_llm_provider": "anthropic", + "litellm_params": {"rust": True}, + "stream": None, + } + kwargs.update(overrides) + return bridge.rust_chat_completions_accepts(**kwargs) + + +class TestGate: + def test_declines_when_the_deployment_did_not_opt_in(self, monkeypatch): + monkeypatch.delenv("LITELLM_RUST", raising=False) + gate = _RecordingDecline() + bridge.set_rust_chat_completions(decline=gate) + assert _accepts(litellm_params={}) is False + assert _accepts(litellm_params=None) is False + assert _accepts(litellm_params={"rust": False}) is False + assert gate.calls == [], "the gate must not be consulted before opt-in" + + def test_accepts_when_the_deployment_opted_in_and_the_core_agrees(self, monkeypatch): + monkeypatch.delenv("LITELLM_RUST", raising=False) + gate = _RecordingDecline() + bridge.set_rust_chat_completions(decline=gate) + assert _accepts() is True + assert gate.calls[0]["model"] == "claude-sonnet-4-5" + assert gate.calls[0]["custom_llm_provider"] == "anthropic" + + def test_the_env_var_opts_in_without_a_per_model_flag(self, monkeypatch): + monkeypatch.setenv("LITELLM_RUST", "true") + bridge.set_rust_chat_completions(decline=_RecordingDecline()) + assert _accepts(litellm_params={}) is True + + def test_declines_streaming_and_providers_off_the_path(self, monkeypatch): + monkeypatch.delenv("LITELLM_RUST", raising=False) + gate = _RecordingDecline() + bridge.set_rust_chat_completions(decline=gate) + assert _accepts(stream=True) is False + assert _accepts(custom_llm_provider="openai") is False + assert _accepts(custom_llm_provider=None) is False + assert gate.calls == [] + + def test_declines_an_anthropic_request_carrying_a_litellm_metadata_user_id(self, monkeypatch): + """`AnthropicConfig.transform_request` copies a valid `user_id` into the Messages body. + + It does that inside the function the Rust route replaces, and the core is + handed `optional_params` only, so accepting here would send the request + to Anthropic with the abuse-detection attribution silently missing. + """ + monkeypatch.delenv("LITELLM_RUST", raising=False) + gate = _RecordingDecline() + bridge.set_rust_chat_completions(decline=gate) + assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}) is False + assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of" + + # Bedrock's Converse transform reads no `user_id`, and an Anthropic request + # whose metadata carries none is one Python would not attribute either. + assert ( + _accepts( + custom_llm_provider="bedrock", + model="bedrock/us-east-1/anthropic.claude-v2", + litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}, + ) + is True + ) + assert _accepts(litellm_params={"rust": True, "metadata": {"trace_id": "t-1"}}) is True + assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": None}}) is True + assert _accepts(litellm_params={"rust": True, "metadata": None}) is True + + def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch): + """`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the + Converse body from `litellm_params`, and owning that field also means + evicting a caller-supplied one. The core can do neither, so an operator + who armed `bedrock_request_metadata_fields` keeps the Python path. + """ + monkeypatch.delenv("LITELLM_RUST", raising=False) + gate = _RecordingDecline() + bridge.set_rust_chat_completions(decline=gate) + bedrock = { + "custom_llm_provider": "bedrock", + "model": "bedrock/us-east-1/anthropic.claude-v2", + } + + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ["user_api_key_team_id"]) + assert _accepts(**bedrock) is False + assert gate.calls == [], "the core must not be consulted for a field it cannot write" + assert _accepts() is True, "arming Bedrock attribution must not decline Anthropic" + + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + assert _accepts(**bedrock) is True, "the decline follows the operator's opt-in alone" + + def test_declines_when_the_core_declines(self, monkeypatch): + monkeypatch.delenv("LITELLM_RUST", raising=False) + bridge.set_rust_chat_completions(decline=_RecordingDecline("streaming")) + assert _accepts() is False + + def test_declines_when_the_bridge_is_unavailable(self, monkeypatch): + monkeypatch.delenv("LITELLM_RUST", raising=False) + _hide_native_bridge(monkeypatch) + assert _accepts() is False + + def test_declines_when_the_gate_itself_raises(self, monkeypatch): + monkeypatch.delenv("LITELLM_RUST", raising=False) + + def exploding(**_kwargs): + raise RuntimeError("boom") + + bridge.set_rust_chat_completions(decline=exploding) + assert _accepts() is False + + +def _call_kwargs(model_response: ModelResponse) -> dict: + return { + "model": "claude-sonnet-4-5", + "messages": MESSAGES, + "optional_params": {"max_tokens": 16}, + "model_response": model_response, + "api_key": "sk-test", + "api_base": None, + "custom_llm_provider": "anthropic", + "extra_headers": {}, + "timeout": 30.0, + "on_response": lambda _rust_response: None, + } + + +class TestSyncCall: + def test_builds_a_model_response_and_stamps_the_rust_header(self): + native = _RecordingCall() + bridge.set_rust_chat_completions(chat_completions=native) + model_response = ModelResponse() + original_id = model_response.id + + result = bridge.chat_completions(**_call_kwargs(model_response)) + + assert result is not None + assert result.choices[0].message.content == "hello from rust" + assert result.choices[0].finish_reason == "stop" + assert result.model == "claude-sonnet-4-5-20260101" + assert result.usage.prompt_tokens == 11 + assert result.usage.completion_tokens == 4 + assert result.usage.total_tokens == 15 + assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + assert result.id == original_id, ( + "the rust path must keep the chatcmpl id litellm already minted" + ) + + def test_passes_the_timeout_through_as_seconds(self): + native = _RecordingCall() + bridge.set_rust_chat_completions(chat_completions=native) + bridge.chat_completions(**_call_kwargs(ModelResponse())) + assert native.calls[0]["timeout_seconds"] == 30.0 + + def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch): + _hide_native_bridge(monkeypatch) + assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None + + def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch): + _fake_native_bridge(monkeypatch) + bridge.set_rust_chat_completions( + chat_completions=_RecordingCall(error=_FakeDeclined("streaming")) + ) + assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None + + +class TestAsyncCall: + @pytest.mark.asyncio + async def test_builds_a_model_response(self): + bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall()) + result = await bridge.achat_completions(**_call_kwargs(ModelResponse())) + assert result is not None + assert result.choices[0].message.content == "hello from rust" + assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + + @pytest.mark.asyncio + async def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch): + _hide_native_bridge(monkeypatch) + assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None + + @pytest.mark.asyncio + async def test_falls_back_when_the_core_declines_before_calling_the_provider( + self, monkeypatch + ): + _fake_native_bridge(monkeypatch) + bridge.set_rust_chat_completions( + achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")) + ) + assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None + + +class TestAsyncFallbackWrapper: + @pytest.mark.asyncio + async def test_returns_the_rust_response_without_running_the_fallback(self): + bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall()) + ran = [] + + async def fallback(): + ran.append(True) + return "python" + + result = await bridge.achat_completions_or_fallback( + **_call_kwargs(ModelResponse()), python_fallback=fallback + ) + assert result.choices[0].message.content == "hello from rust" + assert ran == [] + + @pytest.mark.asyncio + async def test_runs_the_fallback_when_the_core_declines(self, monkeypatch): + _fake_native_bridge(monkeypatch) + bridge.set_rust_chat_completions( + achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")) + ) + + async def fallback(): + return "python" + + result = await bridge.achat_completions_or_fallback( + **_call_kwargs(ModelResponse()), python_fallback=fallback + ) + assert result == "python" + + @pytest.mark.asyncio + async def test_runs_the_fallback_when_the_bridge_is_unavailable(self, monkeypatch): + _hide_native_bridge(monkeypatch) + + async def fallback(): + return "python" + + result = await bridge.achat_completions_or_fallback( + **_call_kwargs(ModelResponse()), python_fallback=fallback + ) + assert result == "python" + + +class TestFailureClassification: + """A failure the provider already saw must not be retried on the Python + path: it would bill the customer for the same work twice.""" + + @pytest.fixture(autouse=True) + def _native_exceptions(self, monkeypatch): + _fake_native_bridge(monkeypatch) + + def test_a_decline_falls_back_because_nothing_was_sent(self): + bridge.set_rust_chat_completions( + chat_completions=_RecordingCall(error=_FakeDeclined("streaming")) + ) + assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None + + def test_an_upstream_failure_is_surfaced_with_its_status(self): + from litellm.exceptions import APIError + + bridge.set_rust_chat_completions( + chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited")) + ) + with pytest.raises(APIError) as raised: + bridge.chat_completions(**_call_kwargs(ModelResponse())) + assert raised.value.status_code == 429 + assert "rate limited" in str(raised.value) + + def test_a_transport_failure_with_no_response_surfaces_as_a_500(self): + from litellm.exceptions import APIError + + bridge.set_rust_chat_completions( + chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset")) + ) + with pytest.raises(APIError) as raised: + bridge.chat_completions(**_call_kwargs(ModelResponse())) + assert raised.value.status_code == 500 + + def test_an_unrecognized_error_is_not_swallowed(self): + bridge.set_rust_chat_completions( + chat_completions=_RecordingCall(error=RuntimeError("something else")) + ) + with pytest.raises(RuntimeError): + bridge.chat_completions(**_call_kwargs(ModelResponse())) + + @pytest.mark.asyncio + async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self): + from litellm.exceptions import APIError + + bridge.set_rust_chat_completions( + achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom")) + ) + ran = [] + + async def fallback(): + ran.append(True) + return "python" + + with pytest.raises(APIError): + await bridge.achat_completions_or_fallback( + **_call_kwargs(ModelResponse()), python_fallback=fallback + ) + assert ran == [], "a request the provider already served must not be re-issued" + + @pytest.mark.asyncio + async def test_the_async_wrapper_falls_back_on_a_decline(self): + bridge.set_rust_chat_completions( + achat_completions=_RecordingAsyncCall(error=_FakeDeclined("blank message text")) + ) + + async def fallback(): + return "python" + + result = await bridge.achat_completions_or_fallback( + **_call_kwargs(ModelResponse()), python_fallback=fallback + ) + assert result == "python"