From 2fa1d8e52b86818bdefb19bc1286c6bb49345f0a Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:46 -0700 Subject: [PATCH] feat(gateway): s3_v2 observer writing the proxy's logging payload to S3 (#47063) Adds the s3_v2 observer, turned on by listing s3_v2 (or s3_log) in litellm_settings.callbacks, success_callback or failure_callback. Each finished request becomes the proxy's StandardLoggingPayload, with the messages from the request body and the response as the client got it, streams put back together. It is queued and written every flush interval under the proxy's key layout, with its Content-MD5, Content-Disposition and encryption headers, one PUT per request or one NDJSON object per folder with s3_batch_file_upload. s3_callback_params is read as the proxy reads it, os.environ references included. Uploads are retried, dropped and kept for the next flush by the proxy's rules, and what is queued is written at shutdown. auth-aws gains an S3 signer that signs the path as sent and every header, as botocore does. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/Cargo.lock | 4 + litellm-rust/crates/auth-aws/src/aws.rs | 102 +- litellm-rust/crates/gateway/Cargo.toml | 6 +- litellm-rust/crates/gateway/src/error.rs | 2 + litellm-rust/crates/gateway/src/lib.rs | 1 + litellm-rust/crates/gateway/src/main.rs | 3 + litellm-rust/crates/gateway/src/prometheus.rs | 33 +- .../crates/gateway/src/s3_log/config.rs | 676 +++++++++ litellm-rust/crates/gateway/src/s3_log/key.rs | 389 ++++++ litellm-rust/crates/gateway/src/s3_log/mod.rs | 506 +++++++ .../crates/gateway/src/s3_log/payload.rs | 1204 +++++++++++++++++ .../crates/gateway/src/s3_log/upload.rs | 278 ++++ litellm-rust/crates/gateway/src/spend.rs | 4 +- litellm-rust/crates/gateway/tests/s3_log.rs | 313 +++++ litellm-rust/docs/customer-proxy/s3-log.svg | 61 + 15 files changed, 3563 insertions(+), 19 deletions(-) create mode 100644 litellm-rust/crates/gateway/src/s3_log/config.rs create mode 100644 litellm-rust/crates/gateway/src/s3_log/key.rs create mode 100644 litellm-rust/crates/gateway/src/s3_log/mod.rs create mode 100644 litellm-rust/crates/gateway/src/s3_log/payload.rs create mode 100644 litellm-rust/crates/gateway/src/s3_log/upload.rs create mode 100644 litellm-rust/crates/gateway/tests/s3_log.rs create mode 100644 litellm-rust/docs/customer-proxy/s3-log.svg diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 20bb696c46c..5605bcd7271 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3754,6 +3754,7 @@ dependencies = [ "futures-util", "http 1.4.2", "http-body-util", + "litellm-auth-aws", "litellm-auth-types", "litellm-config", "litellm-db", @@ -3770,11 +3771,13 @@ dependencies = [ "litellm-router", "litellm-secrets", "litellm-tracing", + "md-5", "rand 0.8.7", "redis", "rstest", "serde", "serde_json", + "sha2 0.10.9", "sqlx", "strum", "tempfile", @@ -3785,6 +3788,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "wiremock", ] [[package]] diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 8fdbba21711..5a17e182a72 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -8,7 +8,8 @@ use sha2::{Digest, Sha256}; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ - SignableBody, SignableRequest, SigningParams, SigningSettings, sign, + PercentEncodingMode, SignableBody, SignableRequest, SigningParams, SigningSettings, + UriPathNormalizationMode, sign, }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; @@ -464,6 +465,58 @@ pub fn sign_post( service: &str, credentials: &Credentials, signing_time: SystemTime, +) -> Result, Error> { + sign_request( + "POST", + url, + body, + headers, + region, + service, + credentials, + signing_time, + SigningSettings::default(), + ) +} + +/// Signs an S3 request as botocore's `S3SigV4Auth` does: the path is signed as sent, encoded +/// once and not normalized, and every header given is covered. +pub fn sign_s3( + method: &str, + url: &str, + body: &[u8], + headers: &BTreeMap, + region: &str, + credentials: &Credentials, + signing_time: SystemTime, +) -> Result, Error> { + let mut settings = SigningSettings::default(); + settings.percent_encoding_mode = PercentEncodingMode::Single; + settings.uri_path_normalization_mode = UriPathNormalizationMode::Disabled; + sign_request( + method, + url, + body, + headers, + region, + "s3", + credentials, + signing_time, + settings, + ) +} + +#[allow(clippy::too_many_arguments)] +fn sign_request( + method: &str, + url: &str, + body: &[u8], + headers: &BTreeMap, + region: &str, + service: &str, + credentials: &Credentials, + signing_time: SystemTime, + settings: SigningSettings, ) -> Result, Error> { let identity: Identity = credentials.clone().into(); let params = v4::SigningParams::builder() @@ -471,14 +524,14 @@ pub fn sign_post( .region(region) .name(service) .time(signing_time) - .settings(SigningSettings::default()) + .settings(settings) .build() .map(SigningParams::from) .map_err(|error| Error::AwsSigningParameters(error.to_string()))?; let header_refs = headers .iter() .map(|(name, value)| (name.as_str(), value.as_str())); - let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) + let request = SignableRequest::new(method, url, header_refs, SignableBody::Bytes(body)) .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; let (instructions, _) = sign(request, ¶ms) .map_err(|error| Error::AwsSigning(error.to_string()))? @@ -935,6 +988,49 @@ mod tests { ); } + /// The PUT Object example of the S3 SigV4 guide, "Signature Calculations for the + /// Authorization Header: Transferring Payload in a Single Chunk" (checked 2026-10). + #[test] + fn s3_requests_are_signed_like_the_s3_guide_example() { + let headers = BTreeMap::from([ + ( + "Date".to_string(), + "Fri, 24 May 2013 00:00:00 GMT".to_string(), + ), + ( + "x-amz-content-sha256".to_string(), + "44ce7dd67c959e0d3524ffac1771dfbba87d2b6b4b4e99e42034a8b803f8b072".to_string(), + ), + ( + "x-amz-storage-class".to_string(), + "REDUCED_REDUNDANCY".to_string(), + ), + ]); + let signed = sign_s3( + "PUT", + "https://examplebucket.s3.amazonaws.com/test%24file.text", + b"Welcome to Amazon S3.", + &headers, + "us-east-1", + &Credentials::new( + "AKIAIOSFODNN7EXAMPLE", + "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + None, + None, + "test", + ), + UNIX_EPOCH + Duration::from_secs(1_369_353_600), + ) + .unwrap(); + + assert_eq!( + signed["Authorization"], + "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20130524/us-east-1/s3/aws4_request, \ + SignedHeaders=date;host;x-amz-content-sha256;x-amz-date;x-amz-storage-class, \ + Signature=98ad721746da40c64f1a55b78f14c238d841ea1380cd77a1b5971af0ece108bd" + ); + } + #[test] fn same_role_comparison_matches_partition_account_and_role() { assert!(same_role_arns( diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index d7e2e59bb53..e1ca8682cab 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -21,7 +21,11 @@ litellm-gateway-auth.workspace = true litellm-gateway-hooks.workspace = true litellm-model-catalog.workspace = true litellm-gateway-ui.workspace = true +litellm-auth-aws.workspace = true litellm-auth-types.workspace = true +base64.workspace = true +md-5 = "0.11.0" +sha2.workspace = true litellm-config.workspace = true litellm-db.workspace = true litellm-http = { workspace = true, features = ["mcp"] } @@ -39,7 +43,7 @@ uuid.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] -base64.workspace = true +wiremock.workspace = true litellm-db-testing.workspace = true sqlx.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index fbab39ef426..742a036e69d 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -22,6 +22,8 @@ pub enum Error { Secret(#[from] litellm_secrets::Error), #[error("litellm_settings.hooks: {0}")] Hooks(String), + #[error("litellm_settings.s3_callback_params: {0}")] + S3Log(String), #[error("coordination Redis: {0}")] Redis(String), #[error("MCP requires a configured master key")] diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index e7bc29372e4..d351afb2061 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -13,6 +13,7 @@ pub mod model_capacity; pub mod prometheus; pub mod rate_limits; pub mod request_hooks; +pub mod s3_log; mod secrets; pub mod spend; pub use error::Error; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index bdf8ffef1da..71482524afc 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -144,6 +144,9 @@ async fn main() -> Result<(), Box> { if let Some(prometheus) = &prometheus { observers.push(prometheus.clone()); } + if let Some(s3_log) = litellm_gateway::s3_log::S3Log::from_config(&config, &inference).await? { + observers.push(s3_log); + } litellm_gateway::hooks::install(&inference, &config, hooks, observers).await?; let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; let address = listener.local_addr()?; diff --git a/litellm-rust/crates/gateway/src/prometheus.rs b/litellm-rust/crates/gateway/src/prometheus.rs index 1a397351fdb..28cf7ef0dcd 100644 --- a/litellm-rust/crates/gateway/src/prometheus.rs +++ b/litellm-rust/crates/gateway/src/prometheus.rs @@ -678,16 +678,31 @@ impl<'a> Values<'a> { /// reports for them; a provider error is named the way the proxy maps its status, prefixed by /// the provider. fn exception_class(record: &Record) -> String { - let status = record.outcome.status; + let class = error_class(record); let Some(deployment) = &record.outcome.deployment else { + return class.to_owned(); + }; + match deployment.provider() { + Some(provider) if !provider.is_empty() => { + let mut letters = provider.chars(); + let first = letters.next().map(|c| c.to_uppercase().collect::()); + format!("{}{}.{class}", first.unwrap_or_default(), letters.as_str()) + } + _ => class.to_owned(), + } +} + +/// The class of the exception the proxy raises for a failed request, without the provider. +pub(crate) fn error_class(record: &Record) -> &'static str { + let status = record.outcome.status; + if record.outcome.deployment.is_none() { return match status { 400 => "BudgetExceededError", 429 | 503 => "HTTPException", _ => "ProxyException", - } - .to_owned(); - }; - let class = match status { + }; + } + match status { 400 => "BadRequestError", 401 => "AuthenticationError", 403 => "PermissionDeniedError", @@ -698,14 +713,6 @@ fn exception_class(record: &Record) -> String { 500 => "InternalServerError", 503 => "ServiceUnavailableError", _ => "APIError", - }; - match deployment.provider() { - Some(provider) if !provider.is_empty() => { - let mut letters = provider.chars(); - let first = letters.next().map(|c| c.to_uppercase().collect::()); - format!("{}{}.{class}", first.unwrap_or_default(), letters.as_str()) - } - _ => class.to_owned(), } } diff --git a/litellm-rust/crates/gateway/src/s3_log/config.rs b/litellm-rust/crates/gateway/src/s3_log/config.rs new file mode 100644 index 00000000000..cbd99d9f5c2 --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/config.rs @@ -0,0 +1,676 @@ +//! `litellm_settings.s3_callback_params`, read the way the proxy's `S3Logger` reads it: values +//! may be `os.environ/` references, and a number or flag the proxy would not parse falls back to +//! its default with a warning instead of stopping startup. + +use std::time::Duration; + +use litellm_auth_aws::AwsAuthConfig; +use litellm_config::{Config, Value}; +use litellm_http::SslVerify; +use litellm_secrets::source::SecretSource; + +use super::key::Granularity; +use crate::Error; + +/// The names a config's callback lists turn this observer on with. +const NAMES: [&str; 2] = ["s3_v2", "s3_log"]; + +const DEFAULT_FLUSH_INTERVAL_SECONDS: u64 = 10; +const DEFAULT_MAX_CONCURRENT_UPLOADS: usize = 16; +const DEFAULT_MAX_QUEUE_SIZE: usize = 50_000; +const DEFAULT_MAX_RETRY_AGE_SECONDS: u64 = 3600; +const DEFAULT_UPLOAD_ATTEMPT_TIMEOUT_SECONDS: u64 = 30; +/// Where the proxy signs non-LLM AWS calls when no region is configured. +const DEFAULT_REGION: &str = "us-west-2"; + +/// A configured value once its `os.environ/` reference is resolved. +#[derive(Clone, Debug, PartialEq)] +enum Param { + Bool(bool), + Int(i64), + Float(f64), + Text(String), +} + +impl Param { + /// Python's `bool(value)`, which the proxy applies to the prefix and style flags: any + /// non-empty string is true, `"false"` included. + fn truthy(&self) -> bool { + match self { + Self::Bool(flag) => *flag, + Self::Int(number) => *number != 0, + Self::Float(number) => *number != 0.0, + Self::Text(text) => !text.is_empty(), + } + } + + /// Pydantic's lax `bool`. + fn boolean(&self) -> Option { + match self { + Self::Bool(flag) => Some(*flag), + Self::Int(0) => Some(false), + Self::Int(1) => Some(true), + Self::Float(number) if *number == 0.0 => Some(false), + Self::Float(number) if *number == 1.0 => Some(true), + Self::Text(text) => match text.trim().to_ascii_lowercase().as_str() { + "0" | "off" | "f" | "false" | "n" | "no" => Some(false), + "1" | "on" | "t" | "true" | "y" | "yes" => Some(true), + _ => None, + }, + _ => None, + } + } + + /// Pydantic's lax `int`. + fn integer(&self) -> Option { + match self { + Self::Bool(flag) => Some(i64::from(*flag)), + Self::Int(number) => Some(*number), + Self::Float(number) if number.fract() == 0.0 && number.is_finite() => { + Some(*number as i64) + } + Self::Text(text) => text.trim().parse().ok(), + Self::Float(_) => None, + } + } + + fn is_blank(&self) -> bool { + matches!(self, Self::Text(text) if text.is_empty()) + } +} + +/// The resolved `s3_callback_params`. +struct Params(Vec<(String, Param)>); + +impl Params { + fn get(&self, name: &str) -> Option<&Param> { + self.0 + .iter() + .find(|(key, _)| key == name) + .map(|(_, value)| value) + } + + /// `params.get(name) or None`: an empty or false value counts as unset. + fn text(&self, name: &str) -> Result, Error> { + match self.get(name) { + None => Ok(None), + Some(Param::Text(text)) => Ok(Some(text.clone()).filter(|text| !text.is_empty())), + Some(Param::Int(number)) => Ok(Some(number.to_string())), + Some(Param::Float(number)) => Ok(Some(number.to_string())), + Some(Param::Bool(_)) => Err(Error::S3Log(format!("{name} must be a string"))), + } + } + + fn truthy(&self, name: &str) -> bool { + self.get(name).is_some_and(Param::truthy) + } + + /// A flag the proxy parses with pydantic; an invalid value warns and takes `invalid`. + fn flag(&self, name: &str, default: bool, invalid: bool) -> bool { + match self.get(name) { + None => default, + Some(value) if value.is_blank() => default, + Some(value) => value.boolean().unwrap_or_else(|| { + tracing::warn!( + setting = name, + ?value, + "s3 logging: not a boolean, using {invalid}" + ); + invalid + }), + } + } + + /// A positive integer; anything else warns and takes `default`. + fn positive(&self, name: &str, default: usize, reject_bool: bool) -> usize { + let Some(value) = self.get(name).filter(|value| !value.is_blank()) else { + return default; + }; + let parsed = match value { + Param::Bool(_) if reject_bool => None, + value => value.integer(), + }; + match parsed.and_then(|number| usize::try_from(number).ok()) { + Some(number) if number >= 1 => number, + _ => { + tracing::warn!( + setting = name, + ?value, + "s3 logging: not a positive integer, using {default}" + ); + default + } + } + } +} + +/// How each request is logged and where it goes. +#[derive(Clone, Debug, PartialEq)] +pub(super) struct Settings { + pub bucket: String, + /// `s3_path`, the prefix every key starts with. + pub path: String, + /// The region requests are signed for and the default endpoint names. + pub region: String, + pub endpoint_url: Option, + pub virtual_hosted_style: bool, + /// `s3_verify`; unset keeps the gateway's own TLS settings. + pub verify: Option, + pub auth: AwsAuthConfig, + pub team_prefix: bool, + pub key_prefix: bool, + pub strip_base64_files: bool, + pub prompts_only: bool, + pub granularity: Granularity, + pub server_side_encryption: Option, + pub sse_kms_key_id: Option, + pub max_concurrent_uploads: usize, + pub max_queue_size: usize, + /// How long an upload that keeps failing is retried; `None` retries it until the queue + /// overflows. + pub max_retry_age: Option, + pub drop_on_terminal_error: bool, + pub batch_file_upload: bool, + pub flush_interval: Duration, + pub attempt_timeout: Duration, + pub log_success: bool, + pub log_failure: bool, + pub turn_off_message_logging: bool, + pub redact_user_api_key_info: bool, + pub user_agent_tags: bool, +} + +impl Settings { + /// `Ok(None)` when no callback list names `s3_v2` or `s3_log`. A config the proxy could + /// never upload with stops the gateway. + pub async fn from_config( + config: &Config, + secrets: &dyn SecretSource, + env: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, Error> { + let litellm = &config.litellm_settings; + let listed = |list: &Option>| { + list.iter().flatten().any(|callback| { + callback + .as_str() + .is_some_and(|name| NAMES.contains(&name.trim())) + }) + }; + let everywhere = listed(&litellm.callbacks); + let log_success = everywhere || listed(&litellm.success_callback); + let log_failure = everywhere || listed(&litellm.failure_callback); + if !log_success && !log_failure { + return Ok(None); + } + let params = resolve(litellm.additional_fields.get("s3_callback_params"), secrets).await?; + for (name, why) in [ + ( + "s3_flush_interval", + "the proxy reads DEFAULT_S3_FLUSH_INTERVAL_SECONDS instead", + ), + ("s3_batch_size", "the proxy flushes on its interval alone"), + ( + "s3_adaptive_concurrency", + "this gateway uploads s3_max_concurrent_uploads at a time", + ), + ] { + if params.get(name).is_some() { + tracing::warn!(setting = name, "s3 logging: ignored, {why}"); + } + } + let bucket = params + .text("s3_bucket_name")? + .ok_or_else(|| Error::S3Log("s3_bucket_name is required".into()))?; + let configured_region = params.text("s3_region_name")?; + let region = configured_region + .clone() + .or_else(|| env("AWS_REGION")) + .or_else(|| env("AWS_REGION_NAME")) + .unwrap_or_else(|| DEFAULT_REGION.into()); + let (server_side_encryption, sse_kms_key_id) = sse( + params.text("s3_server_side_encryption")?, + params.text("s3_sse_kms_key_id")?, + ); + let verify = match params.get("s3_verify") { + None => None, + Some(Param::Bool(true)) => Some(SslVerify::Enabled), + Some(Param::Bool(false)) => Some(SslVerify::Disabled), + Some(Param::Text(text)) => Some(SslVerify::parse(text)), + Some(other) => { + return Err(Error::S3Log(format!( + "s3_verify must be a boolean or a CA bundle path, not {other:?}" + ))); + } + }; + let env_number = |name: &str, default: u64| { + env(name) + .and_then(|value| value.trim().parse::().ok()) + .unwrap_or(default) + }; + let max_retry_age = match params.get("s3_max_retry_age_seconds") { + None => Some(DEFAULT_MAX_RETRY_AGE_SECONDS), + Some(value) if value.is_blank() => Some(DEFAULT_MAX_RETRY_AGE_SECONDS), + Some(value) => match value { + Param::Bool(_) => None, + value => value.integer(), + } + .filter(|seconds| *seconds >= 0) + .map(|seconds| seconds as u64) + .or_else(|| { + tracing::warn!(?value, "s3 logging: s3_max_retry_age_seconds is not a non-negative integer, using {DEFAULT_MAX_RETRY_AGE_SECONDS}"); + Some(DEFAULT_MAX_RETRY_AGE_SECONDS) + }), + } + .filter(|seconds| *seconds > 0) + .map(Duration::from_secs); + let prompts_only = match params.get("s3_log_prompts_only").cloned() { + None => env("S3_LOG_PROMPTS_ONLY").map(Param::Text), + value => value, + } + .filter(|value| !value.is_blank()) + .is_some_and(|value| { + value.boolean().unwrap_or_else(|| { + tracing::warn!( + ?value, + "s3 logging: s3_log_prompts_only is not a boolean, logging prompts only" + ); + true + }) + }); + let granularity = match params.get("s3_partition_granularity").cloned() { + None => env("S3_PARTITION_GRANULARITY").map(Param::Text), + value => value, + }; + let granularity = match granularity { + Some(Param::Text(text)) if text == "hour" => Granularity::Hour, + None => Granularity::Day, + Some(Param::Text(text)) if text.is_empty() || text == "day" => Granularity::Day, + Some(value) => { + tracing::warn!( + ?value, + "s3 logging: s3_partition_granularity is not one of day, hour, using day" + ); + Granularity::Day + } + }; + let litellm_flag = + |name: &str| matches!(litellm.additional_fields.get(name), Some(Value::Bool(true))); + Ok(Some(Self { + bucket, + path: params.text("s3_path")?.unwrap_or_default(), + region: region.clone(), + endpoint_url: params.text("s3_endpoint_url")?, + virtual_hosted_style: params.truthy("s3_use_virtual_hosted_style"), + verify, + auth: AwsAuthConfig { + access_key_id: params.text("s3_aws_access_key_id")?, + secret_access_key: params.text("s3_aws_secret_access_key")?, + session_token: params.text("s3_aws_session_token")?, + region_name: Some(region), + session_name: params.text("s3_aws_session_name")?, + profile_name: params.text("s3_aws_profile_name")?, + role_name: params.text("s3_aws_role_name")?, + web_identity_token: params.text("s3_aws_web_identity_token")?, + sts_endpoint: params.text("s3_aws_sts_endpoint")?, + external_id: None, + }, + team_prefix: params.truthy("s3_use_team_prefix"), + key_prefix: params.truthy("s3_use_key_prefix"), + strip_base64_files: params.truthy("s3_strip_base64_files"), + prompts_only, + granularity, + server_side_encryption, + sse_kms_key_id, + max_concurrent_uploads: params.positive( + "s3_max_concurrent_uploads", + env_number( + "DEFAULT_S3_MAX_CONCURRENT_UPLOADS", + DEFAULT_MAX_CONCURRENT_UPLOADS as u64, + ) as usize, + false, + ), + max_queue_size: params.positive("s3_max_queue_size", DEFAULT_MAX_QUEUE_SIZE, true), + max_retry_age, + drop_on_terminal_error: params.flag("s3_drop_on_terminal_error", true, true), + batch_file_upload: params.flag("s3_batch_file_upload", false, false), + flush_interval: Duration::from_secs( + env_number( + "DEFAULT_S3_FLUSH_INTERVAL_SECONDS", + DEFAULT_FLUSH_INTERVAL_SECONDS, + ) + .max(1), + ), + attempt_timeout: Duration::from_secs( + env_number( + "S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS", + DEFAULT_UPLOAD_ATTEMPT_TIMEOUT_SECONDS, + ) + .clamp(1, 3600), + ), + log_success, + log_failure, + turn_off_message_logging: litellm_flag("turn_off_message_logging"), + redact_user_api_key_info: litellm_flag("redact_user_api_key_info"), + user_agent_tags: !litellm_flag("disable_add_user_agent_to_request_tags"), + })) + } +} + +/// The encryption headers to send: a KMS key alone means `aws:kms`, and a key next to another +/// algorithm is dropped, as the proxy's `resolve_sse_params` does. +fn sse(algorithm: Option, key_id: Option) -> (Option, Option) { + let algorithm = algorithm.or_else(|| key_id.as_ref().map(|_| "aws:kms".into())); + match (algorithm, key_id) { + (None, _) => (None, None), + (Some(algorithm), Some(_)) if !algorithm.starts_with("aws:kms") => { + tracing::warn!(%algorithm, "s3 logging: ignoring s3_sse_kms_key_id, set s3_server_side_encryption to aws:kms to use it"); + (Some(algorithm), None) + } + (algorithm, key_id) => (algorithm, key_id), + } +} + +async fn resolve(value: Option<&Value>, secrets: &dyn SecretSource) -> Result { + let entries = match value { + None | Some(Value::Null) => return Ok(Params(Vec::new())), + Some(Value::Mapping(entries)) => entries, + Some(_) => return Err(Error::S3Log("s3_callback_params must be a mapping".into())), + }; + let mut params = Vec::with_capacity(entries.len()); + for (key, value) in entries { + let Some(name) = key.as_str() else { + return Err(Error::S3Log( + "s3_callback_params keys must be strings".into(), + )); + }; + let param = match value { + Value::Null => continue, + Value::Bool(flag) => Param::Bool(*flag), + Value::Number(number) => match number.as_i64() { + Some(number) => Param::Int(number), + None => Param::Float(number.as_f64().unwrap_or(f64::NAN)), + }, + Value::String(text) => match text.strip_prefix("os.environ/") { + Some(variable) => match secrets.get_secret_str(variable).await? { + Some(secret) => Param::Text(secret.expose().to_owned()), + None => continue, + }, + None => Param::Text(text.clone()), + }, + _ => return Err(Error::S3Log(format!("{name} must be a scalar"))), + }; + params.push((name.to_owned(), param)); + } + Ok(Params(params)) +} + +#[cfg(test)] +pub(super) mod tests { + use std::collections::HashMap; + + use futures_util::future::BoxFuture; + use litellm_auth_types::SecretValue; + use rstest::rstest; + + use super::*; + + struct Secrets(HashMap<&'static str, &'static str>); + + impl SecretSource for Secrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { Ok(self.0.get(name).map(|value| SecretValue::new(*value))) }) + } + } + + pub(crate) async fn settings_with( + yaml: &str, + environment: &[(&'static str, &'static str)], + ) -> Result, String> { + let config = Config::from_yaml(&format!("model_list: []\n{yaml}")).unwrap(); + let environment: HashMap<&str, &str> = environment.iter().copied().collect(); + let env = |name: &str| environment.get(name).map(|value| (*value).to_owned()); + Settings::from_config(&config, &Secrets(environment.clone()), &env) + .await + .map_err(|error| error.to_string()) + } + + async fn settings(params: &str) -> Settings { + settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}" + ), + &[], + ) + .await + .unwrap() + .unwrap() + } + + #[rstest] + #[case::none("", None)] + #[case::other_callback(" success_callback: [prometheus]\n", None)] + #[case::callbacks(" callbacks: [s3_v2]\n", Some((true, true)))] + #[case::gateway_name(" callbacks: [s3_log]\n", Some((true, true)))] + #[case::success_only(" success_callback: [s3_v2]\n", Some((true, false)))] + #[case::failure_only(" failure_callback: s3_v2\n", Some((false, true)))] + #[tokio::test] + async fn on_for_the_lists_that_name_it( + #[case] lists: &str, + #[case] expected: Option<(bool, bool)>, + ) { + let yaml = + format!("litellm_settings:\n{lists} s3_callback_params:\n s3_bucket_name: logs\n"); + let settings = settings_with(&yaml, &[]).await.unwrap(); + assert_eq!( + settings.map(|settings| (settings.log_success, settings.log_failure)), + expected + ); + } + + #[tokio::test] + async fn defaults_are_the_proxys() { + let settings = settings("").await; + assert_eq!(settings.path, ""); + assert_eq!(settings.region, "us-west-2"); + assert_eq!(settings.granularity, Granularity::Day); + assert_eq!(settings.max_concurrent_uploads, 16); + assert_eq!(settings.max_queue_size, 50_000); + assert_eq!(settings.max_retry_age, Some(Duration::from_secs(3600))); + assert!(settings.drop_on_terminal_error); + assert!(!settings.batch_file_upload); + assert!(!settings.prompts_only); + assert_eq!(settings.flush_interval, Duration::from_secs(10)); + assert_eq!(settings.attempt_timeout, Duration::from_secs(30)); + assert_eq!(settings.verify, None); + assert_eq!(settings.server_side_encryption, None); + } + + #[tokio::test] + async fn references_are_resolved_from_the_secret_source() { + let settings = settings_with( + "litellm_settings: + callbacks: [s3_v2] + s3_callback_params: + s3_bucket_name: os.environ/BUCKET + s3_region_name: eu-west-1 + s3_aws_access_key_id: os.environ/KEY_ID + s3_aws_secret_access_key: os.environ/SECRET + s3_aws_session_token: os.environ/UNSET + s3_use_team_prefix: os.environ/TEAM_PREFIX +", + &[ + ("BUCKET", "audit"), + ("KEY_ID", "AKID"), + ("SECRET", "shh"), + ("TEAM_PREFIX", "false"), + ], + ) + .await + .unwrap() + .unwrap(); + assert_eq!(settings.bucket, "audit"); + assert_eq!(settings.region, "eu-west-1"); + assert_eq!(settings.auth.access_key_id.as_deref(), Some("AKID")); + assert_eq!(settings.auth.secret_access_key.as_deref(), Some("shh")); + assert_eq!(settings.auth.session_token, None); + assert_eq!(settings.auth.region_name.as_deref(), Some("eu-west-1")); + // The proxy takes `bool()` of the string, so any non-empty text turns it on. + assert!(settings.team_prefix); + } + + #[rstest] + #[case::configured(" s3_region_name: ap-south-1\n", &[("AWS_REGION", "us-east-2")], "ap-south-1")] + #[case::aws_region_wins("", &[("AWS_REGION", "us-east-2"), ("AWS_REGION_NAME", "eu-west-3")], "us-east-2")] + #[case::aws_region_name("", &[("AWS_REGION_NAME", "eu-west-3")], "eu-west-3")] + #[tokio::test] + async fn the_signing_region_falls_back_like_the_proxy( + #[case] params: &str, + #[case] environment: &[(&'static str, &'static str)], + #[case] expected: &str, + ) { + let settings = settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}" + ), + environment, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(settings.region, expected); + } + + #[rstest] + #[case::kms_key_alone(" s3_sse_kms_key_id: key-1\n", Some("aws:kms"), Some("key-1"))] + #[case::aes(" s3_server_side_encryption: AES256\n", Some("AES256"), None)] + #[case::key_with_aes( + " s3_server_side_encryption: AES256\n s3_sse_kms_key_id: key-1\n", + Some("AES256"), + None + )] + #[case::dsse( + " s3_server_side_encryption: aws:kms:dsse\n s3_sse_kms_key_id: key-1\n", + Some("aws:kms:dsse"), + Some("key-1") + )] + #[tokio::test] + async fn encryption_settings_combine_like_the_proxy( + #[case] params: &str, + #[case] algorithm: Option<&str>, + #[case] key: Option<&str>, + ) { + let settings = settings(params).await; + assert_eq!(settings.server_side_encryption.as_deref(), algorithm); + assert_eq!(settings.sse_kms_key_id.as_deref(), key); + } + + #[rstest] + #[case::text_number(" s3_max_concurrent_uploads: ' 4 '\n", 4, 50_000)] + #[case::zero_falls_back(" s3_max_concurrent_uploads: 0\n", 16, 50_000)] + #[case::garbage_falls_back(" s3_max_queue_size: lots\n", 16, 50_000)] + #[case::bool_queue_falls_back(" s3_max_queue_size: true\n", 16, 50_000)] + #[case::bool_width_is_one(" s3_max_concurrent_uploads: true\n", 1, 50_000)] + #[case::queue(" s3_max_queue_size: 10\n", 16, 10)] + #[tokio::test] + async fn numbers_parse_or_fall_back( + #[case] params: &str, + #[case] uploads: usize, + #[case] queue: usize, + ) { + let settings = settings(params).await; + assert_eq!(settings.max_concurrent_uploads, uploads); + assert_eq!(settings.max_queue_size, queue); + } + + #[rstest] + #[case::zero_disables(" s3_max_retry_age_seconds: 0\n", None)] + #[case::seconds(" s3_max_retry_age_seconds: '60'\n", Some(60))] + #[case::negative_falls_back(" s3_max_retry_age_seconds: -1\n", Some(3600))] + #[case::blank_is_default(" s3_max_retry_age_seconds: ''\n", Some(3600))] + #[tokio::test] + async fn retry_age(#[case] params: &str, #[case] seconds: Option) { + assert_eq!( + settings(params).await.max_retry_age, + seconds.map(Duration::from_secs) + ); + } + + #[rstest] + #[case::configured(" s3_log_prompts_only: 'yes'\n", &[], true)] + #[case::invalid_logs_prompts_only(" s3_log_prompts_only: sometimes\n", &[], true)] + #[case::environment("", &[("S3_LOG_PROMPTS_ONLY", "true")], true)] + #[case::configured_beats_environment(" s3_log_prompts_only: false\n", &[("S3_LOG_PROMPTS_ONLY", "true")], false)] + #[case::unset("", &[], false)] + #[tokio::test] + async fn prompts_only( + #[case] params: &str, + #[case] environment: &[(&'static str, &'static str)], + #[case] expected: bool, + ) { + let settings = settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}" + ), + environment, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(settings.prompts_only, expected); + } + + #[rstest] + #[case::hour(" s3_partition_granularity: hour\n", &[], Granularity::Hour)] + #[case::environment("", &[("S3_PARTITION_GRANULARITY", "hour")], Granularity::Hour)] + #[case::unknown_is_day(" s3_partition_granularity: minute\n", &[], Granularity::Day)] + #[tokio::test] + async fn partition_granularity( + #[case] params: &str, + #[case] environment: &[(&'static str, &'static str)], + #[case] expected: Granularity, + ) { + let settings = settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}" + ), + environment, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(settings.granularity, expected); + } + + #[rstest] + #[case::no_bucket( + "litellm_settings:\n callbacks: [s3_v2]\n", + "s3_bucket_name is required" + )] + #[case::unset_bucket_reference( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: os.environ/NOPE\n", + "s3_bucket_name is required" + )] + #[case::not_a_mapping( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params: [a]\n", + "s3_callback_params must be a mapping" + )] + #[case::nested_value( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_path: [a, b]\n", + "s3_path must be a scalar" + )] + #[case::bad_verify( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_verify: 3\n", + "s3_verify must be a boolean or a CA bundle path" + )] + #[tokio::test] + async fn a_config_that_cannot_upload_stops_the_gateway( + #[case] yaml: &str, + #[case] message: &str, + ) { + let error = settings_with(yaml, &[]).await.unwrap_err(); + assert!(error.contains(message), "{error}"); + } +} diff --git a/litellm-rust/crates/gateway/src/s3_log/key.rs b/litellm-rust/crates/gateway/src/s3_log/key.rs new file mode 100644 index 00000000000..19bc6be2d34 --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/key.rs @@ -0,0 +1,389 @@ +//! Object keys, download names and URLs, laid out as the proxy's `get_s3_object_key` and +//! `S3Logger._build_object_url` lay them out, length bounds included. + +use std::time::{SystemTime, UNIX_EPOCH}; + +use sha2::{Digest, Sha256}; + +use super::config::Settings; +use crate::spend::civil; + +/// S3's limit on a key, in UTF-8 bytes. +const MAX_KEY_BYTES: usize = 1024; +/// The proxy's limit on the last segment of a key. +const MAX_FILE_NAME_BYTES: usize = 255; +/// The proxy's limit on the `Content-Disposition` file name. +const MAX_DOWNLOAD_NAME_BYTES: usize = 1024; +/// How much of an over-long name is kept before its digest. +const BOUNDED_HEAD_BYTES: i64 = 64; +const MIN_BOUNDED_FILE_NAME_BYTES: usize = 64; +const PREFIX_DIGEST_CHARS: usize = 16; +const EXTENSION: &str = ".json"; + +/// `s3_partition_granularity`: a folder per day, or per hour within it. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(super) enum Granularity { + #[default] + Day, + Hour, +} + +/// A wall-clock time, UTC, split the way the key names it. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct Utc { + year: i64, + month: u32, + day: u32, + hour: u64, + minute: u64, + second: u64, + micros: u32, +} + +impl Utc { + pub fn new(time: SystemTime) -> Self { + let elapsed = time.duration_since(UNIX_EPOCH).unwrap_or_default(); + let seconds = elapsed.as_secs(); + let (year, month, day) = civil((seconds / 86_400) as i64); + let of_day = seconds % 86_400; + Self { + year, + month, + day, + hour: of_day / 3600, + minute: of_day / 60 % 60, + second: of_day % 60, + micros: elapsed.subsec_micros(), + } + } + + fn date(&self) -> String { + format!("{:04}-{:02}-{:02}", self.year, self.month, self.day) + } + + /// `%H-%M-%S`. + pub fn clock(&self) -> String { + format!("{:02}-{:02}-{:02}", self.hour, self.minute, self.second) + } + + /// `%H-%M-%S-%f`. + fn clock_micros(&self) -> String { + format!("{}-{:06}", self.clock(), self.micros) + } +} + +/// The proxy's `get_logging_id`: the request's time and the payload id. +pub(super) fn logging_id(start: &Utc, id: &str) -> String { + format!("time-{}_{id}", start.clock_micros()) +} + +/// `{s3_path}/[prefix]YYYY-MM-DD[/HH]/{file_name}.json`, shortened as the proxy shortens it when +/// it would pass S3's limits: the file name first, the configured prefix only if that is what +/// does not fit. +pub(super) fn object_key( + path: &str, + prefix: &str, + start: &Utc, + file_name: &str, + granularity: Granularity, +) -> String { + let sanitized = file_name.replace(['/', ':'], "_"); + let configured = match path.is_empty() { + true => prefix.to_owned(), + false => format!("{}/{prefix}", path.trim_end_matches('/')), + }; + let date = match granularity { + Granularity::Day => format!("{}/", start.date()), + Granularity::Hour => format!("{}/{:02}/", start.date(), start.hour), + }; + let file_name_budget = MAX_FILE_NAME_BYTES - EXTENSION.len(); + let key = format!("{configured}{date}{sanitized}{EXTENSION}"); + if key.len() <= MAX_KEY_BYTES && sanitized.len() <= file_name_budget { + return key; + } + let budget = MAX_KEY_BYTES - date.len() - EXTENSION.len(); + if configured.len() + MIN_BOUNDED_FILE_NAME_BYTES <= budget { + let bounded = bounded_name( + file_name, + &sanitized, + file_name_budget.min(budget - configured.len()), + ); + return format!("{configured}{date}{bounded}{EXTENSION}"); + } + let shortest = bounded_name(file_name, &sanitized, MIN_BOUNDED_FILE_NAME_BYTES); + let prefix = bounded_prefix(&configured, budget.saturating_sub(shortest.len())); + format!("{prefix}{date}{shortest}{EXTENSION}") +} + +/// The proxy's `get_s3_object_download_filename`. +pub(super) fn download_name(start: &Utc, id: &str) -> String { + let stamp = format!("time-{}T{}", start.date(), start.clock_micros()); + let name = format!("{stamp}_{id}"); + let sanitized = format!("{stamp}_{}", id.replace(['/', '"'], "_")); + let budget = MAX_DOWNLOAD_NAME_BYTES - EXTENSION.len(); + match sanitized.len() <= budget { + true => format!("{sanitized}{EXTENSION}"), + false => format!("{}{EXTENSION}", bounded_name(&name, &sanitized, budget)), + } +} + +/// As much of `sanitized` as fits, then the sha256 of the unsanitized name. +fn bounded_name(name: &str, sanitized: &str, max_bytes: usize) -> String { + let digest = hex_sha256(name); + let head_budget = BOUNDED_HEAD_BYTES.min(max_bytes as i64 - digest.len() as i64 - 1); + let head = truncate(sanitized, head_budget); + match head.is_empty() { + true => digest, + false => format!("{head}_{digest}"), + } +} + +/// As much of the configured prefix as fits, then a digest segment naming all of it. +fn bounded_prefix(prefix: &str, max_bytes: usize) -> String { + let digest = format!("{}/", &hex_sha256(prefix)[..PREFIX_DIGEST_CHARS]); + if max_bytes < digest.len() { + return String::new(); + } + let head = truncate(prefix, max_bytes as i64 - digest.len() as i64 - 1).trim_end_matches('/'); + match head.is_empty() { + true => digest, + false => format!("{head}/{digest}"), + } +} + +/// `value` cut to at most `max_bytes` of UTF-8, never inside a character. +fn truncate(value: &str, max_bytes: i64) -> &str { + let Ok(max_bytes) = usize::try_from(max_bytes) else { + return ""; + }; + if value.len() <= max_bytes { + return value; + } + let end = (0..=max_bytes) + .rev() + .find(|end| value.is_char_boundary(*end)) + .unwrap_or(0); + &value[..end] +} + +fn hex_sha256(value: &str) -> String { + Sha256::digest(value.as_bytes()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +/// Everything before the last `/`, the folder batch files are grouped by. +pub(super) fn parent(key: &str) -> &str { + key.rsplit_once('/').map_or("", |(parent, _)| parent) +} + +/// The URL that is signed and sent, with the key percent-encoded once, as the proxy builds it. +pub(super) fn object_url(settings: &Settings, key: &str) -> String { + let key = encode(key); + let bucket = &settings.bucket; + match &settings.endpoint_url { + Some(endpoint) if settings.virtual_hosted_style => { + let (scheme, host) = match endpoint.strip_prefix("https://") { + Some(host) => ("https://", host), + None => ("http://", endpoint.trim_start_matches("http://")), + }; + format!("{scheme}{bucket}.{host}/{key}") + } + Some(endpoint) => format!("{endpoint}/{bucket}/{key}"), + None => format!( + "https://{bucket}.s3.{region}.{suffix}/{key}", + region = settings.region, + suffix = dns_suffix(&settings.region) + ), + } +} + +/// Python's `quote(key, safe="/")`. +fn encode(key: &str) -> String { + key.bytes() + .map(|byte| match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' | b'/' => { + char::from(byte).to_string() + } + byte => format!("%{byte:02X}"), + }) + .collect() +} + +/// The AWS partition's domain for a region. +fn dns_suffix(region: &str) -> &'static str { + [ + ("cn-", "amazonaws.com.cn"), + ("us-isob-", "sc2s.sgov.gov"), + ("us-isof-", "csp.hci.ic.gov"), + ("us-iso-", "c2s.ic.gov"), + ("eu-isoe-", "cloud.adc-e.uk"), + ] + .into_iter() + .find(|(prefix, _)| region.starts_with(prefix)) + .map_or("amazonaws.com", |(_, suffix)| suffix) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use rstest::rstest; + + use super::*; + + /// 2026-03-04 05:06:07.000089 UTC. + fn start() -> Utc { + Utc::new(UNIX_EPOCH + Duration::from_micros(1_772_600_767_000_089)) + } + + #[rstest] + #[case::bare( + "", + "", + Granularity::Day, + "2026-03-04/time-05-06-07-000089_chatcmpl-1.json" + )] + #[case::path( + "logs/", + "", + Granularity::Day, + "logs/2026-03-04/time-05-06-07-000089_chatcmpl-1.json" + )] + #[case::prefixes( + "logs", + "research/team-key/", + Granularity::Day, + "logs/research/team-key/2026-03-04/time-05-06-07-000089_chatcmpl-1.json" + )] + #[case::hour( + "logs", + "", + Granularity::Hour, + "logs/2026-03-04/05/time-05-06-07-000089_chatcmpl-1.json" + )] + fn keys_are_laid_out_like_the_proxy( + #[case] path: &str, + #[case] prefix: &str, + #[case] granularity: Granularity, + #[case] expected: &str, + ) { + let file_name = logging_id(&start(), "chatcmpl-1"); + assert_eq!( + object_key(path, prefix, &start(), &file_name, granularity), + expected + ); + } + + #[test] + fn slashes_and_colons_in_the_id_stay_in_one_segment() { + let file_name = logging_id(&start(), "a/b:c"); + assert_eq!( + object_key("", "", &start(), &file_name, Granularity::Day), + "2026-03-04/time-05-06-07-000089_a_b_c.json" + ); + assert_eq!( + download_name(&start(), "a/b\"c"), + "time-2026-03-04T05-06-07-000089_a_b_c.json" + ); + } + + #[test] + fn a_long_id_is_shortened_to_a_digest_within_s3_limits() { + let id = "é".repeat(400); + let file_name = logging_id(&start(), &id); + let key = object_key("logs", "", &start(), &file_name, Granularity::Day); + let (folder, last) = key.rsplit_once('/').unwrap(); + assert_eq!(folder, "logs/2026-03-04"); + assert!(last.len() <= MAX_FILE_NAME_BYTES, "{last}"); + assert!(last.ends_with(&format!("_{}.json", hex_sha256(&file_name)))); + assert!(last.starts_with("time-05-06-07-000089_é")); + } + + #[test] + fn a_long_prefix_keeps_its_head_and_a_digest_of_all_of_it() { + let path = "p".repeat(2000); + let file_name = logging_id(&start(), "chatcmpl-1"); + let key = object_key(&path, "", &start(), &file_name, Granularity::Day); + assert!(key.len() <= MAX_KEY_BYTES, "{}", key.len()); + let digest = &hex_sha256(&format!("{path}/"))[..PREFIX_DIGEST_CHARS]; + assert!(key.starts_with("ppp"), "{key}"); + assert!(key.contains(&format!("/{digest}/2026-03-04/")), "{key}"); + } + + fn settings(endpoint: Option<&str>, virtual_hosted: bool, region: &str) -> Settings { + Settings { + bucket: "logs".into(), + path: String::new(), + region: region.into(), + endpoint_url: endpoint.map(str::to_owned), + virtual_hosted_style: virtual_hosted, + verify: None, + auth: Default::default(), + team_prefix: false, + key_prefix: false, + strip_base64_files: false, + prompts_only: false, + granularity: Granularity::Day, + server_side_encryption: None, + sse_kms_key_id: None, + max_concurrent_uploads: 1, + max_queue_size: 1, + max_retry_age: None, + drop_on_terminal_error: true, + batch_file_upload: false, + flush_interval: Duration::from_secs(1), + attempt_timeout: Duration::from_secs(1), + log_success: true, + log_failure: true, + turn_off_message_logging: false, + redact_user_api_key_info: false, + user_agent_tags: true, + } + } + + #[rstest] + #[case::aws( + None, + false, + "eu-west-1", + "https://logs.s3.eu-west-1.amazonaws.com/a/b%20c%3D%2B.json" + )] + #[case::china( + None, + false, + "cn-north-1", + "https://logs.s3.cn-north-1.amazonaws.com.cn/a/b%20c%3D%2B.json" + )] + #[case::path_style( + Some("http://minio:9000"), + false, + "us-east-1", + "http://minio:9000/logs/a/b%20c%3D%2B.json" + )] + #[case::virtual_hosted( + Some("https://s3.example"), + true, + "us-east-1", + "https://logs.s3.example/a/b%20c%3D%2B.json" + )] + fn urls_are_built_like_the_proxy( + #[case] endpoint: Option<&str>, + #[case] virtual_hosted: bool, + #[case] region: &str, + #[case] expected: &str, + ) { + assert_eq!( + object_url(&settings(endpoint, virtual_hosted, region), "a/b c=+.json"), + expected + ); + } + + #[rstest] + #[case("a/b/c.json", "a/b")] + #[case("c.json", "")] + fn parents(#[case] key: &str, #[case] expected: &str) { + assert_eq!(parent(key), expected); + } +} diff --git a/litellm-rust/crates/gateway/src/s3_log/mod.rs b/litellm-rust/crates/gateway/src/s3_log/mod.rs new file mode 100644 index 00000000000..73437dbdcfe --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/mod.rs @@ -0,0 +1,506 @@ +//! The `s3_v2` callback: every request's `StandardLoggingPayload` written to S3 as a JSON +//! object, under the proxy's key layout, as the proxy's `S3Logger` writes it. +//! +//! It is an observer, so nothing here runs on the request path. Each finished request is turned +//! into its object and queued; every flush interval the queue is uploaded with bounded +//! concurrency, or written as one NDJSON object per folder with `s3_batch_file_upload`. What +//! fails for now goes back on the queue for the next flush, until it is older than +//! `s3_max_retry_age_seconds` or the queue overflows. + +mod config; +mod key; +mod payload; +mod upload; + +use std::{ + collections::{BTreeMap, VecDeque}, + sync::{Arc, Mutex, Weak}, + time::{Instant, SystemTime}, +}; + +use axum::body::Bytes; +use futures_util::{StreamExt, future::BoxFuture}; +use litellm_config::Config; +use litellm_gateway_hooks::{Facts, Observer, Record}; +use litellm_gateway_inference::Gateway; +use serde_json::Value; + +use self::{ + config::Settings, + key::{Utc, download_name, logging_id, object_key, parent}, + payload::{Redaction, prompts_only, standard_logging_payload, strip_base64_files}, + upload::{Upload, Uploader}, +}; +use crate::Error; + +const JSON: &str = "application/json"; +const NDJSON: &str = "application/x-ndjson"; + +/// One object waiting to be written. +#[derive(Clone, Debug)] +struct Element { + key: String, + download_name: String, + body: Bytes, + content_type: &'static str, + /// A batch file, already holding many requests; never grouped again. + batch: bool, + /// When this object first failed to upload. + retrying_since: Option, +} + +#[derive(Default)] +struct Queue { + elements: VecDeque, + /// How many at the front were put back by the last failed flush. + requeued: usize, +} + +/// The `s3_v2` observer. +pub struct S3Log { + settings: Arc, + uploader: Uploader, + queue: Mutex, + /// Serializes flushes, so the periodic one and the one at shutdown never interleave. + flushing: tokio::sync::Mutex<()>, +} + +impl S3Log { + /// Starts the observer and its flush timer, or `Ok(None)` when no callback list names + /// `s3_v2`. Must be called inside a tokio runtime. + pub async fn from_config( + config: &Config, + inference: &Gateway, + ) -> Result>, Error> { + let env = |name: &str| std::env::var(name).ok(); + let Some(settings) = + Settings::from_config(config, inference.secrets.as_ref(), &env).await? + else { + return Ok(None); + }; + let settings = Arc::new(settings); + let log = Arc::new(Self { + uploader: Uploader::new( + settings.clone(), + inference.resources.pool.clone(), + &inference.http, + ), + settings, + queue: Mutex::default(), + flushing: tokio::sync::Mutex::new(()), + }); + tokio::spawn(run(Arc::downgrade(&log))); + Ok(Some(log)) + } + + fn element(&self, record: &Record) -> Element { + let settings = &self.settings; + let redaction = Redaction { + messages: settings.turn_off_message_logging, + user_api_key_info: settings.redact_user_api_key_info, + user_agent_tags: settings.user_agent_tags, + }; + let mut payload = standard_logging_payload(record, redaction); + if settings.strip_base64_files { + strip_base64_files(&mut payload); + } + let metadata = |field: &str| { + payload["metadata"][field] + .as_str() + .filter(|value| !value.is_empty()) + .map(str::to_owned) + }; + let prefix: String = [ + settings + .team_prefix + .then(|| metadata("user_api_key_team_alias")), + settings.key_prefix.then(|| metadata("user_api_key_alias")), + ] + .into_iter() + .flatten() + .flatten() + .map(|segment| format!("{segment}/")) + .collect(); + let id = payload["id"].as_str().unwrap_or_default().to_owned(); + let start = Utc::new(record.request.started_at); + if settings.prompts_only { + prompts_only(&mut payload); + } + Element { + key: object_key( + &settings.path, + &prefix, + &start, + &logging_id(&start, &id), + settings.granularity, + ), + download_name: download_name(&start, &id), + body: Bytes::from(Value::Object(payload).to_string()), + content_type: JSON, + batch: false, + retrying_since: None, + } + } + + fn enqueue(&self, element: Element) { + let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner()); + queue.elements.push_back(element); + let excess = queue + .elements + .len() + .saturating_sub(self.settings.max_queue_size); + if excess > 0 { + queue.elements.drain(..excess); + queue.requeued = queue.requeued.saturating_sub(excess); + tracing::warn!( + dropped = excess, + max_queue_size = self.settings.max_queue_size, + "s3 logging: queue full; dropped the oldest events" + ); + } + } + + /// Uploads everything queued; what fails for now is put back at the front of the queue. + pub async fn flush(&self) { + let _flushing = self.flushing.lock().await; + let (batch, requeued) = { + let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner()); + let requeued = std::mem::take(&mut queue.requeued); + (Vec::from(std::mem::take(&mut queue.elements)), requeued) + }; + if batch.is_empty() { + return; + } + let batch_len = batch.len(); + let uploads = match self.settings.batch_file_upload { + true => batch_files(batch), + false => batch, + }; + // What failed last time goes last, so a backlog never starves new requests. + let stale = match uploads.len() == batch_len { + true => requeued.min(uploads.len()), + false => 0, + }; + let order: Vec = (stale..uploads.len()).chain(0..stale).collect(); + let mut outcomes = vec![Upload::Retry; uploads.len()]; + let mut finished = futures_util::stream::iter(order) + .map(|index| { + let uploads = &uploads; + async move { (index, self.uploader.upload(&uploads[index]).await) } + }) + .buffer_unordered(self.settings.max_concurrent_uploads); + while let Some((index, outcome)) = finished.next().await { + outcomes[index] = outcome; + } + drop(finished); + let delivered = outcomes + .iter() + .filter(|outcome| **outcome == Upload::Delivered) + .count(); + let total = uploads.len(); + let failed = total - delivered; + if failed == 0 { + return; + } + let now = Instant::now(); + let max_age = self.settings.max_retry_age; + let failures = uploads + .into_iter() + .zip(outcomes) + .filter(|(_, outcome)| *outcome != Upload::Delivered); + // When nothing got through the bucket itself is failing: keep everything, unaged. + let kept: Vec = match delivered { + 0 => failures.map(|(element, _)| element).collect(), + _ => failures + .filter(|(element, outcome)| { + *outcome == Upload::Retry + && !matches!( + (max_age, element.retrying_since), + (Some(max_age), Some(since)) if now.duration_since(since) > max_age + ) + }) + .map(|(element, _)| Element { + retrying_since: element.retrying_since.or(max_age.map(|_| now)), + ..element + }) + .collect(), + }; + tracing::warn!( + failed, + kept = kept.len(), + total, + "s3 logging: uploads failed this flush; the kept ones are retried on the next, the rest \ + were rejected by S3 or retried longer than s3_max_retry_age_seconds" + ); + if kept.is_empty() { + return; + } + let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner()); + let arrivals = std::mem::take(&mut queue.elements); + let kept_len = kept.len(); + let mut elements: VecDeque = kept.into_iter().chain(arrivals).collect(); + let overflow = elements.len().saturating_sub(self.settings.max_queue_size); + if overflow > 0 { + elements.drain(..overflow); + tracing::warn!( + dropped = overflow, + max_queue_size = self.settings.max_queue_size, + "s3 logging: queue exceeded max_queue_size after a failed flush; dropped the oldest events" + ); + } + queue.elements = elements; + queue.requeued = kept_len.saturating_sub(overflow); + } +} + +/// `s3_batch_file_upload`: one NDJSON object per folder, named for the flush. +fn batch_files(batch: Vec) -> Vec { + let now = Utc::new(SystemTime::now()); + let (ready, singles): (Vec, Vec) = + batch.into_iter().partition(|element| element.batch); + let mut folders: BTreeMap> = BTreeMap::new(); + for element in singles { + folders + .entry(parent(&element.key).to_owned()) + .or_default() + .push(element); + } + let files = folders.into_iter().map(|(folder, elements)| { + let name = format!("batch_{}_{}", now.clock(), uuid::Uuid::new_v4().simple()); + let lines: Vec<&[u8]> = elements.iter().map(|element| &element.body[..]).collect(); + Element { + key: match folder.is_empty() { + true => format!("{name}.jsonl"), + false => format!("{folder}/{name}.jsonl"), + }, + download_name: format!("{name}.jsonl"), + body: Bytes::from(lines.join(&b'\n')), + content_type: NDJSON, + batch: true, + retrying_since: elements + .iter() + .filter_map(|element| element.retrying_since) + .min(), + } + }); + ready.into_iter().chain(files).collect() +} + +impl Observer for S3Log { + fn name(&self) -> &'static str { + "s3_v2" + } + + fn reads(&self) -> Facts { + Facts::KEY + | Facts::END_USER + | Facts::CLIENT_HEADERS + | Facts::DEPLOYMENT + | Facts::USAGE + | Facts::COST + | Facts::TIMINGS + | Facts::ANNOTATIONS + | Facts::REQUEST_BODY + | Facts::RESPONSE_BODY + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + let logged = match record.succeeded() { + true => self.settings.log_success, + false => self.settings.log_failure, + }; + if logged { + self.enqueue(self.element(record)); + } + Box::pin(async {}) + } + + fn flush(&self) -> BoxFuture<'_, ()> { + Box::pin(S3Log::flush(self)) + } +} + +/// Flushes every flush interval until the observer is dropped. +async fn run(log: Weak) { + let interval = match log.upgrade() { + Some(log) => log.settings.flush_interval, + None => return, + }; + loop { + tokio::time::sleep(interval).await; + let Some(log) = log.upgrade() else { + return; + }; + log.flush().await; + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use litellm_gateway_hooks::Route; + use serde_json::json; + + use super::*; + + fn element(key: &str, body: &str) -> Element { + Element { + key: key.into(), + download_name: "d.json".into(), + body: Bytes::from(body.to_owned()), + content_type: JSON, + batch: false, + retrying_since: None, + } + } + + #[test] + fn batch_files_hold_one_line_per_request_per_folder() { + let earlier = Instant::now(); + let mut retried = element("logs/2026-03-04/b.json", "{\"b\":1}"); + retried.retrying_since = Some(earlier); + let files = batch_files(vec![ + element("logs/2026-03-04/a.json", "{\"a\":1}"), + element("logs/team/2026-03-04/c.json", "{\"c\":1}"), + retried, + ]); + assert_eq!(files.len(), 2); + let day = &files[0]; + let (folder, name) = day.key.rsplit_once('/').unwrap(); + assert_eq!(folder, "logs/2026-03-04"); + assert!( + name.starts_with("batch_") && name.ends_with(".jsonl"), + "{name}" + ); + assert_eq!(day.download_name, name); + assert_eq!(&day.body[..], b"{\"a\":1}\n{\"b\":1}"); + assert_eq!(day.content_type, NDJSON); + assert_eq!(day.retrying_since, Some(earlier)); + assert!(files[1].key.starts_with("logs/team/2026-03-04/batch_")); + // A batch file kept for a retry is sent again as it is. + let again = batch_files(files.clone()); + assert_eq!( + again.iter().map(|file| &file.key).collect::>(), + files.iter().map(|file| &file.key).collect::>() + ); + } + + async fn log(params: &str) -> S3Log { + let settings = config::tests::settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_path: audit\n{params}" + ), + &[], + ) + .await + .unwrap() + .unwrap(); + let settings = Arc::new(settings); + S3Log { + uploader: Uploader::new( + settings.clone(), + Arc::new(litellm_http::HttpClientPool::new(Arc::new( + litellm_http::media::PublicDnsResolver, + ))), + &litellm_http::Resolution::from(&litellm_http::HttpSettings::default()).config, + ), + settings, + queue: Mutex::default(), + flushing: tokio::sync::Mutex::new(()), + } + } + + fn record() -> Record { + payload::tests::record( + Route::ChatCompletions, + 200, + false, + json!({"messages": [{"role": "user", "content": "hi"}]}), + Some(payload::tests::COMPLETION), + ) + } + + #[rstest::rstest] + #[case::plain("", "audit/1970-01-01/time-00-16-40-000000_chatcmpl-1.json")] + #[case::team( + " s3_use_team_prefix: true\n", + "audit/research/1970-01-01/time-00-16-40-000000_chatcmpl-1.json" + )] + #[case::team_and_key( + " s3_use_team_prefix: true\n s3_use_key_prefix: true\n s3_partition_granularity: hour\n", + "audit/research/team-key/1970-01-01/00/time-00-16-40-000000_chatcmpl-1.json" + )] + #[tokio::test] + async fn objects_are_keyed_by_team_key_and_time(#[case] params: &str, #[case] expected: &str) { + let element = log(params).await.element(&record()); + assert_eq!(element.key, expected); + assert_eq!( + element.download_name, + "time-1970-01-01T00-16-40-000000_chatcmpl-1.json" + ); + } + + #[tokio::test] + async fn prompts_only_objects_have_no_response() { + let element = log(" s3_log_prompts_only: true\n") + .await + .element(&record()); + let body: Value = serde_json::from_slice(&element.body).unwrap(); + assert_eq!(body["response"], Value::Null); + assert_eq!(body["messages"], json!([{"role": "user", "content": "hi"}])); + assert!(element.key.ends_with("_chatcmpl-1.json")); + } + + #[tokio::test] + async fn only_the_listed_outcomes_are_queued() { + let settings = config::tests::settings_with( + "litellm_settings:\n failure_callback: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n", + &[], + ) + .await + .unwrap() + .unwrap(); + let mut log = log("").await; + log.settings = Arc::new(settings); + log.observe(&record()).await; + assert!(log.queue.lock().unwrap().elements.is_empty()); + let mut failed = record(); + failed.outcome.status = 500; + log.observe(&failed).await; + assert_eq!(log.queue.lock().unwrap().elements.len(), 1); + } + + #[tokio::test] + async fn a_full_queue_drops_the_oldest() { + let log = log(" s3_max_queue_size: 2\n").await; + for key in ["a", "b", "c"] { + log.enqueue(element(key, "{}")); + } + let queue = log.queue.lock().unwrap(); + let keys: Vec<_> = queue + .elements + .iter() + .map(|element| element.key.as_str()) + .collect(); + assert_eq!(keys, ["b", "c"]); + } + + #[tokio::test] + async fn an_unreachable_bucket_keeps_everything_for_the_next_flush() { + let log = log(" s3_endpoint_url: http://127.0.0.1:9\n s3_aws_access_key_id: k\n s3_aws_secret_access_key: s\n s3_max_retry_age_seconds: 1\n").await; + let mut old = element("a", "{}"); + old.retrying_since = Some(Instant::now() - Duration::from_secs(60)); + log.enqueue(old); + log.enqueue(element("b", "{}")); + log.flush().await; + let queue = log.queue.lock().unwrap(); + let kept: Vec<_> = queue + .elements + .iter() + .map(|element| element.key.as_str()) + .collect(); + assert_eq!(kept, ["a", "b"]); + assert_eq!(queue.requeued, 2); + } +} diff --git a/litellm-rust/crates/gateway/src/s3_log/payload.rs b/litellm-rust/crates/gateway/src/s3_log/payload.rs new file mode 100644 index 00000000000..8b67803dfb5 --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/payload.rs @@ -0,0 +1,1204 @@ +//! The proxy's `StandardLoggingPayload` for a request, from what the gateway knows of it, and +//! the transforms the proxy applies before logging it: base64 truncation, message redaction, +//! `s3_strip_base64_files` and `s3_log_prompts_only`. + +use std::{ + collections::BTreeMap, + time::{SystemTime, UNIX_EPOCH}, +}; + +use litellm_gateway_hooks::{Record, Route}; +use serde_json::{Map, Value, json}; + +use crate::{ + prometheus::error_class, + spend::{logged_model, spend_logs_metadata, usage_values}, +}; + +const REDACTED: &str = "redacted-by-litellm"; +/// The proxy's `MAX_BASE64_LENGTH_FOR_LOGGING`: longer data URIs are logged as their size. +const MAX_BASE64_LENGTH: usize = 64; +const MAX_TRUNCATION_DEPTH: usize = 20; +/// `DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER`, the depth `s3_strip_base64_files` walks. +const MAX_STRIP_DEPTH: usize = 10; + +/// The `StandardLoggingMetadata` keys, every one present and null when unset. +const METADATA_KEYS: [&str; 35] = [ + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_spend", + "user_api_key_max_budget", + "user_api_key_budget_reset_at", + "user_api_key_user_spend", + "user_api_key_user_max_budget", + "user_api_key_team_spend", + "user_api_key_team_max_budget", + "user_api_key_team_id", + "user_api_key_org_id", + "user_api_key_org_alias", + "user_api_key_project_id", + "user_api_key_project_alias", + "user_api_key_user_id", + "user_api_key_team_alias", + "user_api_key_user_email", + "user_api_key_end_user_id", + "user_api_key_request_route", + "spend_logs_metadata", + "requester_ip_address", + "user_agent", + "requester_metadata", + "prompt_management_metadata", + "applied_guardrails", + "mcp_tool_call_metadata", + "vector_store_request_metadata", + "routing_decision", + "usage_object", + "requester_custom_headers", + "cold_storage_object_key", + "user_api_key_auth_metadata", + "team_alias", + "team_id", + "used_client_oauth_token", +]; + +/// What the proxy's settings change about every payload. +#[derive(Clone, Copy, Debug, Default)] +pub(super) struct Redaction { + /// `litellm_settings.turn_off_message_logging`. + pub messages: bool, + /// `litellm_settings.redact_user_api_key_info`. + pub user_api_key_info: bool, + /// The proxy's request tags carry the client's user agent unless + /// `disable_add_user_agent_to_request_tags` is set. + pub user_agent_tags: bool, +} + +/// The request's `StandardLoggingPayload`. +pub(super) fn standard_logging_payload( + record: &Record, + redaction: Redaction, +) -> Map { + let request = &record.request; + let outcome = &record.outcome; + let deployment = outcome.deployment.as_ref(); + let succeeded = record.succeeded(); + let body: Map = serde_json::from_slice(request.body()).unwrap_or_default(); + let response = match succeeded { + true => response(record), + false => Value::Object(Map::new()), + }; + let id = response + .get("id") + .and_then(Value::as_str) + .map_or_else(|| request.request_id.to_string(), str::to_owned); + let usage = outcome.usage.unwrap_or_default(); + let start = seconds(request.started_at); + let end = seconds(outcome.ended_at); + let completion_start = match (request.stream, outcome.time_to_first_byte) { + (true, Some(first_byte)) => seconds(request.started_at + first_byte), + _ => end, + }; + let status = match succeeded { + true => "success", + false => "failure", + }; + let model = logged_model(record); + let api_base = deployment + .and_then(|deployment| deployment.api_base.as_deref()) + .map(strip_trailing_slash) + .unwrap_or_default(); + let user_agent = request + .client_headers + .get(http::header::USER_AGENT) + .and_then(|value| value.to_str().ok()); + let error_message = (!succeeded) + .then(|| error_message(record)) + .flatten() + .unwrap_or_default(); + let messages = body.get("messages").map(|messages| { + let messages = with_system_prompt(messages.clone(), body.get("system")); + truncate_base64(messages) + }); + let hidden_params = json!({ + "model_id": deployment.map(|deployment| &*deployment.id), + "cache_key": null, + "api_base": deployment.and_then(|deployment| deployment.api_base.as_deref()), + "response_cost": outcome.cost, + "litellm_overhead_time_ms": null, + "additional_headers": null, + "batch_models": null, + "batch_successful_requests": null, + "batch_failed_requests": null, + "litellm_model_name": deployment.map(|_| &model), + "usage_object": null, + }); + let error_information = json!({ + "error_code": if succeeded { String::new() } else { outcome.status.to_string() }, + "error_class": if succeeded { "" } else { error_class(record) }, + "llm_provider": deployment + .filter(|_| !succeeded) + .and_then(|deployment| deployment.provider()) + .unwrap_or_default(), + "traceback": "", + "error_message": &error_message, + }); + let mut payload = json!({ + "id": id, + "litellm_call_id": &*request.request_id, + "trace_id": &*request.request_id, + "session_id": "", + "call_type": request.route.call_type(), + "cache_hit": if request.stream { Value::Bool(false) } else { Value::Null }, + "stream": if request.stream && succeeded { Value::Bool(true) } else { Value::Null }, + "status": status, + "status_fields": {"llm_api_status": status, "guardrail_status": "not_run"}, + "custom_llm_provider": deployment.and_then(|deployment| deployment.provider()), + "saved_cache_cost": 0.0, + "startTime": start, + "endTime": end, + "completionStartTime": completion_start, + "response_time": match request.stream { + true => completion_start - start, + false => end - start, + }, + "model": &model, + "metadata": metadata(record, &body, user_agent, redaction), + "cache_key": null, + "response_cost": outcome.cost.unwrap_or(0.0), + "cost_breakdown": null, + "autorouter_savings": null, + "autorouter_savings_estimate": null, + "autorouter_baseline_observation": null, + }); + let rest = json!({ + "total_tokens": usage.total_tokens(), + "prompt_tokens": usage.prompt_tokens(), + "completion_tokens": usage.output_tokens, + "request_tags": request_tags(&body, user_agent.filter(|_| redaction.user_agent_tags)), + "request_model_access_groups": [], + "end_user": request.end_user.as_deref(), + "api_base": api_base, + "model_group": deployment.map_or("", |deployment| &deployment.model_group), + "model_id": deployment.map_or("", |deployment| &deployment.id), + "requester_ip_address": null, + "user_agent": user_agent, + "messages": messages, + "response": response, + "model_parameters": model_parameters(&body), + "hidden_params": hidden_params, + "model_map_information": {"model_map_key": &model, "model_map_value": null}, + "error_str": (!succeeded && !error_message.is_empty()).then_some(&error_message), + "error_information": error_information, + "response_cost_failure_debug_info": null, + "zero_cost_diagnostic": null, + "guardrail_information": null, + "standard_built_in_tools_params": null, + }); + if let (Value::Object(payload), Value::Object(rest)) = (&mut payload, rest) { + payload.extend(rest); + } + if redaction.messages && message_logging_off(record) { + redact(&mut payload); + } + match payload { + Value::Object(payload) => payload, + _ => unreachable!("the payload is built as an object"), + } +} + +fn seconds(time: SystemTime) -> f64 { + time.duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64() +} + +fn strip_trailing_slash(api_base: &str) -> &str { + match api_base.ends_with("//") { + true => api_base.trim_end_matches('/'), + false => api_base.strip_suffix('/').unwrap_or(api_base), + } +} + +/// The proxy's `append_system_prompt_messages`: a string `system` becomes the first message. +fn with_system_prompt(messages: Value, system: Option<&Value>) -> Value { + let Some(Value::String(system)) = system else { + return messages; + }; + let system_message = json!({"role": "system", "content": system}); + match messages { + Value::Null => Value::Array(vec![system_message]), + Value::Array(messages) => { + if messages.first().is_some_and(|first| { + first.get("role") == Some(&json!("system")) + && first.get("content") == Some(&json!(system)) + }) { + return Value::Array(messages); + } + Value::Array(std::iter::once(system_message).chain(messages).collect()) + } + Value::String(text) => json!([system_message, {"role": "user", "content": text}]), + other => other, + } +} + +/// The proxy redacts unless the client opts out with `litellm-disable-message-redaction`. +fn message_logging_off(record: &Record) -> bool { + !record + .request + .client_headers + .get("litellm-disable-message-redaction") + .is_some_and(|value| !value.is_empty()) +} + +/// The proxy's `_redact_standard_logging_object`. +fn redact(payload: &mut Value) { + if !payload["messages"].is_null() { + payload["messages"] = json!([{"role": "user", "content": REDACTED}]); + } + let response = &mut payload["response"]; + match response { + Value::Null => {} + Value::Object(fields) if fields.get("choices").is_some() => { + if let Some(Value::Array(choices)) = fields.get_mut("choices") { + choices.iter_mut().for_each(redact_choice); + } + } + Value::String(_) => *response = Value::from(REDACTED), + _ => *response = json!({"text": REDACTED}), + } +} + +fn redact_choice(choice: &mut Value) { + let Value::Object(choice) = choice else { + return; + }; + let message = match (choice.contains_key("message"), choice.contains_key("delta")) { + (true, _) => choice.get_mut("message"), + (false, true) => choice.get_mut("delta"), + (false, false) => { + if choice.get("text").is_some_and(|text| !text.is_null()) { + choice.insert("text".into(), Value::from(REDACTED)); + } + return; + } + }; + let Some(Value::Object(message)) = message else { + return; + }; + for field in ["content", "reasoning_content"] { + if message.get(field).is_some_and(|value| !value.is_null()) { + message.insert(field.into(), Value::from(REDACTED)); + } + } + for field in ["thinking_blocks", "audio"] { + if message.contains_key(field) { + message.insert(field.into(), Value::Null); + } + } + if let Some(Value::Array(calls)) = message.get_mut("tool_calls") { + for call in calls { + if let Some(Value::Object(function)) = call.get_mut("function") { + function.insert("arguments".into(), Value::from(REDACTED)); + } + } + } + if let Some(Value::Object(function)) = message.get_mut("function_call") + && function.contains_key("arguments") + { + function.insert("arguments".into(), Value::from(REDACTED)); + } +} + +fn metadata( + record: &Record, + body: &Map, + user_agent: Option<&str>, + redaction: Redaction, +) -> Value { + let mut metadata: Map = METADATA_KEYS + .iter() + .map(|key| ((*key).to_owned(), Value::Null)) + .collect(); + let mut set = |key: &str, value: Value| { + metadata.insert(key.to_owned(), value); + }; + let request = &record.request; + if let Some(key) = record.key().filter(|_| !redaction.user_api_key_info) { + let budgets = &key.budgets; + set("user_api_key_hash", json!(key.hash)); + set("user_api_key_alias", json!(key.alias)); + set("user_api_key_spend", json!(budgets.key.spend)); + set("user_api_key_max_budget", json!(budgets.key.max_budget)); + set( + "user_api_key_user_spend", + json!(budgets.user.map(|user| user.spend)), + ); + set( + "user_api_key_user_max_budget", + json!(budgets.user.and_then(|user| user.max_budget)), + ); + set( + "user_api_key_team_spend", + json!(budgets.team.map(|team| team.spend)), + ); + set( + "user_api_key_team_max_budget", + json!(budgets.team.and_then(|team| team.max_budget)), + ); + set("user_api_key_team_id", json!(key.team_id)); + set("user_api_key_org_id", json!(key.org_id)); + set("user_api_key_user_id", json!(key.user_id)); + set("user_api_key_team_alias", json!(key.team_alias)); + } + if !redaction.user_api_key_info { + set( + "user_api_key_end_user_id", + json!(request.end_user.as_deref()), + ); + set( + "user_api_key_request_route", + json!(request.route.endpoint()), + ); + } + set("spend_logs_metadata", spend_logs_metadata(record)); + set("user_agent", json!(user_agent)); + if let Some(requester) = body.get("metadata").filter(|value| value.is_object()) { + set("requester_metadata", requester.clone()); + } + set( + "usage_object", + record + .outcome + .usage + .map_or_else(|| json!({}), |usage| usage_values(&usage).0), + ); + Value::Object(metadata) +} + +/// The request's `metadata.tags`, then the proxy's user agent tags. +fn request_tags(body: &Map, user_agent: Option<&str>) -> Vec { + let tags = body + .get("metadata") + .and_then(|metadata| metadata.get("tags")) + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let agent = user_agent.into_iter().flat_map(|agent| { + agent + .split_once('/') + .map(|(product, _)| format!("User-Agent: {product}")) + .into_iter() + .chain(std::iter::once(format!("User-Agent: {agent}"))) + .map(Value::from) + }); + tags.into_iter().chain(agent).collect() +} + +/// The request's model parameters: everything but the model, the prompt and the metadata. +fn model_parameters(body: &Map) -> Map { + body.iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "model" | "messages" | "system" | "prompt" | "input" | "metadata" + ) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} + +/// The provider's or the gateway's error message, from an error body when there is one. +fn error_message(record: &Record) -> Option { + let body: Value = serde_json::from_slice(record.response_body.as_deref()?).ok()?; + let error = body.get("error").unwrap_or(&body); + error + .get("message") + .and_then(Value::as_str) + .or_else(|| error.as_str()) + .map(str::to_owned) +} + +/// The response the client got: the body itself, or the stream put back together the way the +/// proxy assembles it. +fn response(record: &Record) -> Value { + let Some(body) = record.response_body.as_deref() else { + return Value::Object(Map::new()); + }; + if let Ok(value @ Value::Object(_)) = serde_json::from_slice::(body) { + return value; + } + let Ok(text) = std::str::from_utf8(body) else { + return Value::Object(Map::new()); + }; + let events = text + .lines() + .filter_map(|line| line.strip_prefix("data:")) + .map(str::trim) + .filter_map(|data| serde_json::from_str::(data).ok()); + let assembled = match record.request.route { + Route::ChatCompletions => chat_completion(events), + Route::Messages => anthropic_message(events), + }; + assembled.unwrap_or_else(|| Value::Object(Map::new())) +} + +#[derive(Default)] +struct ChoiceParts { + role: Option, + content: Option, + reasoning: Option, + finish_reason: Value, + tool_calls: BTreeMap, +} + +/// A chat completion from its chunks. +fn chat_completion(events: impl Iterator) -> Option { + let mut head: Option = None; + let mut usage = Value::Null; + let mut choices: BTreeMap = BTreeMap::new(); + for chunk in events { + if head.is_none() { + head = Some(json!({ + "id": chunk.get("id"), + "created": chunk.get("created"), + "model": chunk.get("model"), + "system_fingerprint": chunk.get("system_fingerprint"), + })); + } + if chunk.get("usage").is_some_and(|usage| !usage.is_null()) { + usage = chunk["usage"].clone(); + } + for choice in chunk["choices"].as_array().into_iter().flatten() { + let parts = choices + .entry(choice["index"].as_u64().unwrap_or(0)) + .or_default(); + let delta = &choice["delta"]; + if let Some(role) = delta.get("role").filter(|role| !role.is_null()) { + parts.role = Some(role.clone()); + } + if let Some(text) = delta["content"].as_str() { + parts.content.get_or_insert_default().push_str(text); + } + if let Some(text) = delta["reasoning_content"].as_str() { + parts.reasoning.get_or_insert_default().push_str(text); + } + if !choice["finish_reason"].is_null() { + parts.finish_reason = choice["finish_reason"].clone(); + } + for call in delta["tool_calls"].as_array().into_iter().flatten() { + let slot = parts + .tool_calls + .entry(call["index"].as_u64().unwrap_or(0)) + .or_insert_with(|| { + json!({"id": null, "type": "function", "function": {"name": null, "arguments": ""}}) + }); + for field in ["id", "type"] { + if let Some(value) = call.get(field).filter(|value| !value.is_null()) { + slot[field] = value.clone(); + } + } + if let Some(name) = call["function"]["name"].as_str() { + slot["function"]["name"] = Value::from(name); + } + if let Some(arguments) = call["function"]["arguments"].as_str() { + let joined = format!( + "{}{arguments}", + slot["function"]["arguments"].as_str().unwrap_or_default() + ); + slot["function"]["arguments"] = Value::from(joined); + } + } + } + } + let mut completion = head?; + completion["object"] = Value::from("chat.completion"); + completion["choices"] = choices + .into_iter() + .map(|(index, parts)| { + let tool_calls: Vec = parts.tool_calls.into_values().collect(); + let mut message = json!({ + "content": parts.content, + "role": parts.role.unwrap_or_else(|| Value::from("assistant")), + "tool_calls": if tool_calls.is_empty() { Value::Null } else { Value::Array(tool_calls) }, + "function_call": null, + }); + if let Some(reasoning) = parts.reasoning { + message["reasoning_content"] = Value::from(reasoning); + } + json!({"finish_reason": parts.finish_reason, "index": index, "message": message}) + }) + .collect(); + completion["usage"] = usage; + Some(completion) +} + +/// An Anthropic message from its stream events. +fn anthropic_message(events: impl Iterator) -> Option { + let mut message: Option = None; + let mut blocks: BTreeMap = BTreeMap::new(); + let mut partial_json: BTreeMap = BTreeMap::new(); + for event in events { + let index = event["index"].as_u64().unwrap_or(0); + match event["type"].as_str() { + Some("message_start") => message = Some(event["message"].clone()), + Some("content_block_start") => { + blocks.insert(index, event["content_block"].clone()); + } + Some("content_block_delta") => { + let block = blocks.entry(index).or_insert_with(|| json!({})); + let delta = &event["delta"]; + let append = |block: &mut Value, field: &str, text: &str| { + let joined = format!("{}{text}", block[field].as_str().unwrap_or_default()); + block[field] = Value::from(joined); + }; + match delta["type"].as_str() { + Some("text_delta") => { + append(block, "text", delta["text"].as_str().unwrap_or_default()) + } + Some("thinking_delta") => append( + block, + "thinking", + delta["thinking"].as_str().unwrap_or_default(), + ), + Some("signature_delta") => block["signature"] = delta["signature"].clone(), + Some("input_json_delta") => partial_json + .entry(index) + .or_default() + .push_str(delta["partial_json"].as_str().unwrap_or_default()), + _ => {} + } + } + Some("message_delta") => { + let Some(message) = message.as_mut() else { + continue; + }; + if let Some(delta) = event["delta"].as_object() { + for (field, value) in delta { + message[field.as_str()] = value.clone(); + } + } + if let (Some(usage), Some(Value::Object(total))) = + (event["usage"].as_object(), message.get_mut("usage")) + { + for (field, value) in usage { + if !value.is_null() { + total.insert(field.clone(), value.clone()); + } + } + } + } + _ => {} + } + } + let mut message = message?; + for (index, json) in partial_json { + if let (Some(block), Ok(input)) = (blocks.get_mut(&index), serde_json::from_str(&json)) { + block["input"] = input; + } + } + message["content"] = Value::Array(blocks.into_values().collect()); + Some(message) +} + +/// Replaces each data URI in `text` that `matcher` accepts; `matcher` gets the text from just +/// after `data:` and returns how much of it the URI spans and its replacement. +fn replace_data_uris(text: &str, matcher: impl Fn(&str) -> Option<(usize, String)>) -> String { + let mut out = String::with_capacity(text.len()); + let mut rest = text; + while let Some(at) = rest.find("data:") { + let after = &rest[at + 5..]; + match matcher(after) { + Some((length, replacement)) => { + out.push_str(&rest[..at]); + out.push_str(&replacement); + rest = &after[length..]; + } + None => { + out.push_str(&rest[..at + 5]); + rest = after; + } + } + } + out.push_str(rest); + out +} + +/// The length of the leading run of `text` whose characters satisfy `accept`. +fn run(text: &str, accept: impl Fn(char) -> bool) -> usize { + text.find(|c: char| !accept(c)).unwrap_or(text.len()) +} + +fn is_base64(c: char) -> bool { + c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=') +} + +/// The proxy's `truncate_base64_in_messages`: a data URI with more than 64 base64 characters +/// is logged as its size. +fn truncate_base64(value: Value) -> Value { + fn walk(value: Value, depth: usize) -> Value { + match value { + Value::String(text) => Value::String(truncate_text(&text)), + Value::Array(items) if depth <= MAX_TRUNCATION_DEPTH => Value::Array( + items + .into_iter() + .map(|item| match item { + Value::String(_) => walk(item, depth), + item => walk(item, depth + 1), + }) + .collect(), + ), + Value::Object(fields) if depth <= MAX_TRUNCATION_DEPTH => Value::Object( + fields + .into_iter() + .map(|(key, item)| match item { + Value::String(_) => (key, walk(item, depth)), + item => (key, walk(item, depth + 1)), + }) + .collect(), + ), + other => other, + } + } + walk(value, 0) +} + +fn truncate_text(text: &str) -> String { + replace_data_uris(text, |after| { + let mime = run(after, |c| !matches!(c, ';' | ',') && !c.is_whitespace()); + if mime == 0 || mime > 255 || !after[mime..].starts_with(";base64,") { + return None; + } + let payload_start = mime + ";base64,".len(); + let payload = run(&after[payload_start..], is_base64); + if payload == 0 { + return None; + } + let length = payload_start + payload; + if payload <= MAX_BASE64_LENGTH { + return Some((length, format!("data:{}", &after[..length]))); + } + Some(( + length, + format!( + "data:{};base64,[base64_data truncated: {}]", + &after[..mime], + base64_size(payload) + ), + )) + }) +} + +/// The proxy's `format_base64_size`. +fn base64_size(characters: usize) -> String { + let bytes = characters as f64 * 3.0 / 4.0; + match bytes { + bytes if bytes >= 1024.0 * 1024.0 => format!("{:.2}MB", bytes / (1024.0 * 1024.0)), + bytes if bytes >= 1024.0 => format!("{:.1}KB", bytes / 1024.0), + bytes => format!("{}B", bytes as u64), + } +} + +/// `s3_strip_base64_files`, as the proxy's `_strip_base64_from_messages_sync`: drops file and +/// non-text content parts and redacts inline base64 in what is left of each message. +pub(super) fn strip_base64_files(payload: &mut Map) { + let Some(Value::Array(messages)) = payload.get_mut("messages") else { + return; + }; + if messages.is_empty() { + return; + } + let stripped: Vec = std::mem::take(messages) + .into_iter() + .filter_map(|message| match message { + Value::Object(message) => Some(Value::Object(strip_message(message))), + _ => None, + }) + .collect(); + *messages = stripped; +} + +fn strip_message(mut message: Map) -> Map { + let content = match message.remove("content") { + Some(Value::Array(parts)) => Value::Array( + parts + .into_iter() + .filter(keep_part) + .map(|part| redact_base64(part, 0)) + .collect(), + ), + other => redact_base64(other.unwrap_or(Value::Null), 0), + }; + let mut stripped: Map = message + .into_iter() + .map(|(key, value)| (key, redact_base64(value, 0))) + .collect(); + stripped.insert("content".into(), content); + stripped +} + +fn keep_part(part: &Value) -> bool { + let Value::Object(part) = part else { + return true; + }; + !part.contains_key("file") + && part + .get("type") + .and_then(Value::as_str) + .is_none_or(|kind| kind == "text") +} + +fn redact_base64(value: Value, depth: usize) -> Value { + if depth > MAX_STRIP_DEPTH { + return Value::from("[MAX_DEPTH_REACHED]"); + } + match value { + Value::String(text) => Value::String(redact_text(&text)), + Value::Array(items) => Value::Array( + items + .into_iter() + .map(|item| redact_base64(item, depth + 1)) + .collect(), + ), + Value::Object(fields) => Value::Object( + fields + .into_iter() + .map(|(key, item)| (key, redact_base64(item, depth + 1))) + .collect(), + ), + other => other, + } +} + +/// Replaces `data:(application|image|audio|video)/;base64,` with +/// `[BASE64_REDACTED]`, the payload running over base64 characters and whitespace. +fn redact_text(text: &str) -> String { + replace_data_uris(text, |after| { + let kind = ["application/", "image/", "audio/", "video/"] + .into_iter() + .find(|kind| after.starts_with(kind))?; + let subtype = run(&after[kind.len()..], |c| { + c.is_ascii_alphanumeric() || matches!(c, '.' | '+' | '-') + }); + let rest = &after[kind.len() + subtype..]; + if subtype == 0 || !rest.starts_with(";base64,") { + return None; + } + let payload_start = kind.len() + subtype + ";base64,".len(); + let payload = run(&after[payload_start..], |c| { + is_base64(c) || c.is_whitespace() + }); + (payload > 0).then(|| (payload_start + payload, "[BASE64_REDACTED]".to_owned())) + }) +} + +/// `s3_log_prompts_only`: the payload without the response. +pub(super) fn prompts_only(payload: &mut Map) { + payload.insert("response".into(), Value::Null); +} + +#[cfg(test)] +pub(super) mod tests { + use std::{sync::Arc, time::Duration}; + + use axum::body::Bytes; + use http::HeaderMap; + use litellm_gateway_hooks::{ + Annotations, DeploymentFacts, KeyFacts, Outcome, RequestFacts, Usage, + }; + use rstest::rstest; + + use super::*; + + pub(crate) const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-x","choices":[{"index":0,"message":{"role":"assistant","content":"hi","tool_calls":[{"id":"t1","type":"function","function":{"name":"f","arguments":"{\"a\":1}"}}]},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":6,"total_tokens":11}}"#; + + pub(crate) fn record( + route: Route, + status: u16, + stream: bool, + request: Value, + response: Option<&str>, + ) -> Record { + let started_at = UNIX_EPOCH + Duration::from_secs(1_000); + let mut headers = HeaderMap::new(); + headers.insert("user-agent", "client/1.0".parse().unwrap()); + let key = KeyFacts { + hash: "hash-1".into(), + alias: Some("team-key".into()), + user_id: Some("user-1".into()), + team_id: Some("team-1".into()), + team_alias: Some("research".into()), + ..KeyFacts::default() + }; + let mut facts = RequestFacts::new( + "req-1".into(), + route, + Some(Arc::new(key)), + Some("end-1".into()), + "gpt".into(), + Arc::new(headers), + stream, + None, + Bytes::from(request.to_string()), + ); + facts.started_at = started_at; + Record { + request: Arc::new(facts), + outcome: Outcome { + deployment: Some(DeploymentFacts { + id: "gpt-1".into(), + model_group: "gpt".into(), + model: "openai/gpt-x".into(), + provider: None, + api_base: Some("https://api.example/v1/".into()), + }), + status, + usage: (status == 200).then_some(Usage { + input_tokens: 5, + output_tokens: 6, + ..Usage::default() + }), + cost: (status == 200).then_some(0.5), + time_to_first_byte: Some(Duration::from_millis(250)), + upstream_started_at: None, + ended_at: started_at + Duration::from_secs(2), + interrupted: false, + }, + annotations: Annotations::default(), + response_body: response.map(|body| Bytes::from(body.to_owned())), + } + } + + fn chat_request() -> Value { + json!({ + "model": "gpt", + "temperature": 0.2, + "metadata": {"tags": ["prod"]}, + "messages": [{"role": "user", "content": "hi"}], + }) + } + + #[test] + fn a_success_carries_the_proxys_fields() { + let record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let payload = Value::Object(standard_logging_payload( + &record, + Redaction { + user_agent_tags: true, + ..Redaction::default() + }, + )); + assert_eq!(payload["id"], "chatcmpl-1"); + assert_eq!(payload["litellm_call_id"], "req-1"); + assert_eq!(payload["call_type"], "acompletion"); + assert_eq!(payload["status"], "success"); + assert_eq!(payload["custom_llm_provider"], "openai"); + assert_eq!(payload["model"], "gpt-x"); + assert_eq!(payload["model_group"], "gpt"); + assert_eq!(payload["model_id"], "gpt-1"); + assert_eq!(payload["api_base"], "https://api.example/v1"); + assert_eq!(payload["response_cost"], 0.5); + assert_eq!(payload["prompt_tokens"], 5); + assert_eq!(payload["completion_tokens"], 6); + assert_eq!(payload["total_tokens"], 11); + assert_eq!(payload["startTime"], 1000.0); + assert_eq!(payload["endTime"], 1002.0); + assert_eq!(payload["completionStartTime"], 1002.0); + assert_eq!(payload["response_time"], 2.0); + assert_eq!(payload["stream"], Value::Null); + assert_eq!(payload["end_user"], "end-1"); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": "hi"}]) + ); + assert_eq!( + payload["response"], + serde_json::from_str::(COMPLETION).unwrap() + ); + assert_eq!(payload["model_parameters"], json!({"temperature": 0.2})); + assert_eq!( + payload["request_tags"], + json!(["prod", "User-Agent: client", "User-Agent: client/1.0"]) + ); + let metadata = &payload["metadata"]; + assert_eq!(metadata.as_object().unwrap().len(), METADATA_KEYS.len()); + assert_eq!(metadata["user_api_key_hash"], "hash-1"); + assert_eq!(metadata["user_api_key_alias"], "team-key"); + assert_eq!(metadata["user_api_key_team_alias"], "research"); + assert_eq!(metadata["user_api_key_team_id"], "team-1"); + assert_eq!(metadata["user_api_key_user_id"], "user-1"); + assert_eq!(metadata["user_api_key_end_user_id"], "end-1"); + assert_eq!(metadata["requester_metadata"], json!({"tags": ["prod"]})); + assert_eq!(metadata["usage_object"]["total_tokens"], 11); + assert_eq!(payload["error_information"]["error_code"], ""); + assert_eq!(payload["error_str"], Value::Null); + } + + #[test] + fn a_failure_carries_the_error_and_an_empty_response() { + let error = r#"{"error":{"message":"slow down","type":"rate_limit_error"}}"#; + let record = record( + Route::ChatCompletions, + 429, + false, + chat_request(), + Some(error), + ); + let payload = Value::Object(standard_logging_payload(&record, Redaction::default())); + assert_eq!(payload["status"], "failure"); + assert_eq!(payload["id"], "req-1"); + assert_eq!(payload["response"], json!({})); + assert_eq!(payload["response_cost"], 0.0); + assert_eq!(payload["total_tokens"], 0); + assert_eq!(payload["error_str"], "slow down"); + assert_eq!( + payload["error_information"], + json!({ + "error_code": "429", + "error_class": "RateLimitError", + "llm_provider": "openai", + "traceback": "", + "error_message": "slow down", + }) + ); + assert_eq!(payload["metadata"]["usage_object"], json!({})); + assert_eq!(payload["request_tags"], json!(["prod"])); + } + + #[test] + fn a_chat_stream_is_put_back_together() { + let stream = concat!( + "data: {\"id\":\"chatcmpl-9\",\"created\":7,\"model\":\"gpt-x\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"he\"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-9\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"llo\",\"tool_calls\":[{\"index\":0,\"id\":\"t1\",\"function\":{\"name\":\"f\",\"arguments\":\"{\\\"a\\\"\"}}]},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-9\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\":1}\"}}]},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n", + "data: [DONE]\n\n", + ); + let record = record( + Route::ChatCompletions, + 200, + true, + chat_request(), + Some(stream), + ); + let payload = Value::Object(standard_logging_payload(&record, Redaction::default())); + assert_eq!(payload["id"], "chatcmpl-9"); + assert_eq!(payload["stream"], true); + assert_eq!(payload["cache_hit"], false); + assert_eq!(payload["completionStartTime"], 1000.25); + assert_eq!(payload["response_time"], 0.25); + let response = &payload["response"]; + assert_eq!(response["object"], "chat.completion"); + assert_eq!(response["created"], 7); + assert_eq!(response["usage"]["total_tokens"], 3); + assert_eq!( + response["choices"], + json!([{ + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": "hello", + "role": "assistant", + "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "f", "arguments": "{\"a\":1}"}}], + "function_call": null, + }, + }]) + ); + } + + #[test] + fn a_messages_stream_is_put_back_together_with_its_system_prompt() { + let stream = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-x\",\"stop_reason\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":1}}}\n\n", + "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n", + "data: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t\",\"name\":\"f\",\"input\":{}}}\n\n", + "data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"a\\\":\"}}\n\n", + "data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"1}\"}}\n\n", + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":42}}\n\n", + ); + let request = json!({"model": "claude", "system": "be brief", "messages": [{"role": "user", "content": "hi"}]}); + let record = record(Route::Messages, 200, true, request, Some(stream)); + let payload = Value::Object(standard_logging_payload(&record, Redaction::default())); + assert_eq!(payload["id"], "msg_1"); + assert_eq!(payload["call_type"], "anthropic_messages"); + assert_eq!( + payload["messages"], + json!([{"role": "system", "content": "be brief"}, {"role": "user", "content": "hi"}]) + ); + let response = &payload["response"]; + assert_eq!(response["stop_reason"], "tool_use"); + assert_eq!( + response["usage"], + json!({"input_tokens": 10, "output_tokens": 42}) + ); + assert_eq!( + response["content"], + json!([{"type": "text", "text": "hi"}, {"type": "tool_use", "id": "t", "name": "f", "input": {"a": 1}}]) + ); + } + + #[test] + fn turning_off_message_logging_redacts_messages_and_response() { + let record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let redaction = Redaction { + messages: true, + ..Redaction::default() + }; + let payload = Value::Object(standard_logging_payload(&record, redaction)); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": REDACTED}]) + ); + let message = &payload["response"]["choices"][0]["message"]; + assert_eq!(message["content"], REDACTED); + assert_eq!(message["tool_calls"][0]["function"]["arguments"], REDACTED); + assert_eq!(message["tool_calls"][0]["function"]["name"], "f"); + assert_eq!(payload["response"]["usage"]["total_tokens"], 11); + } + + #[test] + fn a_messages_response_is_redacted_to_text() { + let message = + r#"{"id":"msg_1","type":"message","content":[{"type":"text","text":"secret"}]}"#; + let record = record( + Route::Messages, + 200, + false, + json!({"messages": []}), + Some(message), + ); + let redaction = Redaction { + messages: true, + ..Redaction::default() + }; + let payload = Value::Object(standard_logging_payload(&record, redaction)); + assert_eq!(payload["response"], json!({"text": REDACTED})); + assert_eq!(payload["id"], "msg_1"); + } + + #[test] + fn the_client_can_opt_out_of_redaction() { + let mut record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let mut headers = HeaderMap::new(); + headers.insert("litellm-disable-message-redaction", "true".parse().unwrap()); + let request = Arc::get_mut(&mut record.request).unwrap(); + request.client_headers = Arc::new(headers); + let redaction = Redaction { + messages: true, + ..Redaction::default() + }; + let payload = Value::Object(standard_logging_payload(&record, redaction)); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": "hi"}]) + ); + } + + #[test] + fn redacting_key_info_leaves_the_fields_null() { + let record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let redaction = Redaction { + user_api_key_info: true, + ..Redaction::default() + }; + let payload = Value::Object(standard_logging_payload(&record, redaction)); + assert_eq!(payload["metadata"]["user_api_key_team_alias"], Value::Null); + assert_eq!(payload["metadata"]["user_api_key_hash"], Value::Null); + } + + #[rstest] + #[case::short("data:image/png;base64,AAAA", "data:image/png;base64,AAAA")] + #[case::long( + &format!("see data:image/png;base64,{} end", "A".repeat(4000)), + "see data:image/png;base64,[base64_data truncated: 2.9KB] end" + )] + #[case::bytes(&format!("data:audio/wav;base64,{}", "A".repeat(100)), "data:audio/wav;base64,[base64_data truncated: 75B]")] + #[case::not_base64("data:text/plain,hello", "data:text/plain,hello")] + fn long_data_uris_in_messages_are_logged_as_their_size( + #[case] text: &str, + #[case] expected: &str, + ) { + assert_eq!( + truncate_base64(json!([{"content": [{"image_url": {"url": text}}]}])), + json!([{"content": [{"image_url": {"url": expected}}]}]) + ); + } + + #[test] + fn stripping_base64_drops_file_parts_and_redacts_inline_data() { + let mut payload = Map::new(); + payload.insert( + "messages".into(), + json!([ + {"role": "user", "content": [ + {"type": "text", "text": "look data:image/png;base64,QUJD"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,QUJD"}}, + {"file": {"file_data": "x"}}, + ]}, + {"role": "assistant", "content": "data:application/pdf;base64,QUJD\nQUJD"}, + {"role": "tool"}, + "not a message", + ]), + ); + strip_base64_files(&mut payload); + assert_eq!( + payload["messages"], + json!([ + {"role": "user", "content": [{"type": "text", "text": "look [BASE64_REDACTED]"}]}, + {"role": "assistant", "content": "[BASE64_REDACTED]"}, + {"role": "tool", "content": null}, + ]) + ); + } + + #[test] + fn prompts_only_drops_the_response() { + let record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let mut payload = standard_logging_payload(&record, Redaction::default()); + prompts_only(&mut payload); + assert_eq!(payload["response"], Value::Null); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": "hi"}]) + ); + assert_eq!(payload["id"], "chatcmpl-1"); + } + + #[test] + fn without_a_response_body_the_id_is_the_request_id() { + let record = record(Route::ChatCompletions, 200, false, chat_request(), None); + let payload = standard_logging_payload(&record, Redaction::default()); + assert_eq!(payload["id"], "req-1"); + assert_eq!(payload["response"], json!({})); + } +} diff --git a/litellm-rust/crates/gateway/src/s3_log/upload.rs b/litellm-rust/crates/gateway/src/s3_log/upload.rs new file mode 100644 index 00000000000..e78f0d89196 --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/upload.rs @@ -0,0 +1,278 @@ +//! One signed PUT per object, retried as the proxy's `_put_with_retries` retries it: three +//! attempts, a transport error or a 403, 429, 500 or 503 retried after a jittered backoff, and +//! a 400 or 403 that names an object-specific error code dropped instead of kept for later. + +use std::{ + collections::BTreeMap, + sync::Arc, + time::{Duration, Instant, SystemTime}, +}; + +use base64::Engine; +use litellm_auth_aws::{AwsAuthService, Credentials, sign_s3}; +use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool, Verify}; +use md5::Md5; +use sha2::{Digest, Sha256}; + +use super::{Element, config::Settings, key::object_url}; + +const ATTEMPTS: u32 = 3; +const RETRYABLE: [u16; 4] = [403, 429, 500, 503]; +/// Codes S3 answers that mean this object can never be stored, so retrying it is pointless. +const TERMINAL_CODES: [&str; 7] = [ + "EntityTooLarge", + "InvalidArgument", + "MalformedXML", + "InvalidDigest", + "KeyTooLongError", + "BadDigest", + "InvalidRequest", +]; +/// Credentials without an expiry are resolved again after this long. +const CREDENTIALS_TTL: Duration = Duration::from_secs(600); +/// Credentials are refreshed this long before they expire. +const CREDENTIALS_MARGIN: Duration = Duration::from_secs(300); + +/// How one object's upload ended. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum Upload { + Delivered, + /// Failed for now; it goes back on the queue. + Retry, + /// S3 rejected the object itself. + Dropped, +} + +/// What one attempt got back. +enum Attempt { + Response { status: u16, body: String }, + Transport(String), +} + +impl Attempt { + fn is_terminal(&self) -> bool { + matches!(self, Self::Response { status, body } if is_terminal(*status, body)) + } +} + +/// A 400 or 403 whose `` S3 reserves for a rejection of this object. +fn is_terminal(status: u16, body: &str) -> bool { + matches!(status, 400 | 403) + && error_code(body).is_some_and(|code| TERMINAL_CODES.contains(&code)) +} + +fn error_code(body: &str) -> Option<&str> { + let start = body.find("")? + "".len(); + let length = body[start..].find("")?; + Some(&body[start..start + length]).filter(|code| !code.is_empty() && !code.contains('<')) +} + +/// `(1 << attempt) * (0.5 + random())` seconds. +fn backoff(attempt: u32) -> Duration { + Duration::from_secs_f64(f64::from(1u32 << attempt) * (0.5 + rand::random::())) +} + +pub(super) struct Uploader { + settings: Arc, + pool: Arc, + http: HttpClientConfig, + auth: AwsAuthService, + credentials: tokio::sync::Mutex>, +} + +impl Uploader { + pub fn new( + settings: Arc, + pool: Arc, + http: &HttpClientConfig, + ) -> Self { + let verify = match &settings.verify { + None => http.verify.clone(), + Some(litellm_http::SslVerify::Enabled) => Verify::BuiltInRoots, + Some(litellm_http::SslVerify::Disabled) => Verify::Disabled, + Some(litellm_http::SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()), + }; + Self { + settings, + pool, + http: HttpClientConfig { + verify, + ..http.clone() + }, + auth: AwsAuthService::default(), + credentials: tokio::sync::Mutex::new(None), + } + } + + pub async fn upload(&self, element: &Element) -> Upload { + let url = object_url(&self.settings, &element.key); + let headers = self.headers(element); + let mut attempt = 0; + loop { + let result = match self.put(&url, element, &headers).await { + Ok(result) => result, + Err(error) => { + tracing::error!(%error, key = %element.key, "s3 logging: cannot sign the upload"); + return Upload::Retry; + } + }; + let retry = match &result { + Attempt::Transport(_) => true, + Attempt::Response { status, .. } => { + RETRYABLE.contains(status) + && !(self.settings.drop_on_terminal_error && result.is_terminal()) + } + }; + if retry && attempt + 1 < ATTEMPTS { + let wait = backoff(attempt); + let reason = match &result { + Attempt::Transport(error) => error.clone(), + Attempt::Response { status, .. } => status.to_string(), + }; + tracing::debug!(attempt = attempt + 1, %reason, ?wait, key = %element.key, "s3 logging: upload failed, retrying"); + tokio::time::sleep(wait).await; + attempt += 1; + continue; + } + return match result { + Attempt::Response { status, .. } if (200..300).contains(&status) => { + Upload::Delivered + } + Attempt::Response { status, body } => { + tracing::error!(status, body = %body.chars().take(512).collect::(), key = %element.key, "s3 logging: upload failed"); + if self.settings.drop_on_terminal_error && is_terminal(status, &body) { + tracing::warn!(status, key = %element.key, "s3 logging: dropping the object, S3 rejected it"); + return Upload::Dropped; + } + Upload::Retry + } + Attempt::Transport(error) => { + tracing::error!(%error, key = %element.key, "s3 logging: upload failed"); + Upload::Retry + } + }; + } + } + + /// The proxy's PUT headers, before signing. + fn headers(&self, element: &Element) -> BTreeMap { + let body = &element.body; + let sha256: String = Sha256::digest(body) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + let md5 = + base64::engine::general_purpose::STANDARD.encode(::digest(body)); + let mut headers = BTreeMap::from([ + ("Content-Type".to_owned(), element.content_type.to_owned()), + ("Content-MD5".to_owned(), md5), + ("x-amz-content-sha256".to_owned(), sha256), + ("Content-Language".to_owned(), "en".to_owned()), + ( + "Content-Disposition".to_owned(), + format!("inline; filename=\"{}\"", element.download_name), + ), + ( + "Cache-Control".to_owned(), + "private, immutable, max-age=31536000, s-maxage=0".to_owned(), + ), + ]); + if let Some(algorithm) = &self.settings.server_side_encryption { + headers.insert("x-amz-server-side-encryption".into(), algorithm.clone()); + } + if let Some(key_id) = &self.settings.sse_kms_key_id { + headers.insert( + "x-amz-server-side-encryption-aws-kms-key-id".into(), + key_id.clone(), + ); + } + headers + } + + async fn put( + &self, + url: &str, + element: &Element, + headers: &BTreeMap, + ) -> Result> { + let credentials = self.credentials().await?; + let signature = sign_s3( + "PUT", + url, + &element.body, + headers, + &self.settings.region, + &credentials, + SystemTime::now(), + )?; + let client = self.pool.client(&self.http, ClientVariant::Provider)?; + let mut request = client + .put(url) + .timeout(self.settings.attempt_timeout) + .body(element.body.clone()); + for (name, value) in headers.iter().chain(&signature) { + request = request.header(name, value); + } + Ok(match request.send().await { + Ok(response) => Attempt::Response { + status: response.status().as_u16(), + body: response.text().await.unwrap_or_default(), + }, + Err(error) => Attempt::Transport(error.to_string()), + }) + } + + async fn credentials(&self) -> Result { + let mut cached = self.credentials.lock().await; + if let Some((credentials, until)) = cached.as_ref() + && Instant::now() < *until + { + return Ok(credentials.clone()); + } + let env = |name: &str| std::env::var(name).ok(); + let credentials = self + .auth + .resolve_credentials(self.settings.auth.clone(), &env) + .await?; + let until = credentials + .expiry() + .and_then(|expiry| expiry.duration_since(SystemTime::now()).ok()) + .map_or(CREDENTIALS_TTL, |left| { + left.saturating_sub(CREDENTIALS_MARGIN) + }); + *cached = Some((credentials.clone(), Instant::now() + until)); + Ok(credentials) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::too_large(400, "EntityTooLarge", true)] + #[case::bad_digest(403, "BadDigest", true)] + #[case::access_denied(403, "AccessDenied", false)] + #[case::slow_down(503, "InvalidRequest", false)] + #[case::not_xml(400, "bad request", false)] + #[case::empty_code(400, "", false)] + fn only_object_rejections_are_terminal( + #[case] status: u16, + #[case] body: &str, + #[case] expected: bool, + ) { + assert_eq!(is_terminal(status, body), expected); + } + + #[rstest] + #[case(0, 0.5, 1.5)] + #[case(1, 1.0, 3.0)] + fn backoff_doubles_with_jitter(#[case] attempt: u32, #[case] low: f64, #[case] high: f64) { + for _ in 0..50 { + let wait = backoff(attempt).as_secs_f64(); + assert!((low..high).contains(&wait), "{wait}"); + } + } +} diff --git a/litellm-rust/crates/gateway/src/spend.rs b/litellm-rust/crates/gateway/src/spend.rs index 56491896179..92e949ed71e 100644 --- a/litellm-rust/crates/gateway/src/spend.rs +++ b/litellm-rust/crates/gateway/src/spend.rs @@ -597,7 +597,7 @@ fn metadata(record: &Record) -> Value { /// `spend_logs_metadata` as the proxy merges it: the team's, then the key's, then what hooks /// attached for the request, later keys winning. -fn spend_logs_metadata(record: &Record) -> Value { +pub(crate) fn spend_logs_metadata(record: &Record) -> Value { let mut merged = Map::new(); if let Some(key) = record.key() { for metadata in [&key.team_metadata, &key.metadata] { @@ -617,7 +617,7 @@ fn spend_logs_metadata(record: &Record) -> Value { /// The proxy's usage object for a response, and its `additional_usage_values`: everything but /// the three totals. -fn usage_values(usage: &Usage) -> (Value, Value) { +pub(crate) fn usage_values(usage: &Usage) -> (Value, Value) { let mut additional = Map::new(); additional.insert( "prompt_tokens_details".into(), diff --git a/litellm-rust/crates/gateway/tests/s3_log.rs b/litellm-rust/crates/gateway/tests/s3_log.rs new file mode 100644 index 00000000000..629c217e75b --- /dev/null +++ b/litellm-rust/crates/gateway/tests/s3_log.rs @@ -0,0 +1,313 @@ +//! The `s3_v2` observer against a mock S3: a request ends up as a signed PUT under the proxy's +//! key layout, a 503 is retried, and an object S3 rejects is dropped instead of kept. + +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use base64::Engine; +use litellm_auth_aws::{Credentials, sign_s3}; +use litellm_config::Config; +use litellm_gateway::s3_log::S3Log; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, +}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; +use tokio::net::TcpListener; +use tower::ServiceExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path_regex}, +}; + +const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-x","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":6,"total_tokens":11}}"#; +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":4}}"#; +const OBJECT_KEY: &str = + r"^/logs/audit/research/\d{4}-\d{2}-\d{2}/time-\d{2}-\d{2}-\d{2}-\d{6}_chatcmpl-1\.json$"; + +/// Answers Anthropic and OpenAI calls. +async fn upstream() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(|request: Request| async move { + let body = match request.uri().path().ends_with("/chat/completions") { + true => COMPLETION, + false => MESSAGE, + }; + ([(CONTENT_TYPE, "application/json")], body).into_response() + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + format!("http://{address}") +} + +struct Keys(HashMap); + +impl KeyLookup for Keys { + fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> { + Box::pin(async move { Ok(self.0.get(hash).cloned()) }) + } +} + +fn keys() -> Arc { + let status = KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: Some("team-1".into()), + profile: Arc::new(KeyProfile { + alias: Some("team-key".into()), + user_id: Some("user-1".into()), + team_alias: Some("research".into()), + ..KeyProfile::default() + }), + }; + let hash = KeyHash::from_token(&litellm_auth_types::SecretValue::new("sk-team")); + Arc::new(Keys(HashMap::from([(hash, status)]))) +} + +struct Harness { + app: Router, + inference: Arc, + s3_log: Arc, +} + +async fn harness(s3: &MockServer) -> Harness { + let upstream = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: gpt + litellm_params: + model: openai/gpt-x + api_base: {upstream}/v1 + api_key: sk-upstream + model_info: + id: gpt-1 + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +litellm_settings: + callbacks: [s3_v2] + s3_callback_params: + s3_bucket_name: logs + s3_path: audit + s3_region_name: us-east-1 + s3_endpoint_url: {endpoint} + s3_aws_access_key_id: AKIDTEST + s3_aws_secret_access_key: secret + s3_use_team_prefix: true + s3_server_side_encryption: AES256 +", + endpoint = s3.uri(), + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let s3_log = S3Log::from_config(&config, &inference) + .await + .unwrap() + .unwrap(); + litellm_gateway::hooks::install(&inference, &config, Vec::new(), vec![s3_log.clone()]) + .await + .unwrap(); + let app = + litellm_gateway::router_with_keys(inference.clone(), &config, None, None, Some(keys())); + Harness { + app, + inference, + s3_log, + } +} + +async fn send(app: &Router, path: &str, body: Value) { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer sk-team") + .header("user-agent", "client/1.0") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + assert_eq!(status, 200, "{body:?}"); +} + +/// `YYYYMMDDTHHMMSSZ` as a time. +fn amz_date(text: &str) -> SystemTime { + let number = |range: std::ops::Range| text[range].parse::().unwrap(); + let (year, month, day) = (number(0..4), number(4..6), number(6..8)); + let year = if month <= 2 { year - 1 } else { year }; + let era = year.div_euclid(400); + let year_of_era = year - era * 400; + let day_of_year = (153 * (month + if month > 2 { -3 } else { 9 }) + 2) / 5 + day - 1; + let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + let days = era * 146_097 + day_of_era - 719_468; + let seconds = days * 86_400 + number(9..11) * 3600 + number(11..13) * 60 + number(13..15); + UNIX_EPOCH + Duration::from_secs(seconds as u64) +} + +#[tokio::test] +async fn a_request_is_written_as_a_signed_put_and_a_503_is_retried() { + let s3 = MockServer::start().await; + Mock::given(method("PUT")) + .and(path_regex(OBJECT_KEY)) + .respond_with(ResponseTemplate::new(503)) + .up_to_n_times(1) + .with_priority(1) + .mount(&s3) + .await; + Mock::given(method("PUT")) + .and(path_regex(OBJECT_KEY)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&s3) + .await; + let harness = harness(&s3).await; + send( + &harness.app, + "/v1/chat/completions", + json!({"model": "gpt", "temperature": 0.5, "messages": [{"role": "user", "content": "hello"}]}), + ) + .await; + + litellm_gateway::hooks::close(&harness.inference).await; + + let requests = s3.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2, "the 503 is retried once"); + assert_eq!(requests[0].url, requests[1].url); + let put = &requests[1]; + let header = |name: &str| { + put.headers + .get(name) + .map(|value| value.to_str().unwrap().to_owned()) + }; + let payload: Value = serde_json::from_slice(&put.body).unwrap(); + assert_eq!(payload["id"], "chatcmpl-1"); + assert_eq!(payload["status"], "success"); + assert_eq!(payload["model_group"], "gpt"); + assert_eq!(payload["model_id"], "gpt-1"); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": "hello"}]) + ); + assert_eq!( + payload["response"]["choices"][0]["message"]["content"], + "hi" + ); + assert_eq!(payload["metadata"]["user_api_key_team_alias"], "research"); + assert_eq!(payload["total_tokens"], 11); + + assert_eq!(header("content-type").as_deref(), Some("application/json")); + assert_eq!(header("content-language").as_deref(), Some("en")); + assert_eq!( + header("cache-control").as_deref(), + Some("private, immutable, max-age=31536000, s-maxage=0") + ); + assert_eq!( + header("x-amz-server-side-encryption").as_deref(), + Some("AES256") + ); + let disposition = header("content-disposition").unwrap(); + assert!( + disposition.starts_with("inline; filename=\"time-") + && disposition.ends_with("_chatcmpl-1.json\""), + "{disposition}" + ); + let sha256: String = Sha256::digest(&put.body) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + assert_eq!(header("x-amz-content-sha256"), Some(sha256)); + let md5 = base64::engine::general_purpose::STANDARD + .encode(::digest(&put.body)); + assert_eq!(header("content-md5"), Some(md5)); + + let signed: BTreeMap = [ + "content-type", + "content-md5", + "x-amz-content-sha256", + "content-language", + "content-disposition", + "cache-control", + "x-amz-server-side-encryption", + ] + .into_iter() + .map(|name| (name.to_owned(), header(name).unwrap())) + .collect(); + let expected = sign_s3( + "PUT", + &format!("{}{}", s3.uri(), put.url.path()), + &put.body, + &signed, + "us-east-1", + &Credentials::new("AKIDTEST", "secret", None, None, "test"), + amz_date(&header("x-amz-date").unwrap()), + ) + .unwrap(); + assert_eq!(header("x-amz-date"), Some(expected["X-Amz-Date"].clone())); + assert_eq!( + header("authorization"), + Some(expected["Authorization"].clone()) + ); + assert!( + expected["Authorization"].contains("SignedHeaders=cache-control;content-disposition;content-language;content-md5;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-server-side-encryption"), + "{}", + expected["Authorization"] + ); +} + +#[tokio::test] +async fn an_object_s3_rejects_is_dropped_while_the_rest_are_written() { + let s3 = MockServer::start().await; + Mock::given(method("PUT")) + .and(path_regex(r"_msg_1\.json$")) + .respond_with( + ResponseTemplate::new(400) + .set_body_string("EntityTooLarge"), + ) + .expect(1) + .mount(&s3) + .await; + Mock::given(method("PUT")) + .and(path_regex(OBJECT_KEY)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&s3) + .await; + let harness = harness(&s3).await; + send( + &harness.app, + "/v1/chat/completions", + json!({"model": "gpt", "messages": [{"role": "user", "content": "hello"}]}), + ) + .await; + send( + &harness.app, + "/v1/messages", + json!({"model": "claude", "max_tokens": 8, "messages": [{"role": "user", "content": "hello"}]}), + ) + .await; + + litellm_gateway::hooks::close(&harness.inference).await; + harness.s3_log.flush().await; + + assert_eq!(s3.received_requests().await.unwrap().len(), 2); +} diff --git a/litellm-rust/docs/customer-proxy/s3-log.svg b/litellm-rust/docs/customer-proxy/s3-log.svg new file mode 100644 index 00000000000..1618e7d8044 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/s3-log.svg @@ -0,0 +1,61 @@ + + + + + + s3_v2 observer: each request's logging payload as an S3 object, off the request path + + + At startup + on when callbacks lists s3_v2 (or s3_log); + success_callback / failure_callback + log only successes / only failures + s3_callback_params, os.environ/ resolved; + bad numbers and flags warn, use defaults + credentials: keys, profile, role, web + identity or the default chain + no s3_bucket_name stops startup + + + Request finishes (observe) + StandardLoggingPayload: key, team, + deployment, usage, cost, timings, + messages, response (streams rebuilt) + turn_off_message_logging, strip_base64, + log_prompts_only applied here + queued; past s3_max_queue_size the + oldest are dropped with a warning + nothing awaits S3 here + + + + Every 10 s, and at shutdown + the whole queue, s3_max_concurrent_uploads + PUTs at a time (16), SigV4 signed + s3_batch_file_upload: one .jsonl per + folder instead of one .json per request + failed for now: back at the front, sent + after newer ones on the next flush + DEFAULT_S3_FLUSH_INTERVAL_SECONDS + + + + Object key + {s3_path}/[team_alias/][key_alias/] + YYYY-MM-DD[/HH]/time-HH-MM-SS-ffffff_{id}.json + prefixes with s3_use_team_prefix / s3_use_key_prefix + /HH with s3_partition_granularity: hour + id: the response's id, else the request's call id + past S3's 1024 / 255 byte limits: head + sha256 + path style, or s3_use_virtual_hosted_style + + + Upload and retries + 3 attempts; transport error, 403, 429, 500, 503 + retried after 0.5-1.5 s, then 1-3 s + 400 / 403 with an object error code (EntityTooLarge, + BadDigest, ...) dropped (s3_drop_on_terminal_error) + kept until s3_max_retry_age_seconds (3600) + nothing delivered: all kept, the bucket is down + Content-MD5, SSE and Content-Disposition as the proxy +