diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b9d363e72ca..d67623feffd 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3138,6 +3138,7 @@ dependencies = [ "litellm-http", "litellm-llms", "litellm-secrets", + "litellm-tracing", "litellm-types", "mime_guess", "moka", @@ -3215,6 +3216,8 @@ name = "litellm-gateway" version = "0.1.0" dependencies = [ "axum", + "futures-util", + "http-body-util", "litellm-config", "litellm-core", "litellm-gateway-auth", @@ -3222,11 +3225,13 @@ dependencies = [ "litellm-http", "litellm-llms", "litellm-secrets", + "litellm-tracing", "rstest", "serde_json", "tokio", - "tower-http 0.7.1", + "tower", "tracing", + "uuid", ] [[package]] @@ -3256,7 +3261,6 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-core", - "litellm-host", "litellm-http", "litellm-llms", "litellm-router", @@ -3677,6 +3681,7 @@ dependencies = [ name = "litellm-tracing" version = "0.1.0" dependencies = [ + "base64 0.22.1", "fancy-regex 0.19.2", "percent-encoding", "rstest", @@ -4880,7 +4885,7 @@ dependencies = [ "tokio-rustls 0.26.4", "tokio-util", "tower", - "tower-http 0.6.11", + "tower-http", "tower-service", "url", "wasm-bindgen", @@ -4922,7 +4927,7 @@ dependencies = [ "tokio-rustls 0.26.4", "tokio-util", "tower", - "tower-http 0.6.11", + "tower-http", "tower-service", "url", "wasm-bindgen", @@ -6163,23 +6168,6 @@ dependencies = [ "url", ] -[[package]] -name = "tower-http" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c" -dependencies = [ - "bitflags 2.13.1", - "bytes", - "http 1.4.2", - "http-body 1.1.0", - "percent-encoding", - "pin-project-lite", - "tower-layer", - "tower-service", - "tracing", -] - [[package]] name = "tower-layer" version = "0.3.3" diff --git a/litellm-rust/crates/auth-types/src/http.rs b/litellm-rust/crates/auth-types/src/http.rs index f3c5254b60e..1519769521f 100644 --- a/litellm-rust/crates/auth-types/src/http.rs +++ b/litellm-rust/crates/auth-types/src/http.rs @@ -7,7 +7,7 @@ pub enum CredentialPlacement { } impl CredentialPlacement { - pub fn header_name(self) -> &'static str { + pub const fn header_name(self) -> &'static str { match self { Self::Bearer => "Authorization", Self::Header(name) => name, diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index a40265729c7..20da74789bb 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -1,4 +1,6 @@ -litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src//` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back. +litellm-core is the LiteLLM SDK in Rust. Each top-level call is a module under `src//` exposing a public entrypoint named after the route. `messages::messages()` returns `MessagesResponse::Message` for a completed response or `MessagesResponse::Stream { headers, chunks }` when the request sets `stream: true`. The chunks are Anthropic SSE bytes in a `Stream>`. Dropping the stream cancels the call. The Python bridge drives `messages::route::messages_machine()` instead, because Python has to answer the call's operations on its own thread; the gateway and the Rust SDK call the plain entrypoint + +A route module has the same five pieces, in the order Python runs them. `types.rs` holds the call, the provider request, and the response. `prepare.rs` resolves the provider and credentials and shapes the request (Python's `validate_environment`, `get_complete_url`, `transform_request`). `handler.rs` resolves auth, offers the wire request to `litellm_host::hooks::RouteHooks::before_send`, sends it, reports the raw response through `emit`, and normalizes the response or stream (`pre_call`, `post`, `post_call`, `transform_response`). `mod.rs` exposes the entrypoint that runs prepare then handler with no hooks (`()`). `route.rs`, where a host needs it, wraps the same two calls in a `CallMachine` whose `HostChannel` is the hooks, and pumps a stream through `open` and `deliver`. A handler takes `&impl RouteHooks` and never a `HostChannel` directly, so it runs without a coroutine. Keep provider transport and transformation details out of the machine driver ## Crate layering @@ -10,7 +12,7 @@ Each crate mirrors one top-level Python package, so a Rust path reads as its Pyt - `litellm-llms` mirrors `litellm/llms/`: `base_llm//transformation.rs`, `//transformation.rs`, and `base_llm/ocr/handler.rs` (the OCR request handler) - `litellm-core` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...): entrypoints, route request types, provider dispatch, the route machine, and hooks -A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate +A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, the provider-level hooks OCR implements over its host until it folds into `litellm_host::hooks::RouteHooks`. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate ## Error placement diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 0051631f40d..6904dcc023c 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -17,6 +17,7 @@ litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-auth-aws.workspace = true litellm-http.workspace = true litellm-llms.workspace = true +litellm-tracing.workspace = true moka.workspace = true mime_guess = "2.0.5" rand.workspace = true diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 4a1cf7e193e..740db2edefb 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,23 +1,68 @@ use std::time::Duration; +use litellm_auth::AuthServices; +use litellm_host::{ + event::{MachineEvent, RawResponse, RequestContext, WireRequest}, + hooks::RouteHooks, +}; use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body}; -use litellm_llms::base_llm::{auth::resolve_auth, chat::transformation::ProviderChatResponseData}; +use litellm_llms::base_llm::{ + auth::{Authenticated, resolve_auth}, + chat::transformation::ProviderChatResponseData, +}; use litellm_types::utils::ChatCompletionsResponse; use serde_json::Value; -use super::{Error, prepare::prepare_provider_request}; +use super::Error; use crate::{ - chat_completions::types::{ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest}, + chat_completions::types::ProviderChatCompletionsRequest, constants::CHAT_COMPLETIONS_TIMEOUT_SECS, }; -pub(super) async fn execute_chat_completions_provider_call( +pub(super) async fn execute( http: &Client, - auth: &litellm_auth::AuthServices, - request: ResolvedChatCompletionsRequest<'_>, + auth: &AuthServices, + request: ProviderChatCompletionsRequest, + hooks: &impl RouteHooks, ) -> Result { - let request = prepare_provider_request(request)?; - let outbound = outbound_request(auth, &request).await?; + let ProviderChatCompletionsRequest { + model, + custom_llm_provider, + config, + url, + body, + optional_params, + environment, + timeout, + api_key, + } = request; + let context = RequestContext { + model: model.clone(), + custom_llm_provider, + optional_params: Value::Object(optional_params), + secret_fields: Vec::new(), + api_key, + }; + let authenticated = resolve_auth(auth, environment, &|key| std::env::var(key).ok()).await?; + let wire = hooks + .before_send( + WireRequest { + url, + headers: authenticated.headers, + body, + }, + context, + ) + .await?; + let outbound = outbound_request( + Authenticated { + headers: wire.headers, + signer: authenticated.signer, + }, + wire.url, + &wire.body, + timeout, + )?; let response = outbound.send(http).await.map_err(|err| { // Failing to establish the connection means the request never went out, @@ -41,13 +86,17 @@ pub(super) async fn execute_chat_completions_provider_call( body: truncate_error_body(&text), })); } + hooks + .emit(MachineEvent::ResponseReceived { + raw: RawResponse { body: text.clone() }, + }) + .await?; let body: Value = serde_json::from_str(&text).map_err(|err| { Error::InvalidResponse(format!("invalid chat completions response JSON: {err}")) })?; - request - .config - .transform_response(&request.model, ProviderChatResponseData { body }) + config + .transform_response(&model, ProviderChatResponseData { body }) .map_err(Error::from) .map_err(as_response_error) } @@ -69,21 +118,17 @@ pub(super) fn as_response_error(err: Error) -> Error { } } -pub(super) async fn outbound_request( - auth: &litellm_auth::AuthServices, - request: &ProviderChatCompletionsRequest, +pub(super) fn outbound_request( + authenticated: Authenticated, + url: String, + body: &Value, + timeout: Option, ) -> Result { - let env_lookup = |key: &str| std::env::var(key).ok(); - let authenticated = resolve_auth(auth, request.environment.clone(), &env_lookup).await?; crate::outbound::outbound_request( authenticated, - request.url.clone(), - &request.body, - Some( - request - .timeout - .unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)), - ), + url, + body, + Some(timeout.unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS))), ) .map_err(|error| match error { // Python drops the caller's copy and prefers a forwarded Authorization @@ -97,7 +142,133 @@ pub(super) async fn outbound_request( #[cfg(test)] mod tests { - use super::{Error, as_response_error}; + use std::sync::Mutex; + + use rstest::rstest; + use serde_json::json; + use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any}; + + use super::*; + use crate::chat_completions::{ + prepare::{prepare_provider_request, resolve_request}, + types::ChatCompletionsRequest, + }; + + const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; + + /// Rewrites the outgoing request and records what the call reports back. + #[derive(Default)] + struct RecordingHooks { + contexts: Mutex>, + raw: Mutex>, + } + + impl RouteHooks for RecordingHooks { + async fn before_send( + &self, + wire: WireRequest, + context: RequestContext, + ) -> Result { + self.contexts.lock().unwrap().push(context); + let mut body = wire.body; + body["system"] = json!("added by the host"); + Ok(WireRequest { + headers: wire + .headers + .into_iter() + .chain([("x-host".to_string(), "seen".to_string())]) + .collect(), + body, + ..wire + }) + } + + async fn emit(&self, event: MachineEvent) -> Result<(), Error> { + let MachineEvent::ResponseReceived { raw } = event; + self.raw.lock().unwrap().push(raw.body); + Ok(()) + } + } + + fn prepared(api_base: &str) -> ProviderChatCompletionsRequest { + prepare_provider_request( + resolve_request(ChatCompletionsRequest { + model: "anthropic/claude-sonnet-4-5", + messages: json!([{"role": "user", "content": "hi"}]), + optional_params: json!({"max_tokens": 16}).as_object().unwrap().clone(), + api_key: Some("sk-test"), + api_base: Some(api_base), + custom_llm_provider: None, + extra_headers: None, + timeout: None, + }) + .unwrap(), + ) + .unwrap() + } + + #[rstest] + #[tokio::test] + async fn the_hooks_rewrite_the_wire_request_and_see_the_raw_response() { + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with( + ResponseTemplate::new(200).set_body_raw(ANTHROPIC_MESSAGE, "application/json"), + ) + .mount(&upstream) + .await; + let hooks = RecordingHooks::default(); + + execute( + &Client::plain_for_test(), + &AuthServices::default(), + prepared(&upstream.uri()), + &hooks, + ) + .await + .expect("chat completions call succeeds"); + + let [request] = <[Request; 1]>::try_from(upstream.received_requests().await.unwrap()) + .unwrap_or_else(|requests| panic!("one request, saw {}", requests.len())); + let sent: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(sent["system"], "added by the host"); + assert_eq!(request.headers["x-host"], "seen"); + assert_eq!(request.headers["x-api-key"], "sk-test"); + let [context] = <[RequestContext; 1]>::try_from(hooks.contexts.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_send runs once, saw {}", seen.len())); + assert_eq!( + (context.model.as_str(), context.custom_llm_provider.as_str()), + ("claude-sonnet-4-5", "anthropic") + ); + assert_eq!(context.optional_params, json!({"max_tokens": 16})); + assert_eq!(hooks.raw.into_inner().unwrap(), [ANTHROPIC_MESSAGE]); + } + + #[rstest] + #[tokio::test] + async fn an_upstream_failure_is_not_reported_as_a_received_response() { + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(500).set_body_string("boom")) + .mount(&upstream) + .await; + let hooks = RecordingHooks::default(); + + let error = execute( + &Client::plain_for_test(), + &AuthServices::default(), + prepared(&upstream.uri()), + &hooks, + ) + .await + .expect_err("the upstream failure fails the call"); + + assert!(matches!( + error, + Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) + )); + assert!(hooks.raw.into_inner().unwrap().is_empty()); + } #[test] fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index d7003b5d22a..dc4e80b816a 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -11,10 +11,9 @@ pub use crate::error::RouteError as Error; mod common_utils; pub(crate) mod handler; mod prepare; -use handler::execute_chat_completions_provider_call; use litellm_http::{ClientVariant, HttpClientConfig}; use litellm_types::utils::ChatCompletionsResponse; -use prepare::{parse_messages, resolve_provider_config, resolve_request}; +use prepare::{parse_messages, prepare_provider_request, resolve_provider_config, resolve_request}; use serde_json::{Map, Value}; use crate::chat_completions::types::ChatCompletionsRequest; @@ -24,9 +23,9 @@ pub async fn chat_completions( config: &HttpClientConfig, request: ChatCompletionsRequest<'_>, ) -> Result { - let request = resolve_request(request)?; let http = resources.pool.client(config, ClientVariant::Provider)?; - execute_chat_completions_provider_call(&http, &resources.auth, request).await + let request = prepare_provider_request(resolve_request(request)?)?; + handler::execute(&http, &resources.auth, request, &()).await } /// Whether the core would accept this request, without resolving credentials or @@ -41,9 +40,10 @@ pub fn chat_completions_decline_reason( messages: Value, optional_params: &Map, ) -> Option<&'static str> { - let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else { + let Ok(resolved) = resolve_provider_config(model, custom_llm_provider) else { return Some("provider is not on the rust chat completions path"); }; + let config = resolved.config; let Ok(messages) = parse_messages(messages) else { return Some("unreadable message list"); }; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 8b091d4dd6c..6ead713eec1 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,3 +1,4 @@ +use litellm_auth::SecretValue; use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_llms::base_llm::{ auth::{ValidatedEnvironment, with_default_headers}, @@ -14,10 +15,16 @@ use crate::chat_completions::types::{ ChatCompletionsRequest, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest, }; +pub(super) struct ResolvedProvider { + pub(super) model: String, + pub(super) custom_llm_provider: String, + pub(super) config: &'static dyn BaseConfig, +} + pub(super) fn resolve_provider_config<'a>( model: &'a str, custom_llm_provider: Option<&'a str>, -) -> Result<(String, &'static dyn BaseConfig), Error> { +) -> Result { let provider_info = get_custom_llm_provider(model, custom_llm_provider) .or_else(|| { custom_llm_provider.map(|provider| CustomLlmProvider { @@ -32,7 +39,11 @@ pub(super) fn resolve_provider_config<'a>( })?; let config = chat_completions_provider_config(provider_info.custom_llm_provider) .ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; - Ok((provider_info.model.to_string(), config)) + Ok(ResolvedProvider { + model: provider_info.model.to_string(), + custom_llm_provider: provider_info.custom_llm_provider.to_string(), + config, + }) } pub(super) fn parse_messages(messages: Value) -> Result, Error> { @@ -43,7 +54,11 @@ pub(super) fn parse_messages(messages: Value) -> Result, Error> pub(super) fn resolve_request( request: ChatCompletionsRequest<'_>, ) -> Result, Error> { - let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?; + let ResolvedProvider { + model, + custom_llm_provider, + config, + } = resolve_provider_config(request.model, request.custom_llm_provider)?; let messages = parse_messages(request.messages)?; if messages.is_empty() { return Err(Error::InvalidRequest( @@ -55,6 +70,7 @@ pub(super) fn resolve_request( } Ok(ResolvedChatCompletionsRequest { model, + custom_llm_provider, config, messages, optional_params: request.optional_params, @@ -99,15 +115,18 @@ pub(super) fn prepare_provider_request( &env_lookup, )?; let transformed = - config.transform_request(&model, request.messages, request.optional_params)?; + config.transform_request(&model, request.messages, request.optional_params.clone())?; Ok(ProviderChatCompletionsRequest { model, + custom_llm_provider: request.custom_llm_provider, config, url, body: transformed.body, + optional_params: request.optional_params, environment, timeout: request.timeout, + api_key: request.api_key.map(|key| SecretValue::new(key.to_string())), }) } @@ -449,11 +468,19 @@ mod tests { json!("abc-123"), )])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let signed = crate::chat_completions::handler::outbound_request( + let authenticated = resolve_auth( &litellm_auth::AuthServices::default(), - &prepared, + prepared.environment, + &|_| None, ) .await + .expect("resolves"); + let signed = crate::chat_completions::handler::outbound_request( + authenticated, + prepared.url, + &prepared.body, + prepared.timeout, + ) .expect("signs"); let authorization = signed @@ -502,11 +529,19 @@ mod tests { 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 = crate::chat_completions::handler::outbound_request( + let authenticated = resolve_auth( &litellm_auth::AuthServices::default(), - &prepared, + prepared.environment, + &|_| None, ) .await + .expect("resolves"); + let error = crate::chat_completions::handler::outbound_request( + authenticated, + prepared.url, + &prepared.body, + prepared.timeout, + ) .expect_err("{forwarded} should decline instead of being signed"); assert!( matches!(error, Error::Unsupported(_)), diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 66e9498c749..8c969dee730 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -1,5 +1,6 @@ use std::time::Duration; +use litellm_auth::SecretValue; use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig}; use litellm_types::llms::openai::ChatMessage; use serde_json::{Map, Value}; @@ -23,6 +24,7 @@ pub struct ChatCompletionsRequest<'a> { pub struct ResolvedChatCompletionsRequest<'a> { pub model: String, + pub custom_llm_provider: String, pub config: &'static dyn BaseConfig, pub messages: Vec, pub optional_params: Map, @@ -34,11 +36,16 @@ pub struct ResolvedChatCompletionsRequest<'a> { pub struct ProviderChatCompletionsRequest { pub model: String, + pub custom_llm_provider: String, pub config: &'static dyn BaseConfig, pub url: String, pub body: Value, + /// The route's parameters before the provider transformation, reported to the host + /// beside the wire request. + pub optional_params: Map, /// The forwarded and default headers plus how the call authenticates; the credential /// itself is applied when the request is sent. pub environment: ValidatedEnvironment, pub timeout: Option, + pub api_key: Option, } diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 650447d5abd..6832a59c7bd 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,20 +1,113 @@ use std::time::Duration; +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, stream::BoxStream}; +use litellm_auth::AuthServices; +use litellm_host::{ + event::{MachineEvent, RawResponse, RequestContext, WireRequest}, + hooks::RouteHooks, +}; use litellm_http::transport::Error as TransportError; use litellm_llms::base_llm::{ - anthropic_messages::transformation::BaseAnthropicMessagesConfig, auth::Authenticated, + anthropic_messages::{ + streaming::{ByteStream, StreamDecoder, encode_anthropic_sse}, + transformation::BaseAnthropicMessagesConfig, + }, + auth::{Authenticated, resolve_auth}, }; +use litellm_tracing::{ByteChunk, debug}; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use serde_json::Value; -use super::{Error, common_utils::truncate_error_body}; +use super::{ + Error, MessagesResponse, common_utils::truncate_error_body, prepare::ProviderMessagesRequest, +}; use crate::{constants::MESSAGES_TIMEOUT_SECS, outbound::outbound_request}; -pub(super) fn network(error: reqwest::Error) -> Error { +pub(super) async fn execute( + http: &litellm_http::Client, + auth: &AuthServices, + request: ProviderMessagesRequest, + hooks: &impl RouteHooks, +) -> Result { + let ProviderMessagesRequest { + provider, + url, + body, + environment, + timeout, + api_key, + } = request; + let stream = body.params.stream == Some(true); + let context = RequestContext { + model: body.model.clone(), + custom_llm_provider: provider.as_str().to_string(), + optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?, + secret_fields: Vec::new(), + api_key, + }; + let authenticated = resolve_auth(auth, environment, &|key| std::env::var(key).ok()).await?; + let wire = hooks + .before_send( + WireRequest { + url, + headers: authenticated.headers, + body: serde_json::to_value(&body).map_err(serialize_failure)?, + }, + context, + ) + .await?; + let provider_name = provider.as_str(); + debug!(provider = provider_name, stream, body = %wire.body, "provider request"); + let response = send( + http, + Authenticated { + headers: wire.headers, + signer: authenticated.signer, + }, + &wire.url, + &wire.body, + timeout, + ) + .await?; + debug!( + provider = provider_name, + status = response.status().as_u16(), + "provider response headers" + ); + if !response.status().is_success() { + return Err(provider_error(response).await); + } + let config = provider.config(); + if stream { + return Ok(streaming_response( + response, + config.stream_decoder(), + provider_name, + )); + } + let text = response.text().await.map_err(network)?; + debug!(body = text.as_str(), "provider response body"); + hooks + .emit(MachineEvent::ResponseReceived { + raw: RawResponse { body: text.clone() }, + }) + .await?; + decode_response(config, &body.model, &text) + .map(|message| MessagesResponse::Message(Box::new(message))) +} + +fn serialize_failure(err: serde_json::Error) -> Error { + Error::InvalidRequest(format!( + "failed to serialize Anthropic messages request: {err}" + )) +} + +fn network(error: reqwest::Error) -> Error { Error::Transport(TransportError::Network(error.to_string())) } -pub(super) async fn send( +async fn send( http: &litellm_http::Client, authenticated: Authenticated, url: &str, @@ -30,18 +123,21 @@ pub(super) async fn send( request.send(http).await.map_err(network) } -pub(super) async fn provider_error(response: reqwest::Response) -> Error { +async fn provider_error(response: reqwest::Response) -> Error { let status = response.status().as_u16(); match response.text().await { - Ok(text) => Error::Transport(TransportError::Http { - status, - body: truncate_error_body(&text), - }), + Ok(text) => { + litellm_tracing::debug!(status, body = text.as_str(), "provider error body"); + Error::Transport(TransportError::Http { + status, + body: truncate_error_body(&text), + }) + } Err(error) => network(error), } } -pub(super) fn decode_response( +fn decode_response( config: &dyn BaseAnthropicMessagesConfig, model: &str, text: &str, @@ -52,3 +148,97 @@ pub(super) fn decode_response( .transform_anthropic_messages_response(model, response) .map_err(Error::from) } + +fn streaming_response( + response: reqwest::Response, + decoder: Option, + provider: &'static str, +) -> MessagesResponse { + let headers = response + .headers() + .iter() + .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) + .collect(); + let chunks = match decoder { + None => futures_util::stream::try_unfold(response, move |mut response| async move { + let chunk = response.chunk().await.map_err(network)?; + Ok(chunk.map(|chunk| { + log_chunk(provider, "provider_response", &chunk); + (chunk, response) + })) + }) + .boxed(), + Some(decode) => decoded_chunks(response, decode, provider), + }; + MessagesResponse::Stream { headers, chunks } +} + +fn decoded_chunks( + response: reqwest::Response, + decode: StreamDecoder, + provider: &'static str, +) -> BoxStream<'static, Result> { + let bytes: ByteStream = response + .bytes_stream() + .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)) + .map_err(std::io::Error::other) + .boxed(); + futures_util::stream::try_unfold(decode(bytes), move |mut events| async move { + let Some(event) = events.try_next().await? else { + return Ok(None); + }; + let chunk = encode_anthropic_sse(&event)?; + log_chunk(provider, "client_response", &chunk); + Ok(Some((chunk, events))) + }) + .boxed() +} + +fn log_chunk(provider: &str, stage: &str, data: &Bytes) { + let chunk = ByteChunk::new(data); + debug!(provider, stage, encoding = chunk.encoding(), chunk = %chunk, "stream chunk"); +} + +#[cfg(test)] +mod tests { + use litellm_llms::base_llm::anthropic_messages::streaming::anthropic_sse_event_stream; + use rstest::rstest; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; + + use super::*; + + #[rstest] + #[case::event( + "data: {\"type\":\"ping\"}\n\n", + Some("event: ping\ndata: {\"type\":\"ping\"}\n\n") + )] + #[case::invalid_event("data: invalid\n\ndata: {\"type\":\"ping\"}\n\n", None)] + #[tokio::test] + async fn decoded_streams_encode_events_and_stop_at_the_first_error( + #[case] body: &'static str, + #[case] expected: Option<&str>, + ) { + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(200).set_body_raw(body, "text/event-stream")) + .mount(&upstream) + .await; + let response = litellm_http::Client::plain_for_test() + .get(upstream.uri()) + .send() + .await + .unwrap(); + let MessagesResponse::Stream { mut chunks, .. } = + streaming_response(response, Some(anthropic_sse_event_stream), "test") + else { + panic!("a streaming response returns chunks"); + }; + + let chunk = chunks.next().await.unwrap(); + match expected { + Some(expected) => assert_eq!(chunk.unwrap().as_ref(), expected.as_bytes()), + None => assert!(matches!(chunk, Err(Error::InvalidResponse(_))), "{chunk:?}"), + } + assert!(chunks.next().await.is_none()); + } +} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 3c081ff7bbd..f3a57da4d32 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -1,39 +1,27 @@ -//! The Anthropic Messages call, the Rust equivalent of Python's -//! `litellm.messages()`. +//! The Anthropic Messages call, the Rust equivalent of Python's `litellm.messages()`. //! -//! [`route`] is the call as a machine a host drives, streaming or not. [`messages`] runs -//! it in process for a caller that already holds the request and wants the message. +//! [`messages`] prepares the provider request and sends it in process. [`route`] runs the +//! same two steps as a machine for a host that answers the call's operations itself. -pub mod types; -pub use crate::error::RouteError as Error; mod common_utils; mod handler; mod prepare; pub mod route; -use std::sync::Arc; +mod types; use litellm_http::{ClientVariant, HttpClientConfig}; -use litellm_secrets::source::EnvironmentSecrets; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; -use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}; +use litellm_secrets::source::SecretSource; + +pub use crate::error::RouteError as Error; +pub use types::{MessagesCall, MessagesResponse, MessagesShaping, messages_body}; pub async fn messages( resources: &crate::resources::CoreResources, config: &HttpClientConfig, + secrets: &dyn SecretSource, call: MessagesCall, -) -> Result { - let secrets = Arc::new(EnvironmentSecrets::python_compatible( - resources.pool.client(config, ClientVariant::Provider)?, - )); - match litellm_host::run::run( - messages_machine(resources, config, secrets)?, - &LocalMessagesHost::new(call), - ) - .await? - { - MessagesOutput::Message(message) => Ok(*message), - MessagesOutput::Streamed => Err(Error::Unsupported( - "streamed responses need a streaming host", - )), - } +) -> Result { + let http = resources.pool.client(config, ClientVariant::Provider)?; + let request = prepare::prepare(call, secrets).await?; + handler::execute(&http, &resources.auth, request, &()).await } diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 7cd01a3a84c..c8e90cb5f2c 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,3 +1,6 @@ +use std::time::Duration; + +use litellm_auth::SecretValue; use litellm_core_utils::{ dot_notation_indexing::delete_nested_value, get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}, @@ -11,21 +14,42 @@ use litellm_llms::{ auth::{ValidatedEnvironment, with_default_headers}, }, }; +use litellm_secrets::source::SecretSource; use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use super::{ - Error, + Error, MessagesCall, common_utils::{MessagesProvider, string_headers}, - route::MessagesCall, - types::ProviderMessagesRequest, + types::invalid_request, }; -pub(super) struct ResolvedProvider { - pub(super) model: String, - pub(super) provider: MessagesProvider, +struct ResolvedProvider { + model: String, + provider: MessagesProvider, } -pub(super) fn resolve_provider( +pub(super) struct ProviderMessagesRequest { + pub(super) provider: MessagesProvider, + pub(super) url: String, + pub(super) body: AnthropicMessagesRequest, + pub(super) environment: ValidatedEnvironment, + pub(super) timeout: Option, + /// The caller's own credential, reported to the host beside the wire request. + pub(super) api_key: Option, +} + +pub(super) async fn prepare( + call: MessagesCall, + secrets: &dyn SecretSource, +) -> Result { + let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?; + let secrets = secrets + .resolve(resolved.provider.config().secret_names()) + .await?; + prepare_provider_request(call, resolved, secrets.as_ref()) +} + +fn resolve_provider( model: &str, custom_llm_provider: Option<&str>, ) -> Result { @@ -53,7 +77,7 @@ pub(super) fn resolve_provider( }) } -pub(super) fn prepare_provider_request( +fn prepare_provider_request( call: MessagesCall, resolved: ResolvedProvider, secrets: &dyn Lookup, @@ -113,13 +137,10 @@ pub(super) fn prepare_provider_request( body: transformed, environment, timeout, + api_key: api_key.map(SecretValue::new), }) } -pub(super) fn invalid_request(err: serde_json::Error) -> Error { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) -} - fn without_additional_drop_params( request: AnthropicMessagesRequest, paths: &[String], @@ -145,7 +166,7 @@ mod tests { use serde_json::{Map, Value, json}; use super::*; - use crate::messages::types::MessagesShaping; + use crate::messages::MessagesShaping; #[fixture] fn shaping() -> MessagesShaping { diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index f9267dce755..5e5bbc927b6 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -1,55 +1,20 @@ use std::{ convert::Infallible, sync::{Arc, Mutex}, - time::Duration, }; use bytes::Bytes; -use futures_util::StreamExt; -use litellm_auth::SecretValue; +use futures_util::TryStreamExt; use litellm_host::{ - event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, machine::{CallMachine, HostChannel, MachineFault}, protocol::Protocol, }; use litellm_http::{Client, ClientVariant, HttpClientConfig}; -use litellm_llms::base_llm::{ - anthropic_messages::streaming::{ByteStream, StreamDecoder, encode_anthropic_sse}, - auth::{Authenticated, resolve_auth}, -}; use litellm_secrets::source::SecretSource; -use litellm_types::{ - llms::anthropic_messages::{ - anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse, - }, - utils::ProviderSpecificHeaders, -}; -use serde_json::{Map, Value}; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; -use super::{ - Error, - handler::{decode_response, network, provider_error, send}, - prepare::{invalid_request, prepare_provider_request, resolve_provider}, - types::MessagesShaping, -}; - -/// The caller's request as the host projects it. -pub struct MessagesCall { - pub body: AnthropicMessagesRequest, - pub api_key: Option, - pub api_base: Option, - pub custom_llm_provider: Option, - pub extra_headers: Option>, - pub provider_specific_header: Option, - pub timeout: Option, - pub shaping: MessagesShaping, -} - -/// Parses a caller's raw body, failing the way the route fails for any invalid request. -pub fn messages_body(body: Map) -> Result { - serde_json::from_value(Value::Object(body)).map_err(invalid_request) -} +use super::{Error, MessagesCall, MessagesResponse, handler::execute, prepare::prepare}; pub enum MessagesOutput { Message(Box), @@ -121,134 +86,35 @@ pub fn messages_machine( let http = resources.pool.client(config, ClientVariant::Provider)?; let auth = resources.auth.clone(); Ok(CallMachine::new(move |host| { - Box::pin(execute(host, http.clone(), auth.clone(), secrets.clone())) + Box::pin(drive(host, http, auth, secrets)) })) } -async fn execute( +/// The call as its host sees it: projection first, then the same prepare and execute as +/// [`super::messages`], with each chunk of a stream handed over as it arrives. +async fn drive( host: MessagesHost, http: Client, auth: Arc, secrets: Arc, ) -> Result { let call = host.project().await?; - let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?; - let secrets = secrets - .resolve(resolved.provider.config().secret_names()) - .await?; - let api_key = call.api_key.clone().map(SecretValue::new); - let request = prepare_provider_request(call, resolved, secrets.as_ref())?; - let context = RequestContext { - model: request.body.model.clone(), - custom_llm_provider: request.provider.as_str().to_string(), - optional_params: serde_json::to_value(&request.body.params).map_err(serialize_failure)?, - secret_fields: Vec::new(), - api_key, - }; - let stream = request.body.params.stream == Some(true); - let config = request.provider.config(); - let body = serde_json::to_value(&request.body).map_err(serialize_failure)?; - let env_lookup = |key: &str| std::env::var(key).ok(); - let authenticated = resolve_auth(&auth, request.environment, &env_lookup).await?; - let wire = host - .before_send( - WireRequest { - url: request.url, - headers: authenticated.headers, - body, - }, - context, - ) - .await?; - let response = send( - &http, - Authenticated { - headers: wire.headers, - signer: authenticated.signer, - }, - &wire.url, - &wire.body, - request.timeout, - ) - .await?; - if !response.status().is_success() { - return Err(provider_error(response).await); - } - if stream { - return relay(&host, response, config.stream_decoder()).await; - } - let text = response.text().await.map_err(network)?; - host.emit(MachineEvent::ResponseReceived { - raw: RawResponse { body: text.clone() }, - }) - .await?; - decode_response(config, &request.body.model, &text) - .map(|message| MessagesOutput::Message(Box::new(message))) -} - -fn serialize_failure(err: serde_json::Error) -> Error { - Error::InvalidRequest(format!( - "failed to serialize Anthropic messages request: {err}" - )) -} - -/// Hands each upstream chunk to the caller as it arrives. A caller that stops reading -/// ends the upstream read, and the call completes with what it delivered. -/// -/// A host on Anthropic SSE is relayed byte for byte. A host on another wire is decoded into -/// Anthropic stream events and re-encoded as Anthropic SSE. -async fn relay( - host: &MessagesHost, - response: reqwest::Response, - decoder: Option, -) -> Result { - let head = MessagesStreamHead { - headers: response - .headers() - .iter() - .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) - .collect(), - }; - if host.open(head).await? == Demand::Detached { - return Ok(MessagesOutput::Streamed); - } - match decoder { - None => relay_bytes(host, response).await, - Some(decode) => relay_events(host, response, decode).await, - } -} - -async fn relay_bytes( - host: &MessagesHost, - mut response: reqwest::Response, -) -> Result { - while let Some(chunk) = response.chunk().await.map_err(network)? { - if host.deliver(chunk).await? == Demand::Detached { - break; + let request = prepare(call, secrets.as_ref()).await?; + match execute(&http, &auth, request, &host).await? { + MessagesResponse::Message(message) => Ok(MessagesOutput::Message(message)), + MessagesResponse::Stream { + headers, + mut chunks, + } => { + if host.open(MessagesStreamHead { headers }).await? == Demand::Detached { + return Ok(MessagesOutput::Streamed); + } + while let Some(chunk) = chunks.try_next().await? { + if host.deliver(chunk).await? == Demand::Detached { + break; + } + } + Ok(MessagesOutput::Streamed) } } - Ok(MessagesOutput::Streamed) -} - -async fn relay_events( - host: &MessagesHost, - response: reqwest::Response, - decode: StreamDecoder, -) -> Result { - let bytes: ByteStream = futures_util::stream::unfold(response, |mut response| async move { - match response.chunk().await { - Ok(Some(chunk)) => Some((Ok(chunk), response)), - Ok(None) => None, - Err(error) => Some((Err(std::io::Error::other(error)), response)), - } - }) - .boxed(); - let mut events = decode(bytes); - while let Some(event) = events.next().await { - let chunk = encode_anthropic_sse(&event?)?; - if host.deliver(chunk).await? == Demand::Detached { - break; - } - } - Ok(MessagesOutput::Streamed) } diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index 006b1db4efb..b09cb96a919 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -1,12 +1,45 @@ use std::time::Duration; -use litellm_llms::{ - anthropic::common_utils::AnthropicModelCapabilities, base_llm::auth::ValidatedEnvironment, +use bytes::Bytes; +use futures_util::stream::BoxStream; +use litellm_llms::anthropic::common_utils::AnthropicModelCapabilities; +use litellm_types::{ + llms::anthropic_messages::{ + anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse, + }, + utils::ProviderSpecificHeaders, }; -use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; -use super::common_utils::MessagesProvider; +use super::Error; + +pub struct MessagesCall { + pub body: AnthropicMessagesRequest, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, + pub provider_specific_header: Option, + pub timeout: Option, + pub shaping: MessagesShaping, +} + +pub fn messages_body(body: Map) -> Result { + serde_json::from_value(Value::Object(body)).map_err(invalid_request) +} + +pub(super) fn invalid_request(err: serde_json::Error) -> Error { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) +} + +pub enum MessagesResponse { + Message(Box), + Stream { + headers: Vec<(String, String)>, + chunks: BoxStream<'static, Result>, + }, +} #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct MessagesShaping { @@ -20,16 +53,6 @@ pub struct MessagesShaping { pub additional_drop_params: Vec, } -pub(crate) struct ProviderMessagesRequest { - pub(crate) provider: MessagesProvider, - pub(crate) url: String, - pub(crate) body: AnthropicMessagesRequest, - /// The forwarded, default and feature headers plus how the call authenticates; the - /// credential itself is applied when the request is sent. - pub(crate) environment: ValidatedEnvironment, - pub(crate) timeout: Option, -} - #[cfg(test)] mod tests { use litellm_llms::anthropic::common_utils::SupportedEffortTiers; diff --git a/litellm-rust/crates/core/tests/messages/main.rs b/litellm-rust/crates/core/tests/messages/main.rs index 719c86990b0..534af6d7d06 100644 --- a/litellm-rust/crates/core/tests/messages/main.rs +++ b/litellm-rust/crates/core/tests/messages/main.rs @@ -1,9 +1,8 @@ use std::{sync::Arc, time::Duration}; use litellm_core::messages::{ - Error, - route::{LocalMessagesHost, MessagesCall, MessagesMachine, MessagesOutput, messages_machine}, - types::MessagesShaping, + Error, MessagesCall, MessagesShaping, + route::{LocalMessagesHost, MessagesMachine, MessagesOutput, messages_machine}, }; use litellm_http::{HttpSettings, Resolution}; use litellm_secrets::source::SecretSource; diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs index 0d44d26d416..f6ee0e6dfbf 100644 --- a/litellm-rust/crates/core/tests/messages/request.rs +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -1,7 +1,5 @@ -use litellm_llms::anthropic::common_utils::{ - ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_OAUTH_BETA_HEADER, AnthropicModelCapabilities, - SupportedEffortTiers, beta, -}; +use litellm_llms::anthropic::common_utils::{AnthropicModelCapabilities, SupportedEffortTiers}; +use litellm_types::llms::anthropic::{AnthropicBeta, BetaSet}; use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; use rstest::rstest; @@ -247,41 +245,37 @@ async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) assert_eq!(sent["top_k"], 3); } -fn sent_betas(request: &wiremock::Request) -> Vec { +fn sent_betas(request: &wiremock::Request) -> BetaSet { let [header] = <[&str; 1]>::try_from(request.header_values("anthropic-beta")) .unwrap_or_else(|values| panic!("expected one anthropic-beta header, got {values:?}")); - header - .split(',') - .map(str::trim) - .map(str::to_string) - .collect() + header.parse().unwrap() } #[rstest] -#[case::structured_output(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])] -#[case::fast_mode(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])] -#[case::compaction(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])] +#[case::structured_output(json!({"output_format": {"type": "json_schema"}}), &[AnthropicBeta::StructuredOutputs20251113])] +#[case::fast_mode(json!({"speed": "fast"}), &[AnthropicBeta::FastMode20260201])] +#[case::compaction(json!({"compaction": {"enabled": true}}), &[AnthropicBeta::Compact20260904])] #[case::context_management_edits( json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}}), - &[beta::CONTEXT_MANAGEMENT_2025_06_27] + &[AnthropicBeta::ContextManagement20250627] )] #[case::per_message_output_config( json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), - &[beta::PER_TURN_CONTROL_2026_07_01] + &[AnthropicBeta::PerTurnControl20260701] )] #[case::advisor_tool( - json!({"tools": [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": MODEL}]}), - &[beta::ADVISOR_TOOL_2026_03_01] + json!({"tools": [{"type": "advisor_20260301", "name": "advisor", "model": MODEL}]}), + &[AnthropicBeta::AdvisorTool20260301] )] #[case::several_features_at_once( json!({"speed": "fast", "output_format": {"type": "json_schema"}}), - &[beta::STRUCTURED_OUTPUT, beta::FAST_MODE_2026_02_01] + &[AnthropicBeta::StructuredOutputs20251113, AnthropicBeta::FastMode20260201] )] #[tokio::test] async fn feature_betas_join_the_callers_betas_in_one_sorted_header( call: MessagesCall, #[case] fields: Value, - #[case] features: &[&str], + #[case] features: &[AnthropicBeta], ) { let upstream = upstream([message_response()]).await; let capabilities = AnthropicModelCapabilities { @@ -305,12 +299,11 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header( .await; let sent = sent_betas(&only_request(&upstream).await); - let mut expected: Vec = features + let expected: BetaSet = features .iter() - .map(|feature| feature.to_string()) - .chain(["caller-beta-2025-01-01".to_string()]) + .cloned() + .chain([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())]) .collect(); - expected.sort(); assert_eq!(sent, expected); } @@ -331,7 +324,10 @@ async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: M request.header("anthropic-dangerous-direct-browser-access"), Some("true") ); - assert_eq!(sent_betas(&request), [ANTHROPIC_OAUTH_BETA_HEADER]); + assert_eq!( + sent_betas(&request), + BetaSet::from_iter([AnthropicBeta::Oauth20250420]) + ); assert_eq!(request.header("x-api-key"), None); } diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs index ed715e22898..14c8eb6b7c6 100644 --- a/litellm-rust/crates/core/tests/messages/response.rs +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -1,6 +1,6 @@ use litellm_core::{ Phase, - messages::{messages, route::messages_body}, + messages::{MessagesResponse, messages, messages_body}, }; use litellm_http::transport::Error as TransportError; use rstest::rstest; @@ -188,9 +188,10 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes ..HttpSettings::default() }; - let message = messages( + let response = messages( &support::resources(), &Resolution::from(&settings).config, + &RecordingSecrets::empty(), MessagesCall { api_key: Some("sk-ant".into()), api_base: Some(base), @@ -200,6 +201,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes .await .expect("messages request succeeds"); + let MessagesResponse::Message(message) = response else { + panic!("a non-streaming request returns a message"); + }; assert_eq!(message.id, "msg_1"); let sent = only_request(&upstream).await; assert_eq!(sent.header("x-api-key"), Some("sk-ant")); diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs index 49a6f7e87a0..f0e55eca8dd 100644 --- a/litellm-rust/crates/core/tests/messages/stream.rs +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -1,12 +1,21 @@ -use std::{convert::Infallible, sync::Mutex}; +use std::{ + convert::Infallible, + sync::{Mutex, mpsc}, +}; use bytes::Bytes; -use litellm_core::messages::route::{Messages, MessagesStreamHead}; +use futures_util::{StreamExt, TryStreamExt}; +use litellm_core::messages::{ + MessagesResponse, messages, + route::{Messages, MessagesStreamHead}, +}; use litellm_host::host::{Demand, Host}; +use litellm_tracing::{Logger, Metadata, Record, Sink}; use rstest::rstest; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, + task::JoinHandle, }; use super::*; @@ -23,6 +32,20 @@ enum Seen { Deliver(Bytes), } +struct TraceSink(mpsc::Sender<(String, Value)>); + +impl Sink for TraceSink { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + metadata.target().starts_with("litellm_core::messages") + } + + fn emit(&self, record: &Record) { + self.0 + .send((record.message.clone(), Value::Object(record.fields.clone()))) + .unwrap(); + } +} + /// Projects like `LocalMessagesHost`, records every stream op in the order the route /// performs it, and detaches after `detach_after` ops. struct RecordingStreamHost { @@ -120,6 +143,37 @@ async fn upstream_headers_are_on_the_stream_head_before_the_first_chunk(call: Me assert_eq!(delivered, SSE_BODY.as_bytes()); } +#[rstest] +#[tokio::test] +async fn debug_trace_keeps_provider_input_and_every_stream_chunk(call: MessagesCall) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + let (sender, receiver) = mpsc::channel(); + + Logger::new(TraceSink(sender)) + .instrument(stream_through(&host)) + .await + .unwrap(); + + let records: Vec<(String, Value)> = receiver.try_iter().collect(); + let request = records + .iter() + .find(|(message, _)| message == "provider request") + .unwrap(); + let body: Value = serde_json::from_str(request.1["body"].as_str().unwrap()).unwrap(); + assert_eq!(body["messages"][0]["content"], "hi"); + assert_eq!(request.1["stream"], true); + let chunks: String = records + .iter() + .filter(|(message, fields)| { + message == "stream chunk" && fields["stage"] == "provider_response" + }) + .map(|(_, fields)| fields["chunk"].as_str().unwrap()) + .collect(); + assert_eq!(chunks, SSE_BODY); + assert!(!format!("{records:?}").contains("sk-ant")); +} + #[rstest] #[case::at_open(1)] #[case::after_the_first_chunk(2)] @@ -192,10 +246,10 @@ async fn a_stream_that_ends_without_message_stop_is_relayed_as_is(call: Messages } /// Serves one SSE chunk and then holds the connection open without ever finishing. -async fn stalling_upstream() -> String { +async fn stalling_upstream() -> (String, JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base = format!("http://{}", listener.local_addr().unwrap()); - tokio::spawn(async move { + let connection = tokio::spawn(async move { let (mut socket, _) = listener.accept().await.unwrap(); let mut request = vec![0; 4096]; let _ = socket.read(&mut request).await; @@ -206,15 +260,15 @@ async fn stalling_upstream() -> String { ) .await .unwrap(); - std::future::pending::<()>().await; + let _ = socket.read_to_end(&mut Vec::new()).await; }); - base + (base, connection) } #[rstest] #[tokio::test] async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { - let base = stalling_upstream().await; + let (base, connection) = stalling_upstream().await; let host = RecordingStreamHost::new( MessagesCall { timeout: Some(Duration::from_millis(300)), @@ -236,6 +290,145 @@ async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { "the chunk before the stall reached the caller, saw {} ops", seen.len() ); + tokio::time::timeout(Duration::from_secs(5), connection) + .await + .expect("timing out closes the upstream connection") + .unwrap(); +} + +#[rstest] +#[case::anthropic("anthropic")] +#[case::azure_ai("azure_ai")] +#[tokio::test] +async fn the_sdk_returns_stream_headers_and_every_sse_byte( + call: MessagesCall, + #[case] provider: &str, +) { + let upstream = upstream([sse_response()]).await; + let response = messages( + &support::resources(), + &http_config(), + &RecordingSecrets::empty(), + MessagesCall { + custom_llm_provider: Some(provider.into()), + ..streaming(call, upstream.uri()) + }, + ) + .await + .unwrap(); + + let MessagesResponse::Stream { headers, chunks } = response else { + panic!("a streaming request returns a stream"); + }; + for (name, value) in UPSTREAM_HEADERS { + assert!(headers.contains(&(name.into(), value.into()))); + } + let delivered = chunks.try_collect::>().await.unwrap().concat(); + assert_eq!(delivered, SSE_BODY.as_bytes()); + assert_eq!(only_request(&upstream).await.json()["stream"], true); +} + +#[rstest] +#[tokio::test] +async fn the_sdk_returns_http_errors_before_opening_a_stream(call: MessagesCall) { + let upstream = upstream([ResponseTemplate::new(429).set_body_string("slow down")]).await; + let error = messages( + &support::resources(), + &http_config(), + &RecordingSecrets::empty(), + streaming(call, upstream.uri()), + ) + .await + .err() + .expect("upstream failure is returned by messages()"); + + assert_eq!( + error, + Error::Transport(litellm_http::transport::Error::Http { + status: 429, + body: "slow down".into(), + }) + ); +} + +#[rstest] +#[case::before_reading(false)] +#[case::after_reading(true)] +#[tokio::test] +async fn dropping_the_sdk_stream_closes_the_unfinished_upstream( + call: MessagesCall, + #[case] read_chunk: bool, +) { + let (base, connection) = stalling_upstream().await; + let response = tokio::time::timeout( + Duration::from_secs(5), + messages( + &support::resources(), + &http_config(), + &RecordingSecrets::empty(), + MessagesCall { + timeout: Some(Duration::from_secs(30)), + ..streaming(call, base) + }, + ), + ) + .await + .expect("messages() returns before the upstream finishes") + .unwrap(); + + let MessagesResponse::Stream { mut chunks, .. } = response else { + panic!("a streaming request returns a stream"); + }; + if read_chunk { + let chunk = tokio::time::timeout(Duration::from_secs(5), chunks.next()) + .await + .expect("the first chunk arrives before the upstream finishes") + .unwrap() + .unwrap(); + assert_eq!(chunk.as_ref(), b"event: message_start\ndata: {}\n\n"); + } + assert!(!connection.is_finished()); + drop(chunks); + tokio::time::timeout(Duration::from_secs(5), connection) + .await + .expect("dropping the stream closes the upstream connection") + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesCall) { + let (base, connection) = stalling_upstream().await; + let response = messages( + &support::resources(), + &http_config(), + &RecordingSecrets::empty(), + MessagesCall { + timeout: Some(Duration::from_millis(300)), + ..streaming(call, base) + }, + ) + .await + .unwrap(); + + let MessagesResponse::Stream { mut chunks, .. } = response else { + panic!("a streaming request returns a stream"); + }; + assert_eq!( + chunks.next().await.unwrap().unwrap().as_ref(), + b"event: message_start\ndata: {}\n\n" + ); + let error = tokio::time::timeout(Duration::from_secs(5), chunks.next()) + .await + .expect("the stalled body times out") + .unwrap() + .unwrap_err(); + assert!(matches!(error, Error::Transport(_)), "{error:?}"); + assert!(chunks.next().await.is_none()); + tokio::time::timeout(Duration::from_secs(5), connection) + .await + .expect("the failed stream closes its upstream connection") + .unwrap(); } #[rstest] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index e3f40ec5354..f5ee5e81ba6 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -12,7 +12,6 @@ bytes.workspace = true futures-util.workspace = true litellm-auth.workspace = true litellm-core.workspace = true -litellm-host.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-router.workspace = true @@ -20,10 +19,10 @@ litellm-secrets.workspace = true litellm-types.workspace = true serde_json.workspace = true thiserror.workspace = true -tokio = { workspace = true, features = ["sync"] } [dev-dependencies] futures-util.workspace = true +tokio = { workspace = true, features = ["io-util"] } rstest.workspace = true tower = { version = "0.5.3", features = ["util"] } wiremock = "0.6.5" diff --git a/litellm-rust/crates/gateway-inference/src/messages/host.rs b/litellm-rust/crates/gateway-inference/src/messages/host.rs deleted file mode 100644 index 57486d9d157..00000000000 --- a/litellm-rust/crates/gateway-inference/src/messages/host.rs +++ /dev/null @@ -1,69 +0,0 @@ -use std::{convert::Infallible, sync::Mutex}; - -use bytes::Bytes; -use litellm_core::messages::{ - Error, - route::{LocalMessagesHost, Messages, MessagesCall, MessagesStreamHead}, -}; -use litellm_host::host::{Demand, Host}; -use tokio::sync::{mpsc, oneshot}; - -/// Hands a streamed response to the HTTP body: the head once, then each chunk. A dropped -/// receiver means the client went away, which detaches the call. -pub(super) struct ChannelHost { - local: LocalMessagesHost, - head: Mutex>>, - pub(super) chunks: mpsc::Sender, -} - -impl ChannelHost { - pub(super) fn new( - call: MessagesCall, - head: oneshot::Sender, - chunks: mpsc::Sender, - ) -> Self { - Self { - local: LocalMessagesHost::new(call), - head: Mutex::new(Some(head)), - chunks, - } - } - - fn take_head(&self) -> Option> { - self.head - .lock() - .unwrap_or_else(|error| error.into_inner()) - .take() - } - - pub(super) fn opened(&self) -> bool { - self.head - .lock() - .unwrap_or_else(|error| error.into_inner()) - .is_none() - } -} - -impl Host for ChannelHost { - async fn project(&self) -> Result { - self.local.project().await - } - - async fn custom_op(&self, op: Infallible) -> Result<(), Error> { - match op {} - } - - async fn open(&self, head: MessagesStreamHead) -> Result { - Ok(match self.take_head().map(|sender| sender.send(head)) { - Some(Ok(())) => Demand::More, - Some(Err(_)) | None => Demand::Detached, - }) - } - - async fn deliver(&self, chunk: Bytes) -> Result { - Ok(match self.chunks.send(chunk).await { - Ok(()) => Demand::More, - Err(_) => Demand::Detached, - }) - } -} diff --git a/litellm-rust/crates/gateway-inference/src/messages/mod.rs b/litellm-rust/crates/gateway-inference/src/messages/mod.rs index e96cae8ba53..5d6a8faa0e8 100644 --- a/litellm-rust/crates/gateway-inference/src/messages/mod.rs +++ b/litellm-rust/crates/gateway-inference/src/messages/mod.rs @@ -1,7 +1,5 @@ //! `POST /v1/messages`, as the Python proxy's `anthropic_response` serves it. -mod host; - use std::{convert::Infallible, sync::Arc}; use axum::{ @@ -11,13 +9,12 @@ use axum::{ http::{HeaderMap, StatusCode, header}, response::{IntoResponse, Response}, }; -use host::ChannelHost; -use litellm_core::messages::route::{ - MessagesCall, MessagesOutput, messages_body, messages_machine, +use futures_util::{StreamExt, stream::BoxStream}; +use litellm_core::messages::{ + Error as RouteError, MessagesCall, MessagesResponse, messages, messages_body, }; use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; use serde_json::{Map, Value}; -use tokio::sync::{mpsc, oneshot}; use crate::{Deployment, Error, Gateway}; @@ -55,31 +52,16 @@ async fn handle(gateway: &Gateway, headers: &HeaderMap, body: &[u8]) -> Result Ok(stream(chunks)), - joined = call => match joined.map_err(|error| Error::Internal(error.to_string()))?? { - MessagesOutput::Message(message) => Ok(Json(message).into_response()), - MessagesOutput::Streamed => Err(Error::Internal("the stream ended before it opened".into())), - }, + match messages( + &gateway.resources, + &gateway.http, + gateway.secrets.as_ref(), + call, + ) + .await? + { + MessagesResponse::Message(message) => Ok(Json(message).into_response()), + MessagesResponse::Stream { chunks, .. } => Ok(stream(chunks)), } } @@ -123,10 +105,13 @@ fn anthropic_api_headers(headers: &HeaderMap) -> Option }) } -fn stream(chunks: mpsc::Receiver) -> Response { - let body = futures_util::stream::unfold(chunks, |mut chunks| async move { - let chunk = chunks.recv().await?; - Some((Ok::<_, Infallible>(chunk), chunks)) +/// A chunk that fails after the stream opened is delivered as an SSE error frame, since +/// the status line already went out; the stream ends on it. +fn stream(chunks: BoxStream<'static, Result>) -> Response { + let body = chunks.map(|chunk| { + Ok::<_, Infallible>( + chunk.unwrap_or_else(|error| Bytes::from(Error::Route(error).sse_frame())), + ) }); ( StatusCode::OK, diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index ef7e7c66681..30836498d49 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -6,6 +6,7 @@ use axum::{ }; use rstest::rstest; use serde_json::json; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; use wiremock::{ Mock, MockServer, ResponseTemplate, @@ -64,3 +65,57 @@ async fn invalid_messages_stays_an_anthropic_error() { assert_eq!(body["type"], "error"); assert_eq!(body["error"]["type"], "invalid_request_error"); } + +/// Answers with the SSE head and one event, then drops the connection short of the +/// announced body length. +async fn truncating_upstream() -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{FIRST_EVENT}", + FIRST_EVENT.len() * 2 + ) + .as_bytes(), + ) + .await + .unwrap(); + }); + base +} + +const FIRST_EVENT: &str = "event: message_start\ndata: {}\n\n"; + +#[tokio::test] +async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame() { + let base = truncating_upstream().await; + let request = Request::post("/v1/messages") + .header("content-type", "application/json") + .body(Body::from( + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, "stream": true}) + .to_string(), + )) + .unwrap(); + + let response = support::app("anthropic/test-model", &base) + .oneshot(request) + .await + .unwrap(); + + assert_eq!(response.status(), 200); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + let text = std::str::from_utf8(&body).unwrap(); + let frame = text + .strip_prefix(FIRST_EVENT) + .and_then(|rest| rest.strip_prefix("event: error\ndata: ")) + .unwrap_or_else(|| panic!("the delivered event then one error frame, got {text:?}")); + let error: serde_json::Value = serde_json::from_str(frame.trim_end()).unwrap(); + assert_eq!(error["type"], "error"); + assert_eq!(error["error"]["type"], "api_error"); +} diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 554186955b4..c27a3f5b17e 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] axum.workspace = true +http-body-util = "0.1" litellm-core.workspace = true litellm-gateway-inference.workspace = true litellm-gateway-auth.workspace = true @@ -14,11 +15,14 @@ litellm-config.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets.workspace = true -tower-http = { version = "0.7.1", default-features = false, features = ["trace"] } +litellm-tracing.workspace = true +serde_json.workspace = true tracing.workspace = true tokio.workspace = true +uuid.workspace = true [dev-dependencies] +futures-util.workspace = true rstest.workspace = true -serde_json.workspace = true tokio = { workspace = true, features = ["sync"] } +tower = { version = "0.5", features = ["util"] } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 3f67f923a5c..fc16dc0de67 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,7 +1,13 @@ -use std::sync::Arc; +use std::{sync::Arc, time::Instant}; -use axum::{Router, extract::Request}; -use tower_http::trace::{DefaultOnResponse, TraceLayer}; +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + middleware::Next, + response::Response, +}; +use http_body_util::BodyExt; use litellm_config::Config; use litellm_core::resources::CoreResources; @@ -12,6 +18,8 @@ use litellm_http::{ }; use litellm_llms::base_llm::ocr::settings::OcrSettings; use litellm_secrets::source::EnvironmentSecrets; +use litellm_tracing::ByteChunk; +use uuid::Uuid; pub fn build_inference(config: &Config) -> Result, litellm_http::Error> { let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); @@ -42,11 +50,125 @@ pub fn router(inference: Arc, config: &Config) -> Router { RequireMasterKey, _, >(auth)) - .layer( - TraceLayer::new_for_http() - .make_span_with(|request: &Request| { - tracing::info_span!("request", method = %request.method(), path = request.uri().path()) - }) - .on_response(DefaultOnResponse::new().level(tracing::Level::INFO)), - ) + .layer(axum::middleware::from_fn(log_request)) +} + +async fn log_request(request: Request, next: Next) -> Response { + let request_id = Uuid::new_v4().to_string(); + let log_body_chunks = tracing::enabled!(tracing::Level::DEBUG); + let method = request.method().clone(); + let path = request.uri().path().to_owned(); + let started = Instant::now(); + let request = if log_body_chunks { + request.map(|body| logged_body(body, request_id.clone(), "input")) + } else { + request + }; + let response = next.run(request).await; + tracing::info!( + %request_id, + %method, + %path, + status = response.status().as_u16(), + time_to_headers_ms = started.elapsed().as_secs_f64() * 1000.0, + "response headers" + ); + if log_body_chunks { + response.map(|body| logged_body(body, request_id, "output")) + } else { + response + } +} + +fn logged_body(body: Body, request_id: String, direction: &'static str) -> Body { + Body::new(body.map_frame(move |frame| { + if let Some(data) = frame.data_ref() { + log_chunk(&request_id, direction, data); + } + frame + })) +} + +fn log_chunk(request_id: &str, direction: &str, data: &Bytes) { + let chunk = ByteChunk::new(data); + tracing::debug!(request_id, direction, encoding = chunk.encoding(), chunk = %chunk, "body chunk"); +} + +#[cfg(test)] +mod tests { + use std::{convert::Infallible, sync::mpsc}; + + use axum::{body::to_bytes, http::StatusCode, routing::post}; + use futures_util::stream; + use litellm_tracing::{Logger, Metadata, Record, Sink}; + use rstest::rstest; + use serde_json::{Value, json}; + use tower::ServiceExt; + + use super::*; + + struct LogSink(mpsc::Sender); + + impl Sink for LogSink { + fn enabled(&self, _: &Metadata<'_>) -> bool { + true + } + + fn emit(&self, record: &Record) { + self.0 + .send(json!({"message": record.message, "fields": record.fields})) + .unwrap(); + } + } + + #[rstest] + #[tokio::test] + async fn logs_each_body_chunk_without_changing_streamed_bytes() { + let app = Router::new() + .route( + "/stream", + post(|_: Bytes| async { + ( + StatusCode::OK, + Body::from_stream(stream::iter([ + Ok::<_, Infallible>(Bytes::from_static(b"event: first\n\n")), + Ok(Bytes::from_static(b"event: second\n\n")), + ])), + ) + }), + ) + .layer(axum::middleware::from_fn(log_request)); + let request_chunks = [ + Ok::<_, Infallible>(Bytes::from_static(b"hello")), + Ok(Bytes::from_static(b" world")), + ]; + let request = Request::post("/stream") + .body(Body::from_stream(stream::iter(request_chunks))) + .unwrap(); + let (sender, receiver) = mpsc::channel(); + let logger = Logger::new(LogSink(sender)); + + let output = logger + .instrument(async { + let response = app.oneshot(request).await.unwrap(); + to_bytes(response.into_body(), 1024).await.unwrap() + }) + .await; + + assert_eq!(output, "event: first\n\nevent: second\n\n"); + let records: Vec = receiver.try_iter().collect(); + assert_eq!(records.len(), 5); + assert_eq!(records[0]["fields"]["chunk"], "hello"); + assert_eq!(records[1]["fields"]["chunk"], " world"); + assert_eq!(records[2]["fields"]["status"], 200); + assert_eq!(records[3]["fields"]["chunk"], "event: first\n\n"); + assert_eq!(records[4]["fields"]["chunk"], "event: second\n\n"); + let request_id = &records[2]["fields"]["request_id"]; + assert!(request_id.as_str().is_some()); + assert!( + records + .iter() + .all(|record| &record["fields"]["request_id"] == request_id) + ); + } } diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index bae711e16c7..40f9cb442d5 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -1,9 +1,46 @@ -use std::error::Error; +use std::{ + error::Error, + time::{SystemTime, UNIX_EPOCH}, +}; use litellm_config::Config; +use litellm_tracing::{Level, Logger, Metadata, Record, Sink}; +use serde_json::json; + +struct StderrSink { + level: Level, +} + +impl Sink for StderrSink { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + *metadata.level() <= self.level && metadata.target().starts_with("litellm") + } + + fn emit(&self, record: &Record) { + let timestamp_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + eprintln!( + "{}", + json!({ + "timestamp_ms": timestamp_ms, + "level": record.metadata.level().as_str(), + "target": record.metadata.target(), + "message": record.message, + "fields": record.fields, + }) + ); + } +} #[tokio::main] async fn main() -> Result<(), Box> { + let level = std::env::var("RUST_LOG") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(Level::INFO); + Logger::new(StderrSink { level }).install_global()?; let config_path = std::env::var("LITELLM_CONFIG").unwrap_or_else(|_| "config.yaml".into()); let config = Config::load(config_path)?; let inference = litellm_gateway::build_inference(&config)?; @@ -13,6 +50,8 @@ async fn main() -> Result<(), Box> { .parse::()?; let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; + tracing::info!(address = %listener.local_addr()?, models = config.model_list.len(), log_level = %level, "gateway listening"); + axum::serve(listener, litellm_gateway::router(inference, &config)).await?; Ok(()) } diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index a19d8c9c4fe..a91051441ea 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -1,11 +1,31 @@ -use std::{sync::Arc, time::Duration}; +use std::{ + sync::{Arc, mpsc}, + time::Duration, +}; +use axum::{body::Body, http::Request}; use litellm_config::Config; use litellm_gateway_inference::{Error, Gateway}; use litellm_http::ClientVariant; +use litellm_tracing::{Logger, Metadata, Record, Sink}; use rstest::{fixture, rstest}; use serde_json::{Value, json}; use tokio::{net::TcpListener, sync::oneshot, time::timeout}; +use tower::ServiceExt; + +struct LogSink(mpsc::Sender); + +impl Sink for LogSink { + fn enabled(&self, _: &Metadata<'_>) -> bool { + true + } + + fn emit(&self, record: &Record) { + self.0 + .send(json!({"message": record.message, "fields": record.fields})) + .unwrap(); + } +} #[fixture] fn inference() -> Arc { @@ -88,3 +108,36 @@ async fn authenticates_before_serving_mounted_inference_routes( .unwrap() .unwrap(); } + +#[rstest] +#[tokio::test] +async fn logs_request_outcome_without_credentials_or_query(inference: Arc) { + let config = + Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: gateway-key\n") + .unwrap(); + let request = Request::builder() + .method("POST") + .uri("/v1/messages?token=query-secret") + .header("authorization", "Bearer header-secret") + .body(Body::empty()) + .unwrap(); + let (sender, receiver) = mpsc::channel(); + let logger = Logger::new(LogSink(sender)); + + let response = logger + .instrument(litellm_gateway::router(inference, &config).oneshot(request)) + .await + .unwrap(); + + assert_eq!(response.status().as_u16(), 401); + let record = receiver.try_recv().unwrap(); + assert_eq!(record["message"], "response headers"); + assert_eq!(record["fields"]["method"], "POST"); + assert_eq!(record["fields"]["path"], "/v1/messages"); + assert_eq!(record["fields"]["status"], 401); + assert!(record["fields"]["time_to_headers_ms"].as_f64().unwrap() >= 0.0); + assert!(record["fields"]["request_id"].as_str().is_some()); + assert!(receiver.try_recv().is_err()); + assert!(!record.to_string().contains("header-secret")); + assert!(!record.to_string().contains("query-secret")); +} diff --git a/litellm-rust/crates/host/src/hooks.rs b/litellm-rust/crates/host/src/hooks.rs new file mode 100644 index 00000000000..14b0f1ea08a --- /dev/null +++ b/litellm-rust/crates/host/src/hooks.rs @@ -0,0 +1,145 @@ +use std::future::Future; + +use crate::{ + event::{MachineEvent, RequestContext, WireRequest}, + machine::{HostChannel, MachineFault}, + protocol::Protocol, +}; + +/// What a route reaches for mid-call: the send-time rewrite and the events it reports. +/// Python's `logging_obj.pre_call` and `post_call`, in that order. +pub trait RouteHooks: Send + Sync { + fn before_send( + &self, + wire: WireRequest, + context: RequestContext, + ) -> impl Future> + Send; + + fn emit(&self, event: MachineEvent) -> impl Future> + Send; +} + +/// No host: the wire request goes out as prepared and nothing observes the call. +impl RouteHooks for () { + async fn before_send(&self, wire: WireRequest, _: RequestContext) -> Result { + Ok(wire) + } + + async fn emit(&self, _: MachineEvent) -> Result<(), E> { + Ok(()) + } +} + +impl RouteHooks for HostChannel +where + R::Error: From, +{ + async fn before_send( + &self, + wire: WireRequest, + context: RequestContext, + ) -> Result { + HostChannel::before_send(self, wire, context).await + } + + async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> { + HostChannel::emit(self, event).await + } +} + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + + use serde_json::json; + + use super::*; + use crate::{ + event::RawResponse, + host::HostOp, + machine::{CallMachine, Machine, MachineStep}, + }; + + struct Unit; + + #[derive(Clone, Debug)] + struct Fault; + + impl Protocol for Unit { + type Response = (WireRequest, ()); + type Error = Fault; + type Projection = (); + type Op = Infallible; + type Chunk = Infallible; + type StreamHead = Infallible; + } + + impl From for Fault { + fn from(_: MachineFault) -> Self { + Fault + } + } + + fn wire(url: &str) -> WireRequest { + WireRequest { + url: url.into(), + headers: Vec::new(), + body: json!({}), + } + } + + fn context() -> RequestContext { + RequestContext { + model: "m".into(), + custom_llm_provider: "p".into(), + optional_params: json!({}), + secret_fields: Vec::new(), + api_key: None, + } + } + + #[tokio::test] + async fn the_channel_yields_each_hook_as_its_op_and_returns_the_answer() { + let mut machine = CallMachine::::new(|channel| { + Box::pin(async move { + let sent = RouteHooks::before_send(&channel, wire("prepared"), context()).await?; + RouteHooks::emit( + &channel, + MachineEvent::ResponseReceived { + raw: RawResponse { body: "raw".into() }, + }, + ) + .await?; + Ok((sent, ())) + }) + }); + + let Ok(MachineStep::Host(HostOp::BeforeSend { wire, reply, .. })) = machine.resume().await + else { + panic!("before_send yields BeforeSend"); + }; + assert_eq!(wire.url, "prepared"); + reply.send(WireRequest { + url: "rewritten".into(), + ..*wire + }); + + let Ok(MachineStep::Host(HostOp::Emit(event, reply))) = machine.resume().await else { + panic!("emit yields Emit"); + }; + assert!(matches!(event, MachineEvent::ResponseReceived { .. })); + reply.send(()); + + let Ok(MachineStep::Complete((sent, ()))) = machine.resume().await else { + panic!("the call completes with the answers"); + }; + assert_eq!(sent.url, "rewritten"); + } + + #[tokio::test] + async fn no_hooks_pass_the_wire_request_through() { + let sent = RouteHooks::::before_send(&(), wire("prepared"), context()) + .await + .unwrap(); + assert_eq!(sent.url, "prepared"); + } +} diff --git a/litellm-rust/crates/host/src/lib.rs b/litellm-rust/crates/host/src/lib.rs index c6b9e59b65a..1df68941fa3 100644 --- a/litellm-rust/crates/host/src/lib.rs +++ b/litellm-rust/crates/host/src/lib.rs @@ -7,6 +7,7 @@ //! may rewrite the wire request before it is sent. pub mod event; +pub mod hooks; pub mod host; pub mod machine; pub mod protocol; diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index fcf296793a5..17f2e652a92 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -87,6 +87,20 @@ pub fn has_header(headers: &[(String, String)], name: &str) -> bool { .any(|(key, _)| key.eq_ignore_ascii_case(name)) } +pub fn header_value<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { + headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) +} + +pub fn without_headers(headers: Vec<(String, String)>, names: &[&str]) -> Vec<(String, String)> { + headers + .into_iter() + .filter(|(key, _)| !names.iter().any(|name| key.eq_ignore_ascii_case(name))) + .collect() +} + pub fn has_bearer_auth(headers: &[(String, String)]) -> bool { headers.iter().any(|(name, value)| { if !name.eq_ignore_ascii_case("authorization") { @@ -194,6 +208,30 @@ mod tests { assert!(!has_header(&headers, "authorization")); } + #[test] + fn header_value_reads_the_first_match_in_any_case() { + let headers = vec![ + ("X-Api-Key".to_string(), "first".to_string()), + ("x-api-key".to_string(), "second".to_string()), + ]; + assert_eq!(header_value(&headers, "x-API-key"), Some("first")); + assert_eq!(header_value(&headers, "authorization"), None); + } + + #[test] + fn without_headers_drops_every_casing_of_the_named_headers_and_keeps_order() { + let headers = vec![ + ("X-Api-Key".to_string(), "k".to_string()), + ("anthropic-version".to_string(), "v".to_string()), + ("AUTHORIZATION".to_string(), "Bearer t".to_string()), + ("x-api-key".to_string(), "k2".to_string()), + ]; + assert_eq!( + without_headers(headers, &["x-api-key", "authorization"]), + vec![("anthropic-version".to_string(), "v".to_string())] + ); + } + #[test] fn auth_header_detection_is_case_insensitive() { let headers = vec![ diff --git a/litellm-rust/crates/litellm/Cargo.toml b/litellm-rust/crates/litellm/Cargo.toml index f6a63227792..41009ceaff1 100644 --- a/litellm-rust/crates/litellm/Cargo.toml +++ b/litellm-rust/crates/litellm/Cargo.toml @@ -1,3 +1,4 @@ [package] name = "litellm" version = "0.0.1" +edition.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs index 395c2376059..d1a8bdff55a 100644 --- a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs @@ -4,7 +4,7 @@ use serde_json::Value; use time::OffsetDateTime; use url::Url; -use crate::{Error, anthropic::messages::transformation::resolve_anthropic_api_base}; +use crate::{Error, anthropic::common_utils::resolve_anthropic_api_base}; const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches"; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index b19443a7ff8..ba77a6ed790 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,4 +1,4 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::SecretValue; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, prompt_templates::factory::{Conversation, build_conversation}, @@ -12,9 +12,11 @@ use serde_json::{Map, Value, json}; use crate::{ Error, anthropic::{ - ANTHROPIC_OAUTH_TOKEN_PREFIX, chat::handler::ModelResponseIterator, - messages::transformation::{complete_anthropic_url, resolve_anthropic_api_key}, + common_utils::{ + API_KEY_PLACEMENT, complete_anthropic_url, forwarded_oauth_bearer, + resolve_anthropic_api_key, + }, }, base_llm::{ anthropic_messages::streaming::anthropic_sse_event_stream, @@ -50,15 +52,6 @@ pub struct AnthropicConfig; pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicConfig = AnthropicConfig; -fn forwards_oauth_bearer(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)) - }) -} - impl BaseConfig for AnthropicConfig { fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] { SUPPORTED_PARAMS @@ -160,14 +153,14 @@ impl BaseConfig for AnthropicConfig { _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - if forwards_oauth_bearer(&headers) { + if forwarded_oauth_bearer(&headers).is_some() { return Ok(ValidatedEnvironment { headers, auth: AuthScheme::Forwarded, }); } let auth = AuthScheme::Credential { - placement: CredentialPlacement::Header("x-api-key"), + placement: API_KEY_PLACEMENT, secret: SecretValue::new(resolve_anthropic_api_key(api_key, env_lookup)?), }; Ok(ValidatedEnvironment { headers, auth }) diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 34c59c6ec5d..d8d24ec0402 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,30 +1,33 @@ -use litellm_types::llms::anthropic_messages::anthropic_request::{ - AnthropicMessage, ContentBlock, EffortLevel, MessageContent, +use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_http::request::{has_header, header_value, without_headers}; +use litellm_types::llms::{ + anthropic::{AnthropicBeta, BetaSet}, + anthropic_messages::anthropic_request::{ + AnthropicMessage, AnthropicTool, ContentBlock, EffortLevel, MessageContent, + }, }; +use litellm_types::recognized::Recognized; use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX; +use crate::{ + anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX, + base_llm::auth::{AuthScheme, Headers}, +}; -pub const ANTHROPIC_OAUTH_BETA_HEADER: &str = "oauth-2025-04-20"; -pub const ANTHROPIC_ADVISOR_TOOL_TYPE: &str = "advisor_20260301"; -pub const ANTHROPIC_TOOL_SEARCH_TOOL_TYPES: [&str; 2] = [ - "tool_search_tool_regex_20251119", - "tool_search_tool_bm25_20251119", -]; +pub const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; +pub const ANTHROPIC_AUTH_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN"; pub const ENCRYPTED_REASONING_SIGNATURE_PREFIX: &str = "litellm_encrypted_reasoning:"; const THOUGHT_SIGNATURE_SEPARATOR: &str = "__thought__"; - -pub mod beta { - pub const CONTEXT_MANAGEMENT_2025_06_27: &str = "context-management-2025-06-27"; - pub const COMPACT_2026_01_12: &str = "compact-2026-01-12"; - pub const COMPACT_2026_09_04: &str = "compact-2026-09-04"; - pub const STRUCTURED_OUTPUT: &str = "structured-outputs-2025-11-13"; - pub const ADVANCED_TOOL_USE_2025_11_20: &str = "advanced-tool-use-2025-11-20"; - pub const FAST_MODE_2026_02_01: &str = "fast-mode-2026-02-01"; - pub const ADVISOR_TOOL_2026_03_01: &str = "advisor-tool-2026-03-01"; - pub const PER_TURN_CONTROL_2026_07_01: &str = "per-turn-control-2026-07-01"; -} +const BETA_HEADER: &str = "anthropic-beta"; +pub const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; +pub const ANTHROPIC_BASE_URL_ENV: &str = "ANTHROPIC_BASE_URL"; +pub const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com"; +pub const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; +pub const API_KEY_PLACEMENT: CredentialPlacement = CredentialPlacement::Header("x-api-key"); +const API_KEY_HEADER: &str = API_KEY_PLACEMENT.header_name(); +const AUTHORIZATION: &str = CredentialPlacement::Bearer.header_name(); +const DIRECT_BROWSER_ACCESS_HEADER: &str = "anthropic-dangerous-direct-browser-access"; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct SupportedEffortTiers { @@ -111,42 +114,207 @@ impl AnthropicModelCapabilities { } } -pub fn is_anthropic_oauth_key(value: &str) -> bool { - value - .strip_prefix("Bearer ") - .unwrap_or(value) - .starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) +pub fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) } -pub fn split_beta_values(header: Option<&str>) -> impl Iterator + '_ { - header - .into_iter() - .flat_map(|value| value.split(',')) - .map(str::trim) - .filter(|piece| !piece.is_empty()) +pub fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option, name: &str) -> Option { + env_lookup(name).filter(|value| !value.trim().is_empty()) +} + +/// An Anthropic OAuth access token, which authenticates as a bearer instead of an `x-api-key`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OauthToken<'a>(&'a str); + +impl<'a> OauthToken<'a> { + /// The raw token, as a caller passes it in `api_key`. + pub fn parse(value: &'a str) -> Option { + value + .starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) + .then_some(Self(value)) + } + + /// A configured key, which users paste either raw or already prefixed with `Bearer `. + pub fn parse_key(value: &'a str) -> Option { + Self::parse(value.strip_prefix("Bearer ").unwrap_or(value)) + } + + pub fn as_str(self) -> &'a str { + self.0 + } + + pub fn into_auth(self) -> AuthScheme { + AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(self.0), + } + } +} + +/// Python's `AnthropicModelInfo.get_api_key`: the param, else `ANTHROPIC_API_KEY`. +pub fn get_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + non_empty(api_key) .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, ANTHROPIC_API_KEY_ENV)) } -pub fn join_beta_values(values: impl IntoIterator) -> String { - let mut values: Vec = values.into_iter().collect(); - values.sort(); - values.dedup(); - values.join(",") +pub fn get_auth_token(env_lookup: &dyn Fn(&str) -> Option) -> Option { + non_empty_env(env_lookup, ANTHROPIC_AUTH_TOKEN_ENV) } -pub fn is_tool_search_used(tools: Option<&[Value]>) -> bool { - tools.into_iter().flatten().any(|tool| { - tool.get("type") - .and_then(Value::as_str) - .is_some_and(|tool_type| ANTHROPIC_TOOL_SEARCH_TOOL_TYPES.contains(&tool_type)) +/// Python's `AnthropicModelInfo.get_auth_header`, naming the credential instead of building +/// the header: the key goes in `x-api-key` unless it is an OAuth token, and without a key +/// `ANTHROPIC_AUTH_TOKEN` is sent as a bearer. +pub fn get_auth_header( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + if let Some(key) = get_api_key(api_key, env_lookup) { + return Some(match OauthToken::parse_key(&key) { + Some(token) => token.into_auth(), + None => AuthScheme::Credential { + placement: API_KEY_PLACEMENT, + secret: SecretValue::new(key), + }, + }); + } + get_auth_token(env_lookup).map(|token| AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(token), }) } -pub fn has_advisor_tool(tools: Option<&[Value]>) -> bool { +pub fn resolve_anthropic_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + get_api_key(api_key, env_lookup).ok_or(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: ANTHROPIC_API_KEY_ENV, + }) +} + +/// Whether the caller already forwarded an Anthropic credential, in either header. +pub fn has_anthropic_credential(headers: &[(String, String)]) -> bool { + has_header(headers, API_KEY_HEADER) || has_header(headers, AUTHORIZATION) +} + +pub fn resolve_anthropic_api_base( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + non_empty(api_base) + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, ANTHROPIC_API_BASE_ENV)) + .or_else(|| non_empty_env(env_lookup, ANTHROPIC_BASE_URL_ENV)) + .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string()) +} + +pub fn complete_anthropic_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + let api_base = resolve_anthropic_api_base(api_base, env_lookup); + + let api_base = api_base.trim_end_matches('/'); + if api_base.ends_with(MESSAGES_PATH_SUFFIX) { + return api_base.to_string(); + } + format!("{api_base}{MESSAGES_PATH_SUFFIX}") +} + +pub fn existing_betas(headers: &[(String, String)]) -> BetaSet { + headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case(BETA_HEADER)) + .flat_map(|(_, value)| { + value + .parse::() + .unwrap_or_else(|never| match never {}) + }) + .collect() +} + +/// Python's `_merge_beta_headers`, over every casing of the header at once: the union of what +/// the caller sent and `added` replaces the header, sorted and deduplicated. Headers without +/// any beta value stay as they are. +pub fn merge_beta_headers(headers: Headers, added: BetaSet) -> Headers { + let merged = existing_betas(&headers).union(added); + if merged.is_empty() { + return headers; + } + without_headers(headers, &[BETA_HEADER]) + .into_iter() + .chain([(BETA_HEADER.to_string(), merged.to_string())]) + .collect() +} + +/// The outcome of Python's `optionally_handle_anthropic_oauth`. +#[derive(Clone, Debug, PartialEq)] +pub enum OauthHandling { + /// An OAuth token is the whole credential. The headers carry its companions and no + /// longer any `x-api-key` or `authorization`, so the bearer is applied on top. + Bearer { + headers: Headers, + token: SecretValue, + }, + Untouched(Headers), +} + +/// The OAuth token a caller forwarded as `Authorization: Bearer sk-ant-oat…`. +pub fn forwarded_oauth_bearer(headers: &[(String, String)]) -> Option> { + header_value(headers, AUTHORIZATION) + .and_then(|value| value.strip_prefix("Bearer ")) + .and_then(OauthToken::parse) +} + +fn with_oauth_companions(headers: Headers, dropped: &[&str]) -> Headers { + merge_beta_headers( + without_headers(headers, dropped), + BetaSet::from_iter([AnthropicBeta::Oauth20250420]), + ) + .into_iter() + .chain([(DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string())]) + .collect() +} + +pub fn optionally_handle_anthropic_oauth(headers: Headers, api_key: Option<&str>) -> OauthHandling { + if let Some(token) = + forwarded_oauth_bearer(&headers).map(|token| SecretValue::new(token.as_str())) + { + return OauthHandling::Bearer { + headers: with_oauth_companions(headers, &[API_KEY_HEADER, AUTHORIZATION]), + token, + }; + } + if let Some(token) = api_key.and_then(OauthToken::parse) { + return OauthHandling::Bearer { + headers: with_oauth_companions(headers, &[API_KEY_HEADER]), + token: SecretValue::new(token.as_str()), + }; + } + OauthHandling::Untouched(headers) +} + +pub fn is_tool_search_used(tools: Option<&[Recognized]>) -> bool { + tools.into_iter().flatten().any(|tool| { + matches!( + tool, + Recognized::Known( + AnthropicTool::ToolSearchRegex { .. } | AnthropicTool::ToolSearchBm25 { .. } + ) + ) + }) +} + +pub fn has_advisor_tool(tools: Option<&[Recognized]>) -> bool { tools .into_iter() .flatten() - .any(|tool| tool.get("type").and_then(Value::as_str) == Some(ANTHROPIC_ADVISOR_TOOL_TYPE)) + .any(|tool| matches!(tool, Recognized::Known(AnthropicTool::Advisor { .. }))) } pub fn requires_native_compaction_beta( @@ -521,8 +689,97 @@ mod tests { serde_json::from_value(messages).unwrap() } - fn tools(value: Option) -> Option> { - value.map(|tools| tools.as_array().unwrap().clone()) + fn tools(value: Option) -> Option>> { + value.map(|tools| serde_json::from_value(tools).unwrap()) + } + + fn headers(pairs: &[(&str, &str)]) -> Headers { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + } + + fn betas(values: &[&str]) -> BetaSet { + values.join(",").parse().unwrap() + } + + fn env(vars: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option { + move |name| { + vars.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + } + } + + const BOTH_BASE_ENVS: &[(&str, &str)] = &[ + (ANTHROPIC_API_BASE_ENV, "https://api-base.example.com"), + (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com"), + ]; + + #[rstest] + #[case::public_endpoint_by_default(None, &[], "https://api.anthropic.com")] + #[case::explicit_api_base_beats_env( + Some("https://explicit.example.com"), + BOTH_BASE_ENVS, + "https://explicit.example.com" + )] + #[case::explicit_api_base_is_trimmed( + Some(" https://explicit.example.com "), + &[], + "https://explicit.example.com" + )] + #[case::blank_api_base_falls_back_to_env( + Some(" "), + BOTH_BASE_ENVS, + "https://api-base.example.com" + )] + #[case::api_base_env_beats_base_url_env(None, BOTH_BASE_ENVS, "https://api-base.example.com")] + #[case::base_url_env_without_api_base_env( + None, + &[(ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], + "https://base-url.example.com" + )] + #[case::blank_api_base_env_falls_back_to_base_url_env( + None, + &[(ANTHROPIC_API_BASE_ENV, " \t "), (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], + "https://base-url.example.com" + )] + #[case::blank_envs_fall_back_to_public_endpoint( + None, + &[(ANTHROPIC_API_BASE_ENV, ""), (ANTHROPIC_BASE_URL_ENV, " ")], + "https://api.anthropic.com" + )] + fn api_base_resolution( + #[case] api_base: Option<&str>, + #[case] vars: &'static [(&'static str, &'static str)], + #[case] expected: &str, + ) { + assert_eq!(resolve_anthropic_api_base(api_base, &env(vars)), expected); + } + + #[rstest] + #[case::forwarded_api_key(&[("X-Api-Key", "k")], true)] + #[case::forwarded_bearer(&[("Authorization", "Bearer t")], true)] + #[case::nothing_forwarded(&[("anthropic-version", "2023-06-01")], false)] + fn forwarded_credential_is_detected_in_either_header( + #[case] forwarded: &[(&str, &str)], + #[case] expected: bool, + ) { + let headers: Headers = forwarded + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + assert_eq!(has_anthropic_credential(&headers), expected); + } + + fn credential(auth: Option) -> Option<(&'static str, String)> { + auth.map(|auth| match auth { + AuthScheme::Credential { placement, secret } => { + (placement.header_name(), secret.expose().to_string()) + } + other => panic!("expected a credential, got {other:?}"), + }) } fn tagged(encrypted: &str) -> String { @@ -1195,50 +1452,268 @@ mod tests { assert_eq!(twice, once); } + const OAUTH_TOKEN: &str = "sk-ant-oat01-token"; + const OAUTH_BEARER: &str = "Bearer sk-ant-oat01-token"; + const REGULAR_KEY: &str = "sk-ant-api03-regular"; + const OAUTH_BETA: &str = "oauth-2025-04-20"; + const BROWSER_ACCESS: (&str, &str) = ("anthropic-dangerous-direct-browser-access", "true"); + #[rstest] - #[case::no_existing_header(None, "b", "b")] - #[case::empty_existing_header(Some(""), "b", "b")] - #[case::whitespace_existing_header(Some(" "), "b", "b")] - #[case::sorted_after_merge(Some("c,a"), "b", "a,b,c")] - #[case::already_present(Some("a,b"), "a", "a,b")] - #[case::trimmed_and_deduplicated(Some("b, a ,b"), "c", "a,b,c")] - #[case::blank_pieces_skipped(Some("a,,b"), "c", "a,b,c")] - fn beta_values_merge_sorted_and_deduplicated( - #[case] existing: Option<&str>, - #[case] new_beta: &str, - #[case] expected: &str, + #[case::no_beta_header(&[("x-api-key", "k")], &[], &[("x-api-key", "k")])] + #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], &[], &[("Anthropic-Beta", " , "), ("x-api-key", "k")])] + #[case::added_to_no_header(&[("x-api-key", "k")], &["b"], &[("x-api-key", "k"), ("anthropic-beta", "b")])] + #[case::added_to_blank_header(&[("anthropic-beta", " ")], &["b"], &[("anthropic-beta", "b")])] + #[case::sorted_after_merge(&[("anthropic-beta", "c,a")], &["b"], &[("anthropic-beta", "a,b,c")])] + #[case::already_present(&[("anthropic-beta", "a,b")], &["a"], &[("anthropic-beta", "a,b")])] + #[case::existing_normalized_without_additions( + &[("Anthropic-Beta", "b, a ,b"), ("x-api-key", "k")], + &[], + &[("x-api-key", "k"), ("anthropic-beta", "a,b")] + )] + #[case::every_casing_unioned_into_one_lowercase_header( + &[("anthropic-beta", "a"), ("ANTHROPIC-BETA", "c"), ("x-api-key", "k")], + &["b"], + &[("x-api-key", "k"), ("anthropic-beta", "a,b,c")] + )] + fn merge_beta_headers_replaces_the_header_with_the_sorted_union( + #[case] input: &[(&str, &str)], + #[case] added: &[&str], + #[case] expected: &[(&str, &str)], ) { assert_eq!( - join_beta_values(split_beta_values(existing).chain([new_beta.to_string()])), + merge_beta_headers(headers(input), betas(added)), + headers(expected) + ); + } + + #[rstest] + #[case::raw_token(OAUTH_TOKEN, Some(OAUTH_TOKEN))] + #[case::bare_prefix(ANTHROPIC_OAUTH_TOKEN_PREFIX, Some(ANTHROPIC_OAUTH_TOKEN_PREFIX))] + #[case::bearer_token(OAUTH_BEARER, None)] + #[case::api_key(REGULAR_KEY, None)] + #[case::empty("", None)] + #[case::uppercase_prefix("sk-ant-OAT01-abc123", None)] + #[case::prefix_not_at_start(" sk-ant-oat01-abc123", None)] + fn oauth_token_parses_only_the_raw_token(#[case] value: &str, #[case] expected: Option<&str>) { + assert_eq!(OauthToken::parse(value).map(OauthToken::as_str), expected); + } + + #[rstest] + #[case::raw_token(OAUTH_TOKEN, Some(OAUTH_TOKEN))] + #[case::bearer_token(OAUTH_BEARER, Some(OAUTH_TOKEN))] + #[case::api_key(REGULAR_KEY, None)] + #[case::bearer_api_key("Bearer sk-ant-api01-abc123", None)] + #[case::empty("", None)] + #[case::shouting_prefix("SK-ANT-OAT01-abc123", None)] + #[case::lowercase_bearer("bearer sk-ant-oat01-abc123", None)] + #[case::bearer_stripped_once("Bearer Bearer sk-ant-oat01-abc123", None)] + fn oauth_key_parses_the_token_behind_an_optional_bearer( + #[case] value: &str, + #[case] expected: Option<&str>, + ) { + assert_eq!( + OauthToken::parse_key(value).map(OauthToken::as_str), expected ); } #[rstest] - #[case::raw_token("sk-ant-oat01-abc123", true)] - #[case::bearer_token("Bearer sk-ant-oat02-xyz789", true)] - #[case::bare_prefix(ANTHROPIC_OAUTH_TOKEN_PREFIX, true)] - #[case::api_key("sk-ant-api01-abc123", false)] - #[case::bearer_api_key("Bearer sk-ant-api01-abc123", false)] - #[case::empty("", false)] - #[case::uppercase_prefix("sk-ant-OAT01-abc123", false)] - #[case::shouting_prefix("SK-ANT-OAT01-abc123", false)] - #[case::lowercase_bearer("bearer sk-ant-oat01-abc123", false)] - #[case::bearer_stripped_once("Bearer Bearer sk-ant-oat01-abc123", false)] - #[case::prefix_not_at_start(" sk-ant-oat01-abc123", false)] - fn anthropic_oauth_key_detection(#[case] value: &str, #[case] expected: bool) { - assert_eq!(is_anthropic_oauth_key(value), expected); + #[case::bearer(&[("authorization", OAUTH_BEARER)], Some(OAUTH_TOKEN))] + #[case::uppercase_header(&[("AUTHORIZATION", OAUTH_BEARER)], Some(OAUTH_TOKEN))] + #[case::non_oauth_bearer(&[("authorization", "Bearer some-proxy-token")], None)] + #[case::token_without_the_bearer_scheme(&[("authorization", OAUTH_TOKEN)], None)] + #[case::lowercase_bearer_scheme(&[("authorization", "bearer sk-ant-oat01-token")], None)] + #[case::token_in_x_api_key(&[("x-api-key", OAUTH_TOKEN)], None)] + #[case::no_headers(&[], None)] + fn forwarded_oauth_bearer_reads_the_authorization_header( + #[case] forwarded: &[(&str, &str)], + #[case] expected: Option<&str>, + ) { + assert_eq!( + forwarded_oauth_bearer(&headers(forwarded)).map(OauthToken::as_str), + expected + ); } #[rstest] - #[case::regex_tool(Some(json!([{"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[0], "name": "tool_search_tool_regex"}])), true)] - #[case::bm25_tool(Some(json!([{"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[1], "name": "tool_search_tool_bm25"}])), true)] + #[case::forwarded_bearer_drops_forwarded_and_deployment_keys( + &[("X-Api-Key", REGULAR_KEY), ("Authorization", OAUTH_BEARER)], + Some(REGULAR_KEY), + &[], + )] + #[case::forwarded_bearer_keeps_unrelated_headers_in_place( + &[("anthropic-version", "2023-06-01"), ("authorization", OAUTH_BEARER)], + None, + &[("anthropic-version", "2023-06-01")], + )] + #[case::forwarded_bearer_wins_over_an_oauth_api_key( + &[("authorization", OAUTH_BEARER)], + Some("sk-ant-oat01-deployment"), + &[], + )] + #[case::api_key_alone(&[], Some(OAUTH_TOKEN), &[])] + #[case::api_key_removes_a_forwarded_x_api_key(&[("x-api-key", OAUTH_TOKEN)], Some(OAUTH_TOKEN), &[])] + #[case::api_key_keeps_a_forwarded_non_oauth_bearer( + &[("Authorization", "Bearer some-proxy-token")], + Some(OAUTH_TOKEN), + &[("Authorization", "Bearer some-proxy-token")], + )] + fn oauth_token_is_the_whole_credential( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + #[case] kept: &[(&str, &str)], + ) { + let expected = kept + .iter() + .copied() + .chain([("anthropic-beta", OAUTH_BETA), BROWSER_ACCESS]) + .collect::>(); + assert_eq!( + optionally_handle_anthropic_oauth(headers(forwarded), api_key), + OauthHandling::Bearer { + headers: headers(&expected), + token: SecretValue::new(OAUTH_TOKEN), + } + ); + } + + #[rstest] + #[case::forwarded_bearer_merges_a_differently_cased_beta_header( + &[("Anthropic-Beta", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], + None, + )] + #[case::forwarded_bearer_dedupes_an_existing_oauth_beta( + &[("anthropic-beta", "web-search-2025-03-05, oauth-2025-04-20"), ("authorization", OAUTH_BEARER)], + None, + )] + #[case::api_key_merges_the_existing_beta_header( + &[("anthropic-beta", " web-search-2025-03-05 ,")], + Some(OAUTH_TOKEN), + )] + #[case::forwarded_bearer_unions_every_beta_header_casing( + &[("anthropic-beta", "oauth-2025-04-20"), ("ANTHROPIC-BETA", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], + None, + )] + fn oauth_beta_merges_into_existing_betas( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + ) { + assert_eq!( + optionally_handle_anthropic_oauth(headers(forwarded), api_key), + OauthHandling::Bearer { + headers: headers(&[ + ("anthropic-beta", "oauth-2025-04-20,web-search-2025-03-05"), + BROWSER_ACCESS, + ]), + token: SecretValue::new(OAUTH_TOKEN), + } + ); + } + + #[rstest] + #[case::x_api_key(&[("x-api-key", "caller-key")], Some("sk-other"))] + #[case::non_oauth_bearer(&[("Authorization", "Bearer some-proxy-token")], Some(REGULAR_KEY))] + #[case::oauth_token_without_the_bearer_scheme(&[("authorization", OAUTH_TOKEN)], None)] + #[case::bearer_prefixed_api_key(&[], Some(OAUTH_BEARER))] + #[case::nothing(&[], None)] + fn without_an_oauth_token_the_headers_are_untouched( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + ) { + assert_eq!( + optionally_handle_anthropic_oauth(headers(forwarded), api_key), + OauthHandling::Untouched(headers(forwarded)) + ); + } + + #[rstest] + #[case::api_key_param(Some("sk-param"), &[], Some(("x-api-key", "sk-param")))] + #[case::api_key_param_over_env_key_and_auth_token( + Some("sk-param"), + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("x-api-key", "sk-param")), + )] + #[case::env_key_without_a_param(None, &[("ANTHROPIC_API_KEY", "sk-env")], Some(("x-api-key", "sk-env")))] + #[case::env_key_when_the_param_is_blank(Some(" "), &[("ANTHROPIC_API_KEY", "sk-env")], Some(("x-api-key", "sk-env")))] + #[case::env_key_over_auth_token( + None, + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("x-api-key", "sk-env")), + )] + #[case::auth_token_as_a_bearer( + None, + &[("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("Authorization", "env-token")), + )] + #[case::auth_token_when_the_env_key_is_blank( + None, + &[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("Authorization", "env-token")), + )] + #[case::oauth_param_as_a_bearer(Some(OAUTH_TOKEN), &[], Some(("Authorization", OAUTH_TOKEN)))] + #[case::bearer_prefixed_oauth_env_key_as_a_bearer_once( + None, + &[("ANTHROPIC_API_KEY", OAUTH_BEARER)], + Some(("Authorization", OAUTH_TOKEN)), + )] + #[case::no_credentials(None, &[], None)] + #[case::blank_everything(Some(""), &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")], None)] + fn auth_header_prefers_the_key_then_the_auth_token( + #[case] api_key: Option<&str>, + #[case] vars: &'static [(&'static str, &'static str)], + #[case] expected: Option<(&str, &str)>, + ) { + assert_eq!( + credential(get_auth_header(api_key, &env(vars))), + expected.map(|(header, secret)| (header, secret.to_string())) + ); + } + + #[rstest] + #[case::param(Some("sk-param"), &[("ANTHROPIC_API_KEY", "sk-env")], Ok("sk-param"))] + #[case::blank_param_falls_back_to_env(Some(" "), &[("ANTHROPIC_API_KEY", "sk-env")], Ok("sk-env"))] + #[case::env_without_param(None, &[("ANTHROPIC_API_KEY", "sk-env")], Ok("sk-env"))] + #[case::blank_env_is_missing(None, &[("ANTHROPIC_API_KEY", " ")], Err(()))] + #[case::nothing_is_missing(None, &[], Err(()))] + fn api_key_resolution( + #[case] api_key: Option<&str>, + #[case] vars: &'static [(&'static str, &'static str)], + #[case] expected: Result<&str, ()>, + ) { + assert_eq!( + resolve_anthropic_api_key(api_key, &env(vars)).map_err(|error| { + assert!(matches!( + error, + litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + } + )); + }), + expected.map(str::to_string) + ); + } + + #[rstest] + #[case::absent(None, None)] + #[case::blank(Some(" \t "), None)] + #[case::padded(Some(" value "), Some("value"))] + fn non_empty_trims_and_drops_blank_values( + #[case] value: Option<&str>, + #[case] expected: Option<&str>, + ) { + assert_eq!(non_empty(value), expected); + } + + #[rstest] + #[case::regex_tool(Some(json!([{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}])), true)] + #[case::bm25_tool(Some(json!([{"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"}])), true)] #[case::after_other_tools( - Some(json!([{"name": "get_weather", "input_schema": {}}, {"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[1]}])), + Some(json!([{"name": "get_weather", "input_schema": {}}, {"type": "tool_search_tool_bm25_20251119"}])), true )] #[case::function_tool(Some(json!([{"type": "function", "function": {"name": "get_weather"}}])), false)] - #[case::name_without_type(Some(json!([{"name": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[0]}])), false)] + #[case::name_without_type(Some(json!([{"name": "tool_search_tool_regex_20251119"}])), false)] #[case::empty_tools(Some(json!([])), false)] #[case::no_tools(None, false)] fn tool_search_detection(#[case] input: Option, #[case] expected: bool) { @@ -1246,8 +1721,8 @@ mod tests { } #[rstest] - #[case::advisor_tool(Some(json!([{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor"}])), true)] - #[case::after_other_tools(Some(json!([{"name": "f", "input_schema": {}}, {"type": ANTHROPIC_ADVISOR_TOOL_TYPE}])), true)] + #[case::advisor_tool(Some(json!([{"type": "advisor_20260301", "name": "advisor"}])), true)] + #[case::after_other_tools(Some(json!([{"name": "f", "input_schema": {}}, {"type": "advisor_20260301"}])), true)] #[case::tool_named_advisor(Some(json!([{"name": "advisor", "input_schema": {}}])), false)] #[case::other_server_tool(Some(json!([{"type": "web_search_20250305", "name": "web_search"}])), false)] #[case::empty_tools(Some(json!([])), false)] diff --git a/litellm-rust/crates/llms/src/anthropic/messages/headers.rs b/litellm-rust/crates/llms/src/anthropic/messages/headers.rs deleted file mode 100644 index bd1b11be92d..00000000000 --- a/litellm-rust/crates/llms/src/anthropic/messages/headers.rs +++ /dev/null @@ -1,677 +0,0 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; -use litellm_types::{ - llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest, recognized::Recognized, -}; -use serde_json::Value; - -use crate::{ - anthropic::{ - ANTHROPIC_OAUTH_TOKEN_PREFIX, - common_utils::{ - ANTHROPIC_OAUTH_BETA_HEADER, beta, has_advisor_tool, is_anthropic_oauth_key, - is_tool_search_used, join_beta_values, requires_native_compaction_beta, - split_beta_values, - }, - }, - base_llm::{ - anthropic_messages::transformation::Headers, - auth::{AuthScheme, ValidatedEnvironment}, - }, -}; - -const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; -const ANTHROPIC_AUTH_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN"; -const BETA_HEADER: &str = "anthropic-beta"; -const AUTHORIZATION: &str = "authorization"; -const API_KEY_HEADER: &str = "x-api-key"; -const DIRECT_BROWSER_ACCESS_HEADER: &str = "anthropic-dangerous-direct-browser-access"; - -fn header_value<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { - headers - .iter() - .find(|(header, _)| header.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.as_str()) -} - -fn without(headers: Headers, names: &[&str]) -> Headers { - headers - .into_iter() - .filter(|(header, _)| !names.iter().any(|name| header.eq_ignore_ascii_case(name))) - .collect() -} - -fn existing_betas(headers: &[(String, String)]) -> impl Iterator + '_ { - headers - .iter() - .filter(|(header, _)| header.eq_ignore_ascii_case(BETA_HEADER)) - .flat_map(|(_, value)| split_beta_values(Some(value))) -} - -/// The OAuth headers Python's `optionally_handle_anthropic_oauth` sets next to the bearer. -fn with_oauth_companions(headers: Headers, dropped: &[&str]) -> Headers { - let beta = - join_beta_values(existing_betas(&headers).chain([ANTHROPIC_OAUTH_BETA_HEADER.to_string()])); - without(headers, &[dropped, &[BETA_HEADER]].concat()) - .into_iter() - .chain([ - (BETA_HEADER.to_string(), beta), - (DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string()), - ]) - .collect() -} - -fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -fn bearer(token: &str) -> AuthScheme { - AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - secret: SecretValue::new(token), - } -} - -pub fn validate_environment( - headers: Headers, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - if let Some(token) = header_value(&headers, AUTHORIZATION) - .and_then(|forwarded| forwarded.strip_prefix("Bearer ")) - .filter(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) - { - let auth = bearer(token); - return Ok(ValidatedEnvironment { - headers: with_oauth_companions(headers, &[API_KEY_HEADER, AUTHORIZATION]), - auth, - }); - } - if let Some(key) = api_key.filter(|key| key.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) { - return Ok(ValidatedEnvironment { - headers: with_oauth_companions(headers, &[API_KEY_HEADER]), - auth: bearer(key), - }); - } - if header_value(&headers, API_KEY_HEADER).is_some() - || header_value(&headers, AUTHORIZATION).is_some() - { - return Ok(ValidatedEnvironment { - headers, - auth: AuthScheme::Forwarded, - }); - } - let resolved_key = non_empty(api_key) - .map(str::to_string) - .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())); - let auth = match resolved_key { - Some(key) if is_anthropic_oauth_key(&key) => bearer(&key), - Some(key) => AuthScheme::Credential { - placement: CredentialPlacement::Header(API_KEY_HEADER), - secret: SecretValue::new(key), - }, - None => match env_lookup(ANTHROPIC_AUTH_TOKEN_ENV).filter(|value| !value.trim().is_empty()) - { - Some(token) => bearer(&token), - None => { - return Err(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: ANTHROPIC_API_KEY_ENV, - }); - } - }, - }; - Ok(ValidatedEnvironment { headers, auth }) -} - -fn context_management_betas( - context_management: Option<&Value>, -) -> impl Iterator { - let edits = context_management - .and_then(|value| value.get("edits")) - .and_then(Value::as_array) - .map(Vec::as_slice) - .unwrap_or(&[]); - let (compact, other) = edits.iter().fold((false, false), |(compact, other), edit| { - match edit.get("type").and_then(Value::as_str) { - Some("compact_20260112") => (true, other), - _ => (compact, true), - } - }); - compact - .then_some(beta::COMPACT_2026_01_12) - .into_iter() - .chain(other.then_some(beta::CONTEXT_MANAGEMENT_2025_06_27)) -} - -fn uses_structured_output(request: &AnthropicMessagesRequest) -> bool { - request.params.output_format.is_some() - || request - .params - .output_config - .as_ref() - .and_then(Recognized::known) - .is_some_and(|config| config.format.is_some()) -} - -fn messages_carry_output_config(request: &AnthropicMessagesRequest) -> bool { - request - .messages - .iter() - .any(|message| message.extra.contains_key("output_config")) -} - -pub fn feature_betas(request: &AnthropicMessagesRequest) -> Vec<&'static str> { - let tools = request.params.tools.as_deref(); - [ - requires_native_compaction_beta(request.params.compaction.as_ref(), &request.messages) - .then_some(beta::COMPACT_2026_09_04), - uses_structured_output(request).then_some(beta::STRUCTURED_OUTPUT), - (request.params.speed.as_deref() == Some("fast")).then_some(beta::FAST_MODE_2026_02_01), - messages_carry_output_config(request).then_some(beta::PER_TURN_CONTROL_2026_07_01), - has_advisor_tool(tools).then_some(beta::ADVISOR_TOOL_2026_03_01), - is_tool_search_used(tools).then_some(beta::ADVANCED_TOOL_USE_2025_11_20), - ] - .into_iter() - .flatten() - .chain(context_management_betas( - request.params.context_management.as_ref(), - )) - .collect() -} - -pub fn with_feature_betas(headers: Headers, request: &AnthropicMessagesRequest) -> Headers { - let existing = existing_betas(&headers).collect::>(); - let features = feature_betas(request); - if existing.is_empty() && features.is_empty() { - return headers; - } - let merged = join_beta_values( - existing - .into_iter() - .chain(features.into_iter().map(str::to_string)), - ); - without(headers, &[BETA_HEADER]) - .into_iter() - .chain([(BETA_HEADER.to_string(), merged)]) - .collect() -} - -#[cfg(test)] -mod tests { - use rstest::{fixture, rstest}; - use serde_json::json; - - use super::*; - use crate::base_llm::auth::resolve_auth; - - const OAUTH_TOKEN: &str = "sk-ant-oat01-token"; - const OAUTH_BEARER: &str = "Bearer sk-ant-oat01-token"; - const REGULAR_KEY: &str = "sk-ant-api03-regular"; - const BROWSER_ACCESS: (&str, &str) = ("anthropic-dangerous-direct-browser-access", "true"); - - type Env = &'static [(&'static str, &'static str)]; - - fn request(fields: Value) -> AnthropicMessagesRequest { - let mut body = - json!({"model": "claude", "messages": [{"role": "user", "content": "Hello"}]}); - body.as_object_mut() - .unwrap() - .extend(fields.as_object().unwrap().clone()); - serde_json::from_value(body).unwrap() - } - - fn headers(pairs: &[(&str, &str)]) -> Headers { - pairs - .iter() - .map(|(name, value)| (name.to_string(), value.to_string())) - .collect() - } - - fn betas(values: &[&str]) -> String { - values.join(",") - } - - #[fixture] - fn no_env() -> Env { - &[] - } - - #[fixture] - fn full_env() -> Env { - &[ - ("ANTHROPIC_API_KEY", "sk-env"), - ("ANTHROPIC_AUTH_TOKEN", "env-token"), - ] - } - - fn authenticate_with( - forwarded: &[(&str, &str)], - api_key: Option<&str>, - env: Env, - ) -> Result { - let lookup = |name: &str| { - env.iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()) - }; - let validated = validate_environment(headers(forwarded), api_key, &lookup)?; - let resolved = tokio::runtime::Builder::new_current_thread() - .build() - .unwrap() - .block_on(resolve_auth( - &litellm_auth::AuthServices::default(), - validated, - &lookup, - )) - .unwrap(); - Ok(resolved.headers) - } - - #[rstest] - #[case::forwarded_bearer_drops_forwarded_and_deployment_keys( - &[("X-Api-Key", REGULAR_KEY), ("Authorization", OAUTH_BEARER)], - Some(REGULAR_KEY), - OAUTH_BEARER, - &[], - )] - #[case::forwarded_bearer_in_uppercase_authorization_header( - &[("AUTHORIZATION", OAUTH_BEARER)], - None, - OAUTH_BEARER, - &[], - )] - #[case::forwarded_bearer_keeps_unrelated_headers_in_place( - &[("anthropic-version", "2023-06-01"), ("authorization", OAUTH_BEARER)], - None, - OAUTH_BEARER, - &[("anthropic-version", "2023-06-01")], - )] - #[case::forwarded_bearer_wins_over_an_oauth_api_key( - &[("authorization", OAUTH_BEARER)], - Some("sk-ant-oat01-deployment"), - OAUTH_BEARER, - &[], - )] - #[case::api_key_authenticates_as_a_bearer(&[], Some(OAUTH_TOKEN), OAUTH_BEARER, &[])] - #[case::api_key_removes_a_forwarded_x_api_key( - &[("x-api-key", OAUTH_TOKEN)], - Some(OAUTH_TOKEN), - OAUTH_BEARER, - &[], - )] - #[case::api_key_replaces_a_forwarded_non_oauth_bearer( - &[("Authorization", "Bearer some-proxy-token")], - Some(OAUTH_TOKEN), - OAUTH_BEARER, - &[], - )] - fn oauth_token_is_the_whole_credential( - #[case] forwarded: &[(&str, &str)], - #[case] api_key: Option<&str>, - #[case] expected_bearer: &str, - #[case] kept: &[(&str, &str)], - full_env: Env, - ) { - let expected = kept - .iter() - .copied() - .chain([ - ("anthropic-beta", ANTHROPIC_OAUTH_BETA_HEADER), - BROWSER_ACCESS, - ("authorization", expected_bearer), - ]) - .collect::>(); - assert_eq!( - authenticate_with(forwarded, api_key, full_env).unwrap(), - headers(&expected) - ); - } - - #[rstest] - #[case::forwarded_bearer_merges_a_differently_cased_beta_header( - &[("Anthropic-Beta", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], - None, - )] - #[case::forwarded_bearer_dedupes_an_existing_oauth_beta( - &[("anthropic-beta", "web-search-2025-03-05, oauth-2025-04-20"), ("authorization", OAUTH_BEARER)], - None, - )] - #[case::api_key_merges_the_existing_beta_header( - &[("anthropic-beta", " web-search-2025-03-05 ,")], - Some(OAUTH_TOKEN), - )] - #[case::forwarded_bearer_unions_every_beta_header_casing( - &[("anthropic-beta", "oauth-2025-04-20"), ("ANTHROPIC-BETA", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], - None, - )] - fn oauth_beta_merges_into_existing_betas( - #[case] forwarded: &[(&str, &str)], - #[case] api_key: Option<&str>, - no_env: Env, - ) { - assert_eq!( - authenticate_with(forwarded, api_key, no_env).unwrap(), - headers(&[ - ( - "anthropic-beta", - &betas(&[ANTHROPIC_OAUTH_BETA_HEADER, "web-search-2025-03-05"]) - ), - BROWSER_ACCESS, - ("authorization", OAUTH_BEARER), - ]) - ); - } - - #[rstest] - #[case::x_api_key_over_the_deployment_key(&[("x-api-key", "caller-key")], Some("sk-other"))] - #[case::uppercase_x_api_key(&[("X-API-KEY", "caller-key")], None)] - #[case::non_oauth_bearer(&[("Authorization", "Bearer some-proxy-token")], None)] - #[case::non_oauth_bearer_over_a_regular_api_key( - &[("authorization", "Bearer sk-ant-api03-forwarded")], - Some(REGULAR_KEY), - )] - #[case::oauth_token_without_the_bearer_scheme(&[("authorization", OAUTH_TOKEN)], None)] - #[case::oauth_token_behind_a_lowercase_bearer_scheme( - &[("authorization", "bearer sk-ant-oat01-token")], - None, - )] - fn forwarded_auth_header_is_kept_untouched( - #[case] forwarded: &[(&str, &str)], - #[case] api_key: Option<&str>, - full_env: Env, - ) { - assert_eq!( - authenticate_with(forwarded, api_key, full_env).unwrap(), - headers(forwarded) - ); - } - - #[rstest] - #[case::api_key_param(Some("sk-param"), &[], ("x-api-key", "sk-param"))] - #[case::api_key_param_over_env_key_and_auth_token( - Some("sk-param"), - &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], - ("x-api-key", "sk-param"), - )] - #[case::env_key_without_a_param(None, &[("ANTHROPIC_API_KEY", "sk-env")], ("x-api-key", "sk-env"))] - #[case::env_key_when_the_param_is_empty(Some(""), &[("ANTHROPIC_API_KEY", "sk-env")], ("x-api-key", "sk-env"))] - #[case::env_key_when_the_param_is_whitespace( - Some(" "), - &[("ANTHROPIC_API_KEY", "sk-env")], - ("x-api-key", "sk-env"), - )] - #[case::env_key_over_auth_token( - None, - &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], - ("x-api-key", "sk-env"), - )] - #[case::auth_token_as_a_bearer( - None, - &[("ANTHROPIC_AUTH_TOKEN", "env-token")], - ("authorization", "Bearer env-token"), - )] - #[case::auth_token_when_the_env_key_is_whitespace( - None, - &[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], - ("authorization", "Bearer env-token"), - )] - #[case::oauth_env_key_as_a_plain_bearer( - None, - &[("ANTHROPIC_API_KEY", "sk-ant-oat01-env")], - ("authorization", "Bearer sk-ant-oat01-env"), - )] - fn credential_is_resolved_after_the_existing_headers( - #[case] api_key: Option<&str>, - #[case] env: Env, - #[case] expected: (&str, &str), - ) { - let forwarded = [("anthropic-beta", "web-search-2025-03-05")]; - assert_eq!( - authenticate_with(&forwarded, api_key, env).unwrap(), - headers(&[forwarded[0], expected]) - ); - } - - #[rstest] - #[case::no_credentials(&[], None, &[])] - #[case::empty_api_key(&[], Some(""), &[])] - #[case::whitespace_only_env_values( - &[], - None, - &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")], - )] - #[case::unrelated_forwarded_headers(&[("anthropic-beta", "web-search-2025-03-05")], None, &[])] - fn missing_credentials_are_an_auth_error( - #[case] forwarded: &[(&str, &str)], - #[case] api_key: Option<&str>, - #[case] env: Env, - ) { - assert!(matches!( - authenticate_with(forwarded, api_key, env), - Err(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: "ANTHROPIC_API_KEY", - }) - )); - } - - #[rstest] - #[case::no_features(json!({}), &[])] - #[case::output_format(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])] - #[case::null_output_format(json!({"output_format": null}), &[])] - #[case::output_config_format( - json!({"output_config": {"format": {"type": "json_schema"}, "effort": "xhigh"}}), - &[beta::STRUCTURED_OUTPUT] - )] - #[case::null_output_config_format(json!({"output_config": {"format": null}}), &[])] - #[case::top_level_output_config_without_format(json!({"output_config": {"effort": "high"}}), &[])] - #[case::fast_speed(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])] - #[case::standard_speed(json!({"speed": "standard"}), &[])] - #[case::compaction_param(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])] - #[case::empty_compaction_param(json!({"compaction": {}}), &[beta::COMPACT_2026_09_04])] - #[case::signed_compaction_block_in_history( - json!({"messages": [ - {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]}, - {"role": "user", "content": "Continue"}, - ]}), - &[beta::COMPACT_2026_09_04] - )] - #[case::unsigned_compaction_block_in_history( - json!({"messages": [ - {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": ""}]}, - {"role": "user", "content": "Continue"}, - ]}), - &[] - )] - #[case::advisor_tool( - json!({"tools": [{"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-6"}]}), - &[beta::ADVISOR_TOOL_2026_03_01] - )] - #[case::no_tools(json!({"tools": []}), &[])] - #[case::regex_tool_search( - json!({"tools": [{"type": "tool_search_tool_regex_20251119"}]}), - &[beta::ADVANCED_TOOL_USE_2025_11_20] - )] - #[case::bm25_tool_search( - json!({"tools": [{"type": "tool_search_tool_bm25_20251119"}]}), - &[beta::ADVANCED_TOOL_USE_2025_11_20] - )] - #[case::unrelated_server_tool(json!({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}), &[])] - #[case::only_compact_edits( - json!({"context_management": {"edits": [{"type": "compact_20260112"}]}}), - &[beta::COMPACT_2026_01_12] - )] - #[case::only_other_edits( - json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919", "keep": {"type": "tool_uses", "value": 3}}]}}), - &[beta::CONTEXT_MANAGEMENT_2025_06_27] - )] - #[case::compact_and_other_edits( - json!({"context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_tool_uses_20250919"}]}}), - &[beta::COMPACT_2026_01_12, beta::CONTEXT_MANAGEMENT_2025_06_27] - )] - #[case::edit_without_a_type(json!({"context_management": {"edits": [{}]}}), &[beta::CONTEXT_MANAGEMENT_2025_06_27])] - #[case::empty_edits(json!({"context_management": {"edits": []}}), &[])] - #[case::context_management_without_edits(json!({"context_management": {}}), &[])] - #[case::per_message_output_config( - json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), - &[beta::PER_TURN_CONTROL_2026_07_01] - )] - #[case::per_message_null_output_config( - json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}), - &[beta::PER_TURN_CONTROL_2026_07_01] - )] - fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { - assert_eq!(feature_betas(&request(fields)), expected); - } - - #[rstest] - #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}))] - #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}))] - fn headers_without_any_beta_value_are_untouched( - #[case] input: &[(&str, &str)], - #[case] fields: Value, - ) { - assert_eq!( - with_feature_betas(headers(input), &request(fields)), - headers(input) - ); - } - - #[rstest] - #[case::feature_beta_is_appended( - &[("x-api-key", "k")], - json!({"speed": "fast"}), - &[("x-api-key", "k"), ("anthropic-beta", beta::FAST_MODE_2026_02_01)], - )] - #[case::existing_betas_are_normalized_without_features( - &[("Anthropic-Beta", "web-search-2025-03-05, interleaved-thinking-2025-05-14 ,web-search-2025-03-05"), ("x-api-key", "k")], - json!({}), - &[("x-api-key", "k"), ("anthropic-beta", "interleaved-thinking-2025-05-14,web-search-2025-03-05")], - )] - #[case::existing_advisor_beta_is_kept_without_an_advisor_tool( - &[("anthropic-beta", beta::ADVISOR_TOOL_2026_03_01)], - json!({"tools": []}), - &[("anthropic-beta", beta::ADVISOR_TOOL_2026_03_01)], - )] - #[case::feature_already_sent_is_not_duplicated( - &[("anthropic-beta", beta::FAST_MODE_2026_02_01)], - json!({"speed": "fast"}), - &[("anthropic-beta", beta::FAST_MODE_2026_02_01)], - )] - fn feature_betas_merge_into_the_headers( - #[case] input: &[(&str, &str)], - #[case] fields: Value, - #[case] expected: &[(&str, &str)], - ) { - assert_eq!( - with_feature_betas(headers(input), &request(fields)), - headers(expected) - ); - } - - #[test] - fn differently_cased_beta_header_is_replaced_by_one_sorted_header() { - let merged = with_feature_betas( - headers(&[("Anthropic-Beta", "interleaved-thinking-2025-05-14")]), - &request( - json!({"messages": [{"role": "system", "content": "env", "output_config": {"effort": "low"}}]}), - ), - ); - assert_eq!( - merged, - headers(&[( - "anthropic-beta", - &betas(&[ - "interleaved-thinking-2025-05-14", - beta::PER_TURN_CONTROL_2026_07_01 - ]) - )]) - ); - } - - #[test] - fn every_beta_header_casing_is_unioned_into_one_header() { - let merged = with_feature_betas( - headers(&[ - ("anthropic-beta", "interleaved-thinking-2025-05-14"), - ("Anthropic-Beta", "web-search-2025-03-05"), - ]), - &request(json!({"speed": "fast"})), - ); - assert_eq!( - merged, - headers(&[( - "anthropic-beta", - &betas(&[ - beta::FAST_MODE_2026_02_01, - "interleaved-thinking-2025-05-14", - "web-search-2025-03-05" - ]) - )]) - ); - } - - #[test] - fn unknown_client_betas_survive_alongside_the_added_one() { - let client_betas = [ - "claude-code-20250219", - "interleaved-thinking-2025-05-14", - beta::CONTEXT_MANAGEMENT_2025_06_27, - beta::PER_TURN_CONTROL_2026_07_01, - "effort-2025-11-24", - ]; - let merged = with_feature_betas( - headers(&[("anthropic-beta", &betas(&client_betas))]), - &request( - json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), - ), - ); - assert_eq!( - merged, - headers(&[( - "anthropic-beta", - &betas(&[ - "claude-code-20250219", - beta::CONTEXT_MANAGEMENT_2025_06_27, - "effort-2025-11-24", - "interleaved-thinking-2025-05-14", - beta::PER_TURN_CONTROL_2026_07_01, - ]) - )]) - ); - } - - #[test] - fn every_feature_merges_with_the_oauth_beta_sorted_and_last() { - let oauth_headers = authenticate_with(&[], Some(OAUTH_TOKEN), &[]).unwrap(); - let all_features = request(json!({ - "compaction": {"enabled": true}, - "output_format": {"type": "json_schema"}, - "speed": "fast", - "tools": [{"type": "advisor_20260301"}, {"type": "tool_search_tool_bm25_20251119"}], - "context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_thinking_20251015"}]}, - "messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}], - })); - assert_eq!( - with_feature_betas(oauth_headers, &all_features), - headers(&[ - BROWSER_ACCESS, - ("authorization", OAUTH_BEARER), - ( - "anthropic-beta", - &betas(&[ - beta::ADVANCED_TOOL_USE_2025_11_20, - beta::ADVISOR_TOOL_2026_03_01, - beta::COMPACT_2026_01_12, - beta::COMPACT_2026_09_04, - beta::CONTEXT_MANAGEMENT_2025_06_27, - beta::FAST_MODE_2026_02_01, - ANTHROPIC_OAUTH_BETA_HEADER, - beta::PER_TURN_CONTROL_2026_07_01, - beta::STRUCTURED_OUTPUT, - ]) - ), - ]) - ); - } -} diff --git a/litellm-rust/crates/llms/src/anthropic/messages/mod.rs b/litellm-rust/crates/llms/src/anthropic/messages/mod.rs index 5adf5fda16f..dff4bf18bd5 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/mod.rs @@ -1,5 +1,4 @@ pub mod handler; -pub mod headers; pub mod streaming_iterator; pub mod thinking; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 280ea63eefa..07c2eb46eaa 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -1,31 +1,35 @@ +use litellm_auth::CredentialPlacement; use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; -use litellm_types::llms::anthropic_messages::anthropic_request::{ - AnthropicMessagesOptionalParams, AnthropicMessagesRequest, +use litellm_types::{ + llms::{ + anthropic::{AnthropicBeta, BetaSet}, + anthropic_messages::anthropic_request::{ + AnthropicMessage, AnthropicMessagesOptionalParams, AnthropicMessagesRequest, + ContextEdit, ContextManagement, Speed, + }, + }, + recognized::Recognized, }; use serde_json::{Map, Value, json}; -use super::{ - headers::{validate_environment, with_feature_betas}, - thinking::{ThinkingBudgets, ThinkingContext, translate_thinking}, -}; +use super::thinking::{ThinkingBudgets, ThinkingContext, translate_thinking}; use crate::{ Error, anthropic::common_utils::{ - AnthropicModelCapabilities, has_advisor_tool, strip_advisor_blocks, - strip_encrypted_reasoning_blocks, + ANTHROPIC_API_BASE_ENV, ANTHROPIC_API_KEY_ENV, ANTHROPIC_AUTH_TOKEN_ENV, + ANTHROPIC_BASE_URL_ENV, AnthropicModelCapabilities, OauthHandling, complete_anthropic_url, + get_auth_header, has_advisor_tool, has_anthropic_credential, is_tool_search_used, + merge_beta_headers, optionally_handle_anthropic_oauth, requires_native_compaction_beta, + strip_advisor_blocks, strip_encrypted_reasoning_blocks, }, - base_llm::anthropic_messages::transformation::{ - BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, ValidatedEnvironment, + base_llm::{ + anthropic_messages::transformation::{ + BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, ValidatedEnvironment, + }, + auth::AuthScheme, }, }; -const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; -const ANTHROPIC_AUTH_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN"; -const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; -const ANTHROPIC_BASE_URL_ENV: &str = "ANTHROPIC_BASE_URL"; -const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com"; -const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; - pub struct AnthropicMessagesConfig; pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig; @@ -66,18 +70,15 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { context: &MessagesTransformContext, ) -> Result { if request.params.max_tokens.is_none() { - return Err(Error::InvalidRequest( - "max_tokens is required for Anthropic /v1/messages API".to_string(), - )); + return Err(Error::MissingField("max_tokens")); } let request = drop_unsupported_params(request, context)?; let request = translate_thinking(request, &context.thinking)?; let context_management = request .params .context_management - .as_ref() - .and_then(map_openai_context_management_to_anthropic) - .or_else(|| request.params.context_management.clone()); + .clone() + .map(map_openai_context_management_to_anthropic); let messages = if has_advisor_tool(request.params.tools.as_deref()) { request.messages } else { @@ -102,6 +103,8 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { ] } + /// Python's `validate_anthropic_messages_environment` up to the beta merge, which + /// `request_headers` does once the request is final. fn validate_environment( &self, headers: Headers, @@ -109,14 +112,99 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { _model: &str, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - validate_environment(headers, api_key, env_lookup).map_err(Error::from) + let headers = match optionally_handle_anthropic_oauth(headers, api_key) { + OauthHandling::Bearer { headers, token } => { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: token, + }, + }); + } + OauthHandling::Untouched(headers) => headers, + }; + if has_anthropic_credential(&headers) { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + let auth = get_auth_header(api_key, env_lookup).ok_or(Error::Auth( + litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: ANTHROPIC_API_KEY_ENV, + }, + ))?; + Ok(ValidatedEnvironment { headers, auth }) } fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers { - with_feature_betas(headers, request) + update_headers_with_anthropic_beta(headers, request) } } +fn update_headers_with_anthropic_beta( + headers: Headers, + request: &AnthropicMessagesRequest, +) -> Headers { + merge_beta_headers(headers, feature_betas(request)) +} + +fn feature_betas(request: &AnthropicMessagesRequest) -> BetaSet { + let params = &request.params; + let tools = params.tools.as_deref(); + [ + requires_native_compaction_beta(params.compaction.as_ref(), &request.messages) + .then_some(AnthropicBeta::Compact20260904), + uses_structured_output(params).then_some(AnthropicBeta::StructuredOutputs20251113), + (params.speed == Some(Recognized::Known(Speed::Fast))) + .then_some(AnthropicBeta::FastMode20260201), + messages_carry_output_config(&request.messages) + .then_some(AnthropicBeta::PerTurnControl20260701), + has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301), + is_tool_search_used(tools).then_some(AnthropicBeta::AdvancedToolUse20251120), + ] + .into_iter() + .flatten() + .chain(context_management_betas(params.context_management.as_ref())) + .collect() +} + +fn is_compact_edit(edit: &Recognized) -> bool { + matches!(edit, Recognized::Known(ContextEdit::Compact { .. })) +} + +fn context_management_betas( + context_management: Option<&Recognized>, +) -> impl Iterator { + let edits = context_management + .and_then(Recognized::known) + .and_then(|context_management| context_management.edits.as_deref()) + .unwrap_or_default(); + let compact = edits.iter().any(is_compact_edit); + let other = edits.iter().any(|edit| !is_compact_edit(edit)); + compact + .then_some(AnthropicBeta::Compact20260112) + .into_iter() + .chain(other.then_some(AnthropicBeta::ContextManagement20250627)) +} + +fn uses_structured_output(params: &AnthropicMessagesOptionalParams) -> bool { + params.output_format.is_some() + || params + .output_config + .as_ref() + .and_then(Recognized::known) + .is_some_and(|config| config.format.is_some()) +} + +fn messages_carry_output_config(messages: &[AnthropicMessage]) -> bool { + messages + .iter() + .any(|message| message.extra.contains_key("output_config")) +} + fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error { Error::InvalidRequest(format!( "{model} does not support {param}={value}. {hint}To drop unsupported params, set `litellm.drop_params = True`." @@ -136,9 +224,9 @@ fn drop_unsupported_params( Err(unsupported_param(&model, param, &value, hint)) }; let params = request.params; - let speed = match params.speed.as_deref() { + let speed = match ¶ms.speed { Some(speed) if !capabilities.supports_speed => { - reject("speed", format!("'{speed}'"), "")?; + reject("speed", format!("'{}'", speed_text(speed)), "")?; None } _ => params.speed.clone(), @@ -178,101 +266,73 @@ fn drop_unsupported_params( }) } -pub fn map_openai_context_management_to_anthropic(context_management: &Value) -> Option { - match context_management { - Value::Object(edits) if edits.contains_key("edits") => Some(context_management.clone()), - Value::Array(entries) => { - let edits: Vec = entries - .iter() - .filter_map(Value::as_object) - .filter(|entry| entry.get("type").and_then(Value::as_str) == Some("compaction")) - .map(|entry| { - let trigger = entry.get("compact_threshold").and_then(Value::as_f64).map( - |threshold| json!({"type": "input_tokens", "value": threshold as i64}), - ); - let passthrough = entry - .iter() - .filter(|(key, _)| !matches!(key.as_str(), "type" | "compact_threshold")) - .map(|(key, value)| (key.clone(), value.clone())); - Value::Object( - [("type".to_string(), json!("compact_20260112"))] - .into_iter() - .chain(trigger.map(|trigger| ("trigger".to_string(), trigger))) - .chain(passthrough) - .collect::>(), - ) - }) - .collect(); - (!edits.is_empty()).then(|| json!({"edits": edits})) - } - _ => None, +fn speed_text(speed: &Recognized) -> String { + match speed { + Recognized::Known(speed) => speed.as_str().to_string(), + Recognized::Unrecognized(Value::String(text)) => text.clone(), + Recognized::Unrecognized(other) => other.to_string(), } } -pub fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -pub fn resolve_anthropic_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - non_empty(api_key) - .map(str::to_string) - .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) - .ok_or(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: ANTHROPIC_API_KEY_ENV, - }) -} - -pub fn complete_anthropic_url( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> String { - let api_base = resolve_anthropic_api_base(api_base, env_lookup); - - let api_base = api_base.trim_end_matches('/'); - if api_base.ends_with(MESSAGES_PATH_SUFFIX) { - return api_base.to_string(); +fn compact_edit_from_openai(entry: &Map) -> Option { + if entry.get("type").and_then(Value::as_str) != Some("compaction") { + return None; } - format!("{api_base}{MESSAGES_PATH_SUFFIX}") + let trigger = entry + .get("compact_threshold") + .and_then(Value::as_f64) + .map(|threshold| json!({"type": "input_tokens", "value": threshold as i64})); + let passthrough = entry + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "compact_threshold")) + .map(|(key, value)| (key.clone(), value.clone())); + Some(ContextEdit::Compact { + extra: trigger + .map(|trigger| ("trigger".to_string(), trigger)) + .into_iter() + .chain(passthrough) + .collect(), + }) } -pub fn resolve_anthropic_api_base( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> String { - let env = |name: &str| env_lookup(name).filter(|value| !value.trim().is_empty()); - non_empty(api_base) - .map(str::to_string) - .or_else(|| env(ANTHROPIC_API_BASE_ENV)) - .or_else(|| env(ANTHROPIC_BASE_URL_ENV)) - .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string()) +/// An OpenAI-style `context_management` list becomes Anthropic `edits` when it holds +/// compaction entries. Anything else, native edits included, is sent as it came. +pub fn map_openai_context_management_to_anthropic( + context_management: Recognized, +) -> Recognized { + let Recognized::Unrecognized(Value::Array(entries)) = &context_management else { + return context_management; + }; + let edits: Vec> = entries + .iter() + .filter_map(Value::as_object) + .filter_map(compact_edit_from_openai) + .map(Recognized::Known) + .collect(); + if edits.is_empty() { + return context_management; + } + Recognized::Known(ContextManagement { + edits: Some(edits), + extra: Map::new(), + }) } #[cfg(test)] mod tests { use std::process::Command; - use litellm_auth::CredentialPlacement; use rstest::{fixture, rstest}; use super::*; - use crate::{ - anthropic::common_utils::{ENCRYPTED_REASONING_SIGNATURE_PREFIX, beta}, - base_llm::auth::AuthScheme, - }; + use crate::anthropic::common_utils::ENCRYPTED_REASONING_SIGNATURE_PREFIX; type Env = &'static [(&'static str, &'static str)]; - const BOTH_BASE_ENVS: Env = &[ - (ANTHROPIC_API_BASE_ENV, "https://api-base.example.com"), - (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com"), - ]; - const API_KEY_ENV: Env = &[(ANTHROPIC_API_KEY_ENV, "sk-env")]; - const MISSING_API_KEY: &str = - "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"; + const OAUTH_TOKEN: &str = "sk-ant-oat01-token"; + const OAUTH_BEARER: &str = "Bearer sk-ant-oat01-token"; + const OAUTH_BETA: &str = "oauth-2025-04-20"; + const BROWSER_ACCESS: (&str, &str) = ("anthropic-dangerous-direct-browser-access", "true"); const LOW_BUDGET_ENV: &str = "DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET"; const PROCESS_ENV_PROBE: &str = "LITELLM_MESSAGES_TRANSFORM_CONTEXT_PROBE"; @@ -377,7 +437,7 @@ mod tests { fn missing_max_tokens_is_rejected(#[case] fields: Value, unmapped: AnthropicModelCapabilities) { assert_eq!( transform(fields, unmapped, false), - invalid("max_tokens is required for Anthropic /v1/messages API") + Err(Error::MissingField("max_tokens")) ); } @@ -569,17 +629,19 @@ mod tests { #[case::empty_list(json!([]), None)] #[case::anthropic_edits_pass_through( json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]}), - Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]})) + None )] #[case::object_without_edits(json!({"type": "compaction"}), None)] #[case::scalar(json!("compaction"), None)] fn openai_context_management_maps_to_anthropic_edits( #[case] context_management: Value, - #[case] expected: Option, + #[case] mapped: Option, ) { + let parsed: Recognized = + serde_json::from_value(context_management.clone()).unwrap(); assert_eq!( - map_openai_context_management_to_anthropic(&context_management), - expected + serde_json::to_value(map_openai_context_management_to_anthropic(parsed)).unwrap(), + mapped.unwrap_or(context_management) ); } @@ -723,47 +785,6 @@ mod tests { ); } - #[rstest] - #[case::public_endpoint_by_default(None, &[], "https://api.anthropic.com")] - #[case::explicit_api_base_beats_env( - Some("https://explicit.example.com"), - BOTH_BASE_ENVS, - "https://explicit.example.com" - )] - #[case::explicit_api_base_is_trimmed( - Some(" https://explicit.example.com "), - &[], - "https://explicit.example.com" - )] - #[case::blank_api_base_falls_back_to_env( - Some(" "), - BOTH_BASE_ENVS, - "https://api-base.example.com" - )] - #[case::api_base_env_beats_base_url_env(None, BOTH_BASE_ENVS, "https://api-base.example.com")] - #[case::base_url_env_without_api_base_env( - None, - &[(ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], - "https://base-url.example.com" - )] - #[case::blank_api_base_env_falls_back_to_base_url_env( - None, - &[(ANTHROPIC_API_BASE_ENV, " \t "), (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], - "https://base-url.example.com" - )] - #[case::blank_envs_fall_back_to_public_endpoint( - None, - &[(ANTHROPIC_API_BASE_ENV, ""), (ANTHROPIC_BASE_URL_ENV, " ")], - "https://api.anthropic.com" - )] - fn api_base_resolution( - #[case] api_base: Option<&str>, - #[case] vars: Env, - #[case] expected: &str, - ) { - assert_eq!(resolve_anthropic_api_base(api_base, &env(vars)), expected); - } - #[rstest] #[case::public_endpoint(None, &[], "https://api.anthropic.com/v1/messages")] #[case::base_url_env( @@ -794,77 +815,274 @@ mod tests { ); } - #[rstest] - #[case::param_beats_env(Some("sk-param"), API_KEY_ENV, Ok("sk-param"))] - #[case::param_is_trimmed(Some(" sk-param "), &[], Ok("sk-param"))] - #[case::blank_param_falls_back_to_env(Some(" "), API_KEY_ENV, Ok("sk-env"))] - #[case::env_without_param(None, API_KEY_ENV, Ok("sk-env"))] - #[case::blank_env_is_missing(None, &[(ANTHROPIC_API_KEY_ENV, " ")], Err(MISSING_API_KEY))] - #[case::nothing_is_missing(None, &[], Err(MISSING_API_KEY))] - fn api_key_resolution( - #[case] api_key: Option<&str>, - #[case] vars: Env, - #[case] expected: Result<&str, &str>, - ) { - assert_eq!( - resolve_anthropic_api_key(api_key, &env(vars)).map_err(|error| error.to_string()), - expected.map(str::to_string).map_err(str::to_string) - ); + fn betas(values: &[&str]) -> BetaSet { + values.join(",").parse().unwrap() } - #[test] - fn config_reports_a_missing_key_as_an_auth_error() { + fn validated( + forwarded: &[(&str, &str)], + api_key: Option<&str>, + vars: Env, + ) -> Result { + ANTHROPIC_MESSAGES_CONFIG.validate_environment( + headers(forwarded), + api_key, + "claude", + &env(vars), + ) + } + + fn credential(auth: &AuthScheme) -> Option<(&'static str, &str)> { + match auth { + AuthScheme::Credential { placement, secret } => { + Some((placement.header_name(), secret.expose())) + } + AuthScheme::Forwarded => None, + other => panic!("unexpected auth scheme {other:?}"), + } + } + + #[rstest] + #[case::forwarded_oauth_bearer( + &[("anthropic-version", "2023-06-01"), ("X-Api-Key", "sk-caller"), ("Authorization", OAUTH_BEARER)], + Some("sk-deployment"), + &[("ANTHROPIC_API_KEY", "sk-env")], + &[("anthropic-version", "2023-06-01"), ("anthropic-beta", OAUTH_BETA), BROWSER_ACCESS], + Some(("Authorization", OAUTH_TOKEN)), + )] + #[case::oauth_api_key( + &[("x-api-key", OAUTH_TOKEN), ("anthropic-beta", "web-search-2025-03-05")], + Some(OAUTH_TOKEN), + &[], + &[("anthropic-beta", "oauth-2025-04-20,web-search-2025-03-05"), BROWSER_ACCESS], + Some(("Authorization", OAUTH_TOKEN)), + )] + #[case::forwarded_x_api_key_is_kept_over_the_deployment_key( + &[("X-API-KEY", "caller-key")], + Some("sk-other"), + &[("ANTHROPIC_API_KEY", "sk-env")], + &[("X-API-KEY", "caller-key")], + None, + )] + #[case::forwarded_non_oauth_bearer_is_kept( + &[("Authorization", "Bearer some-proxy-token")], + Some("sk-ant-api03-regular"), + &[], + &[("Authorization", "Bearer some-proxy-token")], + None, + )] + #[case::oauth_token_without_the_bearer_scheme_is_kept( + &[("authorization", OAUTH_TOKEN)], + None, + &[], + &[("authorization", OAUTH_TOKEN)], + None, + )] + #[case::api_key_param( + &[("anthropic-beta", "web-search-2025-03-05")], + Some("sk-param"), + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + &[("anthropic-beta", "web-search-2025-03-05")], + Some(("x-api-key", "sk-param")), + )] + #[case::env_key_when_the_param_is_blank( + &[], + Some(" "), + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + &[], + Some(("x-api-key", "sk-env")), + )] + #[case::auth_token_when_no_key_is_set( + &[], + None, + &[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + &[], + Some(("Authorization", "env-token")), + )] + #[case::oauth_env_key_as_a_bearer( + &[], + None, + &[("ANTHROPIC_API_KEY", "sk-ant-oat01-env")], + &[], + Some(("Authorization", "sk-ant-oat01-env")), + )] + fn validate_environment_shapes_the_headers_and_names_the_credential( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + #[case] vars: Env, + #[case] expected_headers: &[(&str, &str)], + #[case] expected_credential: Option<(&str, &str)>, + ) { + let environment = validated(forwarded, api_key, vars).unwrap(); + assert_eq!(environment.headers, headers(expected_headers)); + assert_eq!(credential(&environment.auth), expected_credential); + } + + #[rstest] + #[case::no_credentials(&[], None, &[])] + #[case::empty_api_key(&[], Some(""), &[])] + #[case::whitespace_only_env_values(&[], None, &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")])] + #[case::unrelated_forwarded_headers(&[("anthropic-beta", "web-search-2025-03-05")], None, &[])] + fn missing_credentials_are_an_auth_error( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + #[case] vars: Env, + ) { assert!(matches!( - ANTHROPIC_MESSAGES_CONFIG.validate_environment(vec![], None, "claude", &no_env), + validated(forwarded, api_key, vars), Err(Error::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", - environment_variable: ANTHROPIC_API_KEY_ENV, + environment_variable: "ANTHROPIC_API_KEY", })) )); } - #[test] - fn config_authenticates_with_the_anthropic_auth_token() { - let validated = ANTHROPIC_MESSAGES_CONFIG - .validate_environment( - vec![], - None, - "claude", - &env(&[("ANTHROPIC_AUTH_TOKEN", "auth-token")]), - ) - .unwrap(); - assert!(matches!( - validated.auth, - AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - ref secret - } if secret.expose() == "auth-token" - )); - } - - #[test] - fn config_requests_the_betas_the_request_features_need() { - assert_eq!( - ANTHROPIC_MESSAGES_CONFIG.request_headers( - headers(&[("x-api-key", "sk")]), - &request(json!({"speed": "fast"})) - ), - headers(&[ - ("x-api-key", "sk"), - ("anthropic-beta", beta::FAST_MODE_2026_02_01) - ]) - ); + #[rstest] + #[case::no_features(json!({}), &[])] + #[case::output_format(json!({"output_format": {"type": "json_schema"}}), &["structured-outputs-2025-11-13"])] + #[case::null_output_format(json!({"output_format": null}), &[])] + #[case::output_config_format( + json!({"output_config": {"format": {"type": "json_schema"}, "effort": "xhigh"}}), + &["structured-outputs-2025-11-13"] + )] + #[case::null_output_config_format(json!({"output_config": {"format": null}}), &[])] + #[case::top_level_output_config_without_format(json!({"output_config": {"effort": "high"}}), &[])] + #[case::fast_speed(json!({"speed": "fast"}), &["fast-mode-2026-02-01"])] + #[case::standard_speed(json!({"speed": "standard"}), &[])] + #[case::unknown_speed(json!({"speed": "turbo"}), &[])] + #[case::compaction_param(json!({"compaction": {"enabled": true}}), &["compact-2026-09-04"])] + #[case::empty_compaction_param(json!({"compaction": {}}), &["compact-2026-09-04"])] + #[case::signed_compaction_block_in_history( + json!({"messages": [ + {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]}, + {"role": "user", "content": "Continue"}, + ]}), + &["compact-2026-09-04"] + )] + #[case::unsigned_compaction_block_in_history( + json!({"messages": [ + {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": ""}]}, + {"role": "user", "content": "Continue"}, + ]}), + &[] + )] + #[case::advisor_tool( + json!({"tools": [{"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-6"}]}), + &["advisor-tool-2026-03-01"] + )] + #[case::no_tools(json!({"tools": []}), &[])] + #[case::regex_tool_search( + json!({"tools": [{"type": "tool_search_tool_regex_20251119"}]}), + &["advanced-tool-use-2025-11-20"] + )] + #[case::bm25_tool_search( + json!({"tools": [{"type": "tool_search_tool_bm25_20251119"}]}), + &["advanced-tool-use-2025-11-20"] + )] + #[case::unrelated_server_tool(json!({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}), &[])] + #[case::only_compact_edits( + json!({"context_management": {"edits": [{"type": "compact_20260112"}]}}), + &["compact-2026-01-12"] + )] + #[case::only_other_edits( + json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919", "keep": {"type": "tool_uses", "value": 3}}]}}), + &["context-management-2025-06-27"] + )] + #[case::compact_and_other_edits( + json!({"context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_tool_uses_20250919"}]}}), + &["compact-2026-01-12", "context-management-2025-06-27"] + )] + #[case::edit_without_a_type(json!({"context_management": {"edits": [{}]}}), &["context-management-2025-06-27"])] + #[case::unknown_edit_type(json!({"context_management": {"edits": [{"type": "future"}]}}), &["context-management-2025-06-27"])] + #[case::empty_edits(json!({"context_management": {"edits": []}}), &[])] + #[case::context_management_without_edits(json!({"context_management": {}}), &[])] + #[case::unmapped_openai_context_management(json!({"context_management": [{"type": "other"}]}), &[])] + #[case::per_message_output_config( + json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), + &["per-turn-control-2026-07-01"] + )] + #[case::per_message_null_output_config( + json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}), + &["per-turn-control-2026-07-01"] + )] + fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { + assert_eq!(feature_betas(&request(fields)), betas(expected)); } #[rstest] - #[case::absent(None, None)] - #[case::blank(Some(" \t "), None)] - #[case::padded(Some(" value "), Some("value"))] - fn non_empty_trims_and_drops_blank_values( - #[case] value: Option<&str>, - #[case] expected: Option<&str>, + #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}), &[("x-api-key", "k"), ("anthropic-version", "2023-06-01")])] + #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}), &[("Anthropic-Beta", " , "), ("x-api-key", "k")])] + #[case::feature_beta_is_appended( + &[("x-api-key", "k")], + json!({"speed": "fast"}), + &[("x-api-key", "k"), ("anthropic-beta", "fast-mode-2026-02-01")], + )] + #[case::existing_betas_are_normalized_without_features( + &[("Anthropic-Beta", "web-search-2025-03-05, interleaved-thinking-2025-05-14 ,web-search-2025-03-05"), ("x-api-key", "k")], + json!({}), + &[("x-api-key", "k"), ("anthropic-beta", "interleaved-thinking-2025-05-14,web-search-2025-03-05")], + )] + #[case::existing_advisor_beta_is_kept_without_an_advisor_tool( + &[("anthropic-beta", "advisor-tool-2026-03-01")], + json!({"tools": []}), + &[("anthropic-beta", "advisor-tool-2026-03-01")], + )] + #[case::feature_already_sent_is_not_duplicated( + &[("anthropic-beta", "fast-mode-2026-02-01")], + json!({"speed": "fast"}), + &[("anthropic-beta", "fast-mode-2026-02-01")], + )] + #[case::differently_cased_beta_header_is_replaced_by_one_sorted_header( + &[("Anthropic-Beta", "interleaved-thinking-2025-05-14")], + json!({"messages": [{"role": "system", "content": "env", "output_config": {"effort": "low"}}]}), + &[("anthropic-beta", "interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")], + )] + #[case::every_beta_header_casing_is_unioned_into_one_header( + &[("anthropic-beta", "interleaved-thinking-2025-05-14"), ("Anthropic-Beta", "web-search-2025-03-05")], + json!({"speed": "fast"}), + &[("anthropic-beta", "fast-mode-2026-02-01,interleaved-thinking-2025-05-14,web-search-2025-03-05")], + )] + #[case::unknown_client_betas_survive_alongside_the_added_one( + &[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14,context-management-2025-06-27,per-turn-control-2026-07-01,effort-2025-11-24")], + json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), + &[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")], + )] + fn request_headers_merge_the_feature_betas( + #[case] input: &[(&str, &str)], + #[case] fields: Value, + #[case] expected: &[(&str, &str)], ) { - assert_eq!(non_empty(value), expected); + assert_eq!( + ANTHROPIC_MESSAGES_CONFIG.request_headers(headers(input), &request(fields)), + headers(expected) + ); + } + + #[test] + fn every_feature_merges_with_the_oauth_beta_sorted() { + let environment = validated(&[], Some(OAUTH_TOKEN), &[]).unwrap(); + let all_features = request(json!({ + "compaction": {"enabled": true}, + "output_format": {"type": "json_schema"}, + "speed": "fast", + "tools": [{"type": "advisor_20260301"}, {"type": "tool_search_tool_bm25_20251119"}], + "context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_thinking_20251015"}]}, + "messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}], + })); + assert_eq!( + ANTHROPIC_MESSAGES_CONFIG.request_headers(environment.headers, &all_features), + headers(&[ + BROWSER_ACCESS, + ( + "anthropic-beta", + "advanced-tool-use-2025-11-20,advisor-tool-2026-03-01,compact-2026-01-12,compact-2026-09-04,context-management-2025-06-27,fast-mode-2026-02-01,oauth-2025-04-20,per-turn-control-2026-07-01,structured-outputs-2025-11-13" + ), + ]) + ); + assert_eq!( + credential(&environment.auth), + Some(("Authorization", OAUTH_TOKEN)) + ); } #[test] diff --git a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs index 8c768a6a66d..0c1434afb5d 100644 --- a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs @@ -1,4 +1,4 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::SecretValue; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_types::llms::anthropic_messages::{ anthropic_request::{ @@ -10,8 +10,9 @@ use litellm_types::llms::anthropic_messages::{ use crate::{ Error, - anthropic::messages::transformation::{ - ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, + anthropic::{ + common_utils::{API_KEY_PLACEMENT, MESSAGES_PATH_SUFFIX, non_empty}, + messages::transformation::{ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig}, }, base_llm::{ anthropic_messages::transformation::{ @@ -24,9 +25,7 @@ use crate::{ const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY"; const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic"; -const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; const SYSTEM_ROLE: &str = "system"; -const API_KEY_HEADER: &str = "x-api-key"; pub struct AzureAnthropicMessagesConfig { anthropic: AnthropicMessagesConfig, @@ -86,14 +85,14 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { _model: &str, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - if has_header(&headers, API_KEY_HEADER) || has_bearer_auth(&headers) { + if has_header(&headers, API_KEY_PLACEMENT.header_name()) || has_bearer_auth(&headers) { return Ok(ValidatedEnvironment { headers, auth: AuthScheme::Forwarded, }); } let auth = AuthScheme::Credential { - placement: CredentialPlacement::Header(API_KEY_HEADER), + placement: API_KEY_PLACEMENT, secret: SecretValue::new(resolve_azure_api_key(api_key, env_lookup)?), }; Ok(ValidatedEnvironment { headers, auth }) @@ -216,6 +215,8 @@ mod tests { use rstest::rstest; use serde_json::json; + use litellm_auth::CredentialPlacement; + use super::*; use crate::anthropic::common_utils::AnthropicModelCapabilities; diff --git a/litellm-rust/crates/llms/src/base_llm/auth.rs b/litellm-rust/crates/llms/src/base_llm/auth.rs index 897b0f62e82..7d1b0014dcd 100644 --- a/litellm-rust/crates/llms/src/base_llm/auth.rs +++ b/litellm-rust/crates/llms/src/base_llm/auth.rs @@ -7,6 +7,7 @@ use litellm_auth::{AuthServices, CredentialPlacement, SecretValue, TokenProviderHandle}; use litellm_auth_aws::{AwsCredentialSource, SigV4Signer}; +use litellm_http::request::without_headers; pub type Headers = Vec<(String, String)>; @@ -107,9 +108,8 @@ fn with_credential(headers: Headers, placement: CredentialPlacement, credential: CredentialPlacement::Bearer => format!("Bearer {credential}"), CredentialPlacement::Header(_) => credential.to_string(), }; - headers + without_headers(headers, &[name]) .into_iter() - .filter(|(header, _)| !header.eq_ignore_ascii_case(name)) .chain([(name.to_ascii_lowercase(), value)]) .collect() } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index bb0ec6671e3..cf634ab3fc3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -2,9 +2,8 @@ use std::convert::Infallible; use bytes::Bytes; use litellm_core::messages::{ - Error, - route::{Messages, MessagesCall, MessagesOutput, MessagesStreamHead, messages_body}, - types::MessagesShaping, + Error, MessagesCall, MessagesShaping, messages_body, + route::{Messages, MessagesOutput, MessagesStreamHead}, }; use litellm_host_python::{InvokeError, ProtocolHost, from_py, lookup, to_py}; use litellm_http::transport::Error as TransportError; @@ -76,6 +75,11 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult { error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; Ok(error) } + Error::MissingField(field) => { + let error = PyValueError::new_err(format!("missing required field: {field}")); + error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; + Ok(error) + } other => Ok(route_error_to_pyerr(other)), } } @@ -314,6 +318,7 @@ mod tests { #[rstest] #[case::rejected_request(Error::InvalidRequest("does not support top_k=5".into()), true)] + #[case::missing_field(Error::MissingField("max_tokens"), true)] #[case::unresolvable_provider(Error::InvalidProvider("openai".into()), false)] #[case::upstream_failure( Error::Transport(TransportError::Http { status: 400, body: "bad".into() }), diff --git a/litellm-rust/crates/router/src/deployment.rs b/litellm-rust/crates/router/src/deployment.rs index 4904bf4ffdd..9234728db28 100644 --- a/litellm-rust/crates/router/src/deployment.rs +++ b/litellm-rust/crates/router/src/deployment.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_core::messages::types::MessagesShaping; +use litellm_core::messages::MessagesShaping; #[derive(Clone, Debug, Default)] pub struct Deployment { diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index 3cbd316cab0..2d16f102de3 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -1,7 +1,7 @@ use std::time::Duration; use litellm_config::Config; -use litellm_core::messages::types::MessagesShaping; +use litellm_core::messages::MessagesShaping; use litellm_router::{Deployment, Router}; use rstest::rstest; diff --git a/litellm-rust/crates/secrets/src/native.rs b/litellm-rust/crates/secrets/src/native.rs index f1dc7ccb732..8edcebe7b12 100644 --- a/litellm-rust/crates/secrets/src/native.rs +++ b/litellm-rust/crates/secrets/src/native.rs @@ -6,8 +6,8 @@ use litellm_http::{HttpClientConfig, HttpClientPool}; use crate::{Error, KeyManagementSettings, KeyManagementSystem, SecretManager}; pub async fn load_native_manager( - pool: &HttpClientPool, - config: &HttpClientConfig, + _pool: &HttpClientPool, + _config: &HttpClientConfig, system: KeyManagementSystem, settings: KeyManagementSettings, environment: Arc, @@ -33,14 +33,14 @@ pub async fn load_native_manager( #[cfg(feature = "azure")] (KeyManagementSystem::AzureKeyVault, _, environment, _) => Ok( SecretManager::AzureKeyVault(crate::azure::AzureKeyVault::new( - pool.client(config, litellm_http::ClientVariant::Provider)?, + _pool.client(_config, litellm_http::ClientVariant::Provider)?, environment, )?), ), #[cfg(feature = "google")] (KeyManagementSystem::GoogleSecretManager, _, environment, enterprise_enabled) => Ok( SecretManager::GoogleSecretManager(crate::google::GoogleSecretManager::new( - pool.client(config, litellm_http::ClientVariant::Provider)?, + _pool.client(_config, litellm_http::ClientVariant::Provider)?, environment, enterprise_enabled, )?), @@ -61,8 +61,8 @@ pub async fn load_native_manager( #[cfg(feature = "cyberark")] (KeyManagementSystem::Cyberark, _, environment, enterprise_enabled) => Ok( SecretManager::Cyberark(crate::cyberark::CyberArkSecretManager::new( - pool, - config, + _pool, + _config, environment, enterprise_enabled, )?), diff --git a/litellm-rust/crates/tracing/Cargo.toml b/litellm-rust/crates/tracing/Cargo.toml index 41ad20afb3e..914d988d301 100644 --- a/litellm-rust/crates/tracing/Cargo.toml +++ b/litellm-rust/crates/tracing/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +base64.workspace = true fancy-regex.workspace = true percent-encoding.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/tracing/src/lib.rs b/litellm-rust/crates/tracing/src/lib.rs index 47f97d6db27..4c6ec104f3a 100644 --- a/litellm-rust/crates/tracing/src/lib.rs +++ b/litellm-rust/crates/tracing/src/lib.rs @@ -5,6 +5,7 @@ use std::{ pin::pin, }; +use base64::{Engine, engine::general_purpose::STANDARD}; use serde_json::{Map, Value}; use tracing::{ Dispatch, Event, Subscriber, @@ -20,6 +21,31 @@ pub use processing::{DiagnosticInput, DiagnosticOutput, Policy, Processor}; pub use redaction::{REDACTED, SecretRedactor}; pub use tracing::{Level, Metadata, debug, error, info, trace, warn}; +pub struct ByteChunk<'a>(&'a [u8]); + +impl<'a> ByteChunk<'a> { + pub fn new(data: &'a [u8]) -> Self { + Self(data) + } + + pub fn encoding(&self) -> &'static str { + if std::str::from_utf8(self.0).is_ok() { + "utf8" + } else { + "base64" + } + } +} + +impl fmt::Display for ByteChunk<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match std::str::from_utf8(self.0) { + Ok(text) => formatter.write_str(text), + Err(_) => formatter.write_str(&STANDARD.encode(self.0)), + } + } +} + pub trait Sink: Send + Sync + 'static { fn enabled(&self, metadata: &Metadata<'_>) -> bool; fn emit(&self, record: &Record); @@ -44,6 +70,10 @@ impl Logger { } } + pub fn install_global(&self) -> Result<(), tracing::dispatcher::SetGlobalDefaultError> { + tracing::dispatcher::set_global_default(self.dispatch.clone()) + } + pub fn scope(&self, operation: impl FnOnce() -> T) -> T { if EMITTING.get() { return operation(); diff --git a/litellm-rust/crates/tracing/tests/logging.rs b/litellm-rust/crates/tracing/tests/logging.rs index 585e442dad1..3387f259f8b 100644 --- a/litellm-rust/crates/tracing/tests/logging.rs +++ b/litellm-rust/crates/tracing/tests/logging.rs @@ -4,7 +4,9 @@ use std::sync::{ mpsc, }; -use litellm_tracing::{Level, Logger, Metadata, Record, Sink, info, warn}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_tracing::{ByteChunk, Level, Logger, Metadata, Record, Sink, info, warn}; +use rstest::rstest; use serde_json::{Value, json}; struct Output { @@ -120,3 +122,18 @@ fn nested_scopes_restore_the_previous_sink() { ["inside"] ); } + +#[rstest] +#[case::utf8(b"event: message_stop\n\n", "utf8")] +#[case::binary(&[0xff, 0x00, 0x80], "base64")] +fn byte_chunk_logging_preserves_exact_bytes(#[case] bytes: &[u8], #[case] encoding: &str) { + let chunk = ByteChunk::new(bytes); + assert_eq!(chunk.encoding(), encoding); + let text = chunk.to_string(); + let recovered = match encoding { + "utf8" => text.into_bytes(), + "base64" => STANDARD.decode(text).unwrap(), + _ => unreachable!(), + }; + assert_eq!(recovered, bytes); +} diff --git a/litellm-rust/crates/types/src/llms/anthropic.rs b/litellm-rust/crates/types/src/llms/anthropic.rs new file mode 100644 index 00000000000..3e0c4369b11 --- /dev/null +++ b/litellm-rust/crates/types/src/llms/anthropic.rs @@ -0,0 +1,240 @@ +use std::{ + cmp::Ordering, + collections::BTreeSet, + convert::Infallible, + fmt, + hash::{Hash, Hasher}, + str::FromStr, +}; + +/// One value of the `anthropic-beta` header. Equality, ordering and hashing follow the wire +/// string, so a value parsed from a caller's header never disagrees with the matching variant. +#[derive(Clone, Debug, strum::AsRefStr, strum::Display, strum::EnumString)] +pub enum AnthropicBeta { + #[strum(serialize = "oauth-2025-04-20")] + Oauth20250420, + #[strum(serialize = "web-fetch-2025-09-10")] + WebFetch20250910, + #[strum(serialize = "web-search-2025-03-05")] + WebSearch20250305, + #[strum(serialize = "context-management-2025-06-27")] + ContextManagement20250627, + #[strum(serialize = "compact-2026-01-12")] + Compact20260112, + #[strum(serialize = "compact-2026-09-04")] + Compact20260904, + #[strum(serialize = "structured-outputs-2025-11-13")] + StructuredOutputs20251113, + #[strum(serialize = "advanced-tool-use-2025-11-20")] + AdvancedToolUse20251120, + #[strum(serialize = "fast-mode-2026-02-01")] + FastMode20260201, + #[strum(serialize = "advisor-tool-2026-03-01")] + AdvisorTool20260301, + #[strum(serialize = "per-turn-control-2026-07-01")] + PerTurnControl20260701, + #[strum(serialize = "dangerous-tool-use-2026-09-03")] + DangerousToolUse20260903, + #[strum(default, transparent)] + Other(String), +} + +impl AnthropicBeta { + pub const KNOWN: [Self; 12] = [ + Self::Oauth20250420, + Self::WebFetch20250910, + Self::WebSearch20250305, + Self::ContextManagement20250627, + Self::Compact20260112, + Self::Compact20260904, + Self::StructuredOutputs20251113, + Self::AdvancedToolUse20251120, + Self::FastMode20260201, + Self::AdvisorTool20260301, + Self::PerTurnControl20260701, + Self::DangerousToolUse20260903, + ]; + + pub fn as_str(&self) -> &str { + self.as_ref() + } +} + +impl PartialEq for AnthropicBeta { + fn eq(&self, other: &Self) -> bool { + self.as_str() == other.as_str() + } +} + +impl Eq for AnthropicBeta {} + +impl Hash for AnthropicBeta { + fn hash(&self, state: &mut H) { + self.as_str().hash(state); + } +} + +impl PartialOrd for AnthropicBeta { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for AnthropicBeta { + fn cmp(&self, other: &Self) -> Ordering { + self.as_str().cmp(other.as_str()) + } +} + +/// The values of one `anthropic-beta` header: sorted, deduplicated, comma-joined on the wire. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct BetaSet(BTreeSet); + +impl BetaSet { + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + pub fn contains(&self, beta: &AnthropicBeta) -> bool { + self.0.contains(beta) + } + + pub fn iter(&self) -> impl Iterator { + self.0.iter() + } + + pub fn union(self, other: Self) -> Self { + self.0.into_iter().chain(other.0).collect() + } +} + +impl FromIterator for BetaSet { + fn from_iter>(betas: I) -> Self { + Self(betas.into_iter().collect()) + } +} + +impl IntoIterator for BetaSet { + type Item = AnthropicBeta; + type IntoIter = std::collections::btree_set::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl FromStr for BetaSet { + type Err = Infallible; + + fn from_str(header: &str) -> Result { + Ok(header + .split(',') + .map(str::trim) + .filter(|piece| !piece.is_empty()) + .map(|piece| AnthropicBeta::from_str(piece).unwrap_or_else(|never| match never {})) + .collect()) + } +} + +impl fmt::Display for BetaSet { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut betas = self.0.iter(); + let Some(first) = betas.next() else { + return Ok(()); + }; + f.write_str(first.as_str())?; + betas.try_for_each(|beta| write!(f, ",{beta}")) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn set(header: &str) -> BetaSet { + header.parse().unwrap_or_else(|never| match never {}) + } + + #[rstest] + fn every_known_beta_parses_back_to_itself( + #[values( + AnthropicBeta::Oauth20250420, + AnthropicBeta::WebFetch20250910, + AnthropicBeta::WebSearch20250305, + AnthropicBeta::ContextManagement20250627, + AnthropicBeta::Compact20260112, + AnthropicBeta::Compact20260904, + AnthropicBeta::StructuredOutputs20251113, + AnthropicBeta::AdvancedToolUse20251120, + AnthropicBeta::FastMode20260201, + AnthropicBeta::AdvisorTool20260301, + AnthropicBeta::PerTurnControl20260701, + AnthropicBeta::DangerousToolUse20260903 + )] + beta: AnthropicBeta, + ) { + let parsed: AnthropicBeta = beta.as_str().parse().unwrap(); + assert!(!matches!(parsed, AnthropicBeta::Other(_))); + assert_eq!(parsed, beta); + assert!(AnthropicBeta::KNOWN.contains(&beta)); + } + + #[test] + fn unknown_values_are_kept_verbatim() { + let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap(); + assert_eq!( + parsed, + AnthropicBeta::Other("claude-code-20250219".to_string()) + ); + assert_eq!(parsed.to_string(), "claude-code-20250219"); + } + + #[test] + fn a_known_value_spelled_as_other_is_the_same_beta() { + let spelled_out = AnthropicBeta::Other("compact-2026-01-12".to_string()); + assert_eq!(spelled_out, AnthropicBeta::Compact20260112); + assert_eq!( + spelled_out.cmp(&AnthropicBeta::Compact20260112), + Ordering::Equal + ); + assert_eq!( + BetaSet::from_iter([spelled_out, AnthropicBeta::Compact20260112]).to_string(), + "compact-2026-01-12" + ); + } + + #[rstest] + #[case::empty("", "")] + #[case::blank_pieces(" , ,", "")] + #[case::single("b", "b")] + #[case::sorted("c,a", "a,c")] + #[case::trimmed_and_deduplicated("b, a ,b", "a,b")] + #[case::blank_pieces_skipped("a,,b", "a,b")] + #[case::known_and_unknown_sort_together( + "web-search-2025-03-05,claude-code-20250219,fast-mode-2026-02-01", + "claude-code-20250219,fast-mode-2026-02-01,web-search-2025-03-05" + )] + fn header_values_round_trip_sorted_and_deduplicated(#[case] header: &str, #[case] wire: &str) { + assert_eq!(set(header).to_string(), wire); + assert_eq!(set(header).is_empty(), wire.is_empty()); + } + + #[rstest] + #[case::disjoint("a,c", "b", "a,b,c")] + #[case::overlapping("a,b", "b,c", "a,b,c")] + #[case::empty_right("a", "", "a")] + #[case::empty_left("", "a", "a")] + fn union_merges_both_sides(#[case] left: &str, #[case] right: &str, #[case] wire: &str) { + assert_eq!(set(left).union(set(right)).to_string(), wire); + } + + #[test] + fn contains_matches_by_wire_value() { + let betas = set("oauth-2025-04-20,claude-code-20250219"); + assert!(betas.contains(&AnthropicBeta::Oauth20250420)); + assert!(betas.contains(&AnthropicBeta::Other("claude-code-20250219".into()))); + assert!(!betas.contains(&AnthropicBeta::FastMode20260201)); + } +} diff --git a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs index 342e891a1e3..335a03c4b5b 100644 --- a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs +++ b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs @@ -111,6 +111,70 @@ impl From for ReasoningEffort { } } +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +#[strum(serialize_all = "lowercase")] +pub enum Speed { + Fast, + Standard, +} + +impl Speed { + pub fn as_str(self) -> &'static str { + self.into() + } +} + +/// The tools whose presence changes how the request is sent. Every other tool, custom or +/// server, deserializes as `Recognized::Unrecognized` and passes through verbatim. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum AnthropicTool { + #[serde(rename = "advisor_20260301")] + Advisor { + #[serde(flatten)] + extra: Map, + }, + #[serde(rename = "tool_search_tool_regex_20251119")] + ToolSearchRegex { + #[serde(flatten)] + extra: Map, + }, + #[serde(rename = "tool_search_tool_bm25_20251119")] + ToolSearchBm25 { + #[serde(flatten)] + extra: Map, + }, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum ContextEdit { + #[serde(rename = "compact_20260112")] + Compact { + #[serde(flatten)] + extra: Map, + }, + #[serde(rename = "clear_tool_uses_20250919")] + ClearToolUses { + #[serde(flatten)] + extra: Map, + }, + #[serde(rename = "clear_thinking_20251015")] + ClearThinking { + #[serde(flatten)] + extra: Map, + }, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct ContextManagement { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub edits: Option>>, + #[serde(flatten)] + pub extra: Map, +} + #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct OutputConfig { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -210,7 +274,7 @@ pub struct AnthropicMessagesOptionalParams { #[serde(skip_serializing_if = "Option::is_none")] pub top_k: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub tools: Option>, + pub tools: Option>>, #[serde(skip_serializing_if = "Option::is_none")] pub tool_choice: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -222,13 +286,13 @@ pub struct AnthropicMessagesOptionalParams { #[serde(skip_serializing_if = "Option::is_none")] pub mcp_servers: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub context_management: Option, + pub context_management: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub output_format: Option, #[serde(skip_serializing_if = "Option::is_none")] pub output_config: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub speed: Option, + pub speed: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub inference_geo: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -397,6 +461,33 @@ mod tests { "thinking": {"type": "future", "budget_tokens": 1}, "output_config": "bogus" }))] + #[case::tools_speed_and_context_management(json!({ + "model": "m", + "messages": [], + "speed": "fast", + "tools": [ + {"name": "get_weather", "input_schema": {"type": "object"}}, + {"type": "custom", "name": "f", "input_schema": {}}, + {"type": "web_search_20250305", "name": "web_search", "max_uses": 3}, + {"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-6"}, + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + {"type": "tool_search_tool_bm25_20251119"} + ], + "context_management": {"edits": [ + {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}, + {"type": "clear_tool_uses_20250919", "keep": {"type": "tool_uses", "value": 3}}, + {"type": "clear_thinking_20251015"}, + {"type": "future_edit"}, + {} + ], "future": true} + }))] + #[case::unrecognized_tools_speed_and_context_management_are_kept_verbatim(json!({ + "model": "m", + "messages": [], + "speed": "turbo", + "tools": ["none", 5], + "context_management": [{"type": "compaction", "compact_threshold": 5}] + }))] fn request_round_trips_unchanged(#[case] request: Value) { assert_eq!(round_trip::(&request), request); } @@ -444,6 +535,114 @@ mod tests { ); } + #[rstest] + #[case::advisor( + json!({"type": "advisor_20260301", "name": "advisor"}), + Recognized::Known(AnthropicTool::Advisor { extra: Map::from_iter([("name".to_string(), json!("advisor"))]) }) + )] + #[case::regex_tool_search( + json!({"type": "tool_search_tool_regex_20251119"}), + Recognized::Known(AnthropicTool::ToolSearchRegex { extra: Map::new() }) + )] + #[case::bm25_tool_search( + json!({"type": "tool_search_tool_bm25_20251119"}), + Recognized::Known(AnthropicTool::ToolSearchBm25 { extra: Map::new() }) + )] + #[case::custom_tool_without_a_type( + json!({"name": "advisor", "input_schema": {}}), + Recognized::Unrecognized(json!({"name": "advisor", "input_schema": {}})) + )] + #[case::other_server_tool( + json!({"type": "web_search_20250305", "name": "web_search"}), + Recognized::Unrecognized(json!({"type": "web_search_20250305", "name": "web_search"})) + )] + #[case::not_an_object(json!("advisor_20260301"), Recognized::Unrecognized(json!("advisor_20260301")))] + fn tools_are_recognized_by_their_exact_type( + #[case] tool: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(tool).unwrap(), + expected + ); + } + + #[rstest] + #[case::compact( + json!({"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1}}), + Recognized::Known(ContextEdit::Compact { + extra: Map::from_iter([("trigger".to_string(), json!({"type": "input_tokens", "value": 1}))]), + }) + )] + #[case::clear_tool_uses( + json!({"type": "clear_tool_uses_20250919"}), + Recognized::Known(ContextEdit::ClearToolUses { extra: Map::new() }) + )] + #[case::clear_thinking( + json!({"type": "clear_thinking_20251015"}), + Recognized::Known(ContextEdit::ClearThinking { extra: Map::new() }) + )] + #[case::unknown_type(json!({"type": "future"}), Recognized::Unrecognized(json!({"type": "future"})))] + #[case::no_type(json!({}), Recognized::Unrecognized(json!({})))] + fn context_edits_are_recognized_by_their_exact_type( + #[case] edit: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(edit).unwrap(), + expected + ); + } + + #[rstest] + #[case::edits( + json!({"edits": [{"type": "compact_20260112"}]}), + Recognized::Known(ContextManagement { + edits: Some(vec![Recognized::Known(ContextEdit::Compact { extra: Map::new() })]), + extra: Map::new(), + }) + )] + #[case::object_without_edits( + json!({"future": 1}), + Recognized::Known(ContextManagement { + edits: None, + extra: Map::from_iter([("future".to_string(), json!(1))]), + }) + )] + #[case::openai_list(json!([{"type": "compaction"}]), Recognized::Unrecognized(json!([{"type": "compaction"}])))] + #[case::edits_not_a_list(json!({"edits": 5}), Recognized::Unrecognized(json!({"edits": 5})))] + #[case::scalar(json!("compaction"), Recognized::Unrecognized(json!("compaction")))] + fn context_management_is_known_only_as_an_edits_object( + #[case] value: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(value).unwrap(), + expected + ); + } + + #[rstest] + #[case::fast(json!("fast"), Recognized::Known(Speed::Fast))] + #[case::standard(json!("standard"), Recognized::Known(Speed::Standard))] + #[case::unknown(json!("turbo"), Recognized::Unrecognized(json!("turbo")))] + #[case::wrong_case(json!("Fast"), Recognized::Unrecognized(json!("Fast")))] + #[case::not_a_string(json!(1), Recognized::Unrecognized(json!(1)))] + fn speed_is_known_only_as_a_documented_value( + #[case] value: Value, + #[case] expected: Recognized, + ) { + assert_eq!( + serde_json::from_value::>(value).unwrap(), + expected + ); + } + + #[rstest] + fn speed_names_match_the_wire(#[values(Speed::Fast, Speed::Standard)] speed: Speed) { + assert_eq!(serde_json::to_value(speed).unwrap(), json!(speed.as_str())); + } + #[rstest] fn effort_level_names_match_the_wire( #[values( diff --git a/litellm-rust/crates/types/src/llms/mod.rs b/litellm-rust/crates/types/src/llms/mod.rs index 09d2207a0ca..19ce0bb77ef 100644 --- a/litellm-rust/crates/types/src/llms/mod.rs +++ b/litellm-rust/crates/types/src/llms/mod.rs @@ -1,2 +1,3 @@ +pub mod anthropic; pub mod anthropic_messages; pub mod openai; diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 6e455817194..46acb94c958 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -109,7 +109,7 @@ RULES: Final[Rules] = ( RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY), RouteRule(Route.OCR, Rollout.RUST_REQUIRED), - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, providers=frozenset({"anthropic"})), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index eeea364b674..2b3edac612e 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -50,12 +50,13 @@ def test_shipped_decisions( monkeypatch.setenv("LITELLM_RUST", environment) context: Final = RouteContext(route, provider=provider, model="test-model", delivery=delivery) - if route is Route.OCR: - assert catalog.rollout(context) is Rollout.RUST_REQUIRED - assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.TRANSCRIPTION and provider == "bedrock": + if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED + elif route is Route.MESSAGES and provider == "anthropic": + assert catalog.rollout(context) is Rollout.RUST_OPT_IN + opted_in: Final = environment == "1" or (environment is None and process is True) + assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON) else: assert catalog.rollout(context) is Rollout.PYTHON_ONLY assert catalog.decision(context) is Decision.PYTHON @@ -145,10 +146,7 @@ def test_ocr_has_no_python_path_to_opt_out_to( monkeypatch.setenv("LITELLM_RUST", environment) assert catalog.decision(RouteContext(Route.OCR, model="m")) is Decision.RUST_REQUIRED - assert ( - catalog.decision(RouteContext(Route.OCR, provider="aws_textract", model="m")) - is Decision.RUST_REQUIRED - ) + assert catalog.decision(RouteContext(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED @pytest.mark.parametrize(