diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 20bb696c46c..4b741c993e4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3750,14 +3750,17 @@ version = "0.1.0" dependencies = [ "axum", "base64 0.22.1", + "bytes", "envy", "futures-util", "http 1.4.2", "http-body-util", + "litellm-auth-aws", "litellm-auth-types", "litellm-config", "litellm-db", "litellm-db-testing", + "litellm-framer", "litellm-gateway-auth", "litellm-gateway-hooks", "litellm-gateway-inference", @@ -3770,11 +3773,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 +3790,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "wiremock", ] [[package]] @@ -4055,6 +4061,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-framer", "litellm-host", "litellm-host-native", "litellm-http", @@ -4068,7 +4075,9 @@ dependencies = [ "rstest", "serde_json", "tokio", + "tokio-util", "tracing", + "uuid", "wiremock", ] 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/core-utils/src/core_helpers.rs b/litellm-rust/crates/core-utils/src/core_helpers.rs index ada1c3ceb1a..94268f4b452 100644 --- a/litellm-rust/crates/core-utils/src/core_helpers.rs +++ b/litellm-rust/crates/core-utils/src/core_helpers.rs @@ -11,6 +11,9 @@ const FINISH_REASONS: &[(&str, &str)] = &[ ("end_turn", "stop"), ("stop_sequence", "stop"), ("max_tokens", "length"), + ("tool_use", "tool_calls"), + ("tool_calls", "tool_calls"), + ("function_call", "function_call"), ("refusal", "content_filter"), ("compaction", "length"), ("guardrail_intervened", "content_filter"), @@ -44,7 +47,9 @@ pub fn usage_from_parts( cached_tokens: cache_read_tokens, cache_creation_tokens, text_tokens: input_tokens, + ..PromptTokensDetails::default() }, + ..ChatCompletionsUsage::default() } } @@ -74,6 +79,7 @@ mod tests { #[case::end_turn("end_turn", "stop")] #[case::stop_sequence("stop_sequence", "stop")] #[case::max_tokens("max_tokens", "length")] + #[case::tool_use("tool_use", "tool_calls")] #[case::refusal("refusal", "content_filter")] #[case::guardrail_intervened("guardrail_intervened", "content_filter")] #[case::content_filtered("content_filtered", "content_filter")] diff --git a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs index acea48799c0..eb352ea3ba1 100644 --- a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs +++ b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs @@ -12,6 +12,9 @@ pub enum LlmProviders { Anthropic, #[strum(serialize = "aws_textract")] AwsTextract, + /// Azure OpenAI (`azure/`), distinct from Azure AI Foundry's `azure_ai`. + #[strum(serialize = "azure")] + Azure, #[strum(serialize = "azure_ai")] AzureAi, #[strum(serialize = "bedrock")] diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs index d696427fd8d..3b05ce4bb91 100644 --- a/litellm-rust/crates/gateway-hooks/src/facts.rs +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -221,6 +221,8 @@ pub struct Outcome { /// In dollars, from usage and the deployment's prices; `None` when either is missing. pub cost: Option, pub time_to_first_byte: Option, + /// When the attempt that served, or failed last, was sent to the provider. + pub upstream_started_at: Option, pub ended_at: SystemTime, /// The stream stopped before its end: the client went away or the upstream failed. pub interrupted: bool, diff --git a/litellm-rust/crates/gateway-hooks/src/hook.rs b/litellm-rust/crates/gateway-hooks/src/hook.rs index 9246f77d1a2..7d92d15769c 100644 --- a/litellm-rust/crates/gateway-hooks/src/hook.rs +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -146,7 +146,7 @@ pub struct RequestEdit { allowed: Targets, writer: &'static str, header_ops: Vec, - body_fields: Vec<(String, Value)>, + body_fields: Vec<(Vec, Value)>, response_headers: Vec<(HeaderName, HeaderValue)>, annotations: Annotations, } @@ -183,8 +183,15 @@ impl RequestEdit { /// Sets a top-level field of the upstream body, replacing what the client sent. pub fn set_upstream_body_field(&mut self, field: impl Into, value: Value) { - if self.may(Targets::UPSTREAM_BODY) { - self.body_fields.push((field.into(), value)); + self.set_upstream_body_path(vec![field.into()], value); + } + + /// Sets a field inside a top-level object of the upstream body, such as + /// `["metadata", "user_id"]`, keeping the object's other fields. A missing or non-object + /// step becomes an object. + pub fn set_upstream_body_path(&mut self, path: Vec, value: Value) { + if !path.is_empty() && self.may(Targets::UPSTREAM_BODY) { + self.body_fields.push((path, value)); } } @@ -202,7 +209,11 @@ impl RequestEdit { /// Whether any change reaches the provider. pub fn changes_upstream(&self) -> bool { - !self.header_ops.is_empty() || !self.body_fields.is_empty() + self.changes_upstream_headers() || !self.body_fields.is_empty() + } + + pub fn changes_upstream_headers(&self) -> bool { + !self.header_ops.is_empty() } /// Applies the header changes in order to the headers about to be sent. @@ -221,8 +232,24 @@ impl RequestEdit { /// Applies the body changes in order to the body about to be sent. pub fn apply_upstream_body(&self, body: &mut Map) { - for (field, value) in &self.body_fields { - body.insert(field.clone(), value.clone()); + for (path, value) in &self.body_fields { + let Some((last, parents)) = path.split_last() else { + continue; + }; + let mut object = &mut *body; + for step in parents { + let slot = object + .entry(step.clone()) + .or_insert_with(|| Value::Object(Map::new())); + if !slot.is_object() { + *slot = Value::Object(Map::new()); + } + let Value::Object(inner) = slot else { + unreachable!("just made an object") + }; + object = inner; + } + object.insert(last.clone(), value.clone()); } } @@ -328,6 +355,25 @@ mod tests { ); } + #[rstest] + fn nested_body_fields_keep_their_siblings() { + let mut edit = RequestEdit::default(); + edit.scope("hook", Targets::UPSTREAM_BODY); + edit.set_upstream_body_path(vec!["metadata".into(), "user_id".into()], Value::from("u")); + edit.set_upstream_body_path(vec!["extra".into(), "a".into()], Value::from(1)); + + let serde_json::Value::Object(mut body) = + serde_json::json!({"metadata": {"tag": "t"}, "extra": "not an object"}) + else { + unreachable!() + }; + edit.apply_upstream_body(&mut body); + assert_eq!( + Value::Object(body), + serde_json::json!({"metadata": {"tag": "t", "user_id": "u"}, "extra": {"a": 1}}) + ); + } + #[rstest] fn undeclared_writes_are_dropped() { let mut edit = RequestEdit::default(); diff --git a/litellm-rust/crates/gateway-hooks/src/observer.rs b/litellm-rust/crates/gateway-hooks/src/observer.rs index 6a8b214d5c3..6b748c4aff5 100644 --- a/litellm-rust/crates/gateway-hooks/src/observer.rs +++ b/litellm-rust/crates/gateway-hooks/src/observer.rs @@ -146,6 +146,7 @@ mod tests { usage: None, cost: None, time_to_first_byte: None, + upstream_started_at: None, ended_at: SystemTime::now(), interrupted: false, }, diff --git a/litellm-rust/crates/gateway-hooks/src/pipeline.rs b/litellm-rust/crates/gateway-hooks/src/pipeline.rs index 04d7076a85e..db18be8b9a1 100644 --- a/litellm-rust/crates/gateway-hooks/src/pipeline.rs +++ b/litellm-rust/crates/gateway-hooks/src/pipeline.rs @@ -80,6 +80,7 @@ impl Pipeline { usage: None, cost: None, time_to_first_byte: None, + upstream_started_at: None, ended_at: SystemTime::now(), interrupted: false, }; @@ -185,6 +186,7 @@ mod tests { }), cost: None, time_to_first_byte: None, + upstream_started_at: None, ended_at: SystemTime::now(), interrupted: false, } diff --git a/litellm-rust/crates/gateway-hooks/src/pricing.rs b/litellm-rust/crates/gateway-hooks/src/pricing.rs index be1bca921ba..b2f91d3c91e 100644 --- a/litellm-rust/crates/gateway-hooks/src/pricing.rs +++ b/litellm-rust/crates/gateway-hooks/src/pricing.rs @@ -18,6 +18,12 @@ use crate::usage::{ServiceTier, Usage}; pub struct Prices { plans: HashMap, PricingPlan>, max_output_tokens: HashMap, u64>, + /// Deployments priced by the model their response names, with the provider route the + /// cost map keys that model under. The proxy prices an Azure deployment that sets neither + /// `base_model` nor prices of its own this way. + by_response_model: HashMap, Arc>, + /// Cost map prices by model name, for the deployments priced by their response's model. + models: HashMap, } impl Prices { @@ -27,18 +33,29 @@ impl Prices { if let Some(limit) = fields.get("max_output_tokens").and_then(Value::as_u64) { self.max_output_tokens.insert(id.clone(), limit); } - let parsed = ParsedRates::parse(fields); - if parsed.standard == Rates::EMPTY && parsed.thresholds.is_empty() { - return; + if let Some(plan) = plan(&id, fields) { + self.plans.insert(id, plan); } - match parsed.compile() { - Ok(plan) => { - self.plans.insert(id, plan); - } - Err(error) => tracing::warn!(deployment = %id, ?error, "ignoring invalid prices"), + } + + /// Prices `deployment` by the model its response names, looked up as + /// `{provider}/{model}`, before its own prices. + pub fn price_by_response_model(&mut self, deployment: Arc, provider: &str) { + self.by_response_model.insert(deployment, provider.into()); + } + + /// Adds the cost map prices of `model`, for deployments priced by their response's model. + pub fn insert_model(&mut self, model: &str, fields: &Map) { + if let Some(plan) = plan(model, fields) { + self.models.insert(model.to_owned(), plan); } } + /// Whether some deployment is priced by its response's model. + pub fn prices_by_response_model(&self) -> bool { + !self.by_response_model.is_empty() + } + pub fn len(&self) -> usize { self.plans.len() } @@ -54,41 +71,78 @@ impl Prices { /// The cost in dollars of `usage` on a deployment, or `None` when its prices are unknown. pub fn cost(&self, deployment: &str, usage: &Usage) -> Option { - let plan = self.plans.get(deployment)?; - let one_hour = usage.cache_write_1h_tokens.min(usage.cache_write_tokens); - let split = (one_hour > 0).then_some(()); - let request = Request { - usage: litellm_cost::Usage { - prompt_tokens: usage.prompt_tokens(), - completion_tokens: usage.output_tokens, - cache_read_tokens: usage.cache_read_tokens, - cache_write_tokens: usage.cache_write_tokens, - cache_write_5m_tokens: split.map(|()| usage.cache_write_tokens - one_hour), - cache_write_1h_tokens: split.map(|()| one_hour), - prompt_convention: PromptConvention::IncludesCache, - }, - service_tier: match usage.service_tier { - ServiceTier::Priority => CostTier::Priority, - ServiceTier::Flex => CostTier::Flex, - ServiceTier::Standard | ServiceTier::Batch => CostTier::Standard, - }, - threshold_policy: ThresholdPolicy::Exclusive, - region_multiplier: None, - billed_at_utc_minute: None, - }; - match plan.calculate(&request) { - Ok(cost) => Some(cost.total()), - Err(error) => { - tracing::debug!(deployment, ?error, "no cost for response"); - None - } + plan_cost(deployment, self.plans.get(deployment)?, usage) + } + + /// [`Self::cost`], first by `response_model` for a deployment priced by its response's + /// model, as the proxy tries `{provider}/{response model}` ahead of the deployment. + pub fn response_cost( + &self, + deployment: &str, + response_model: Option<&str>, + usage: &Usage, + ) -> Option { + self.by_response_model + .get(deployment) + .zip(response_model.filter(|model| !model.is_empty())) + .and_then(|(provider, model)| { + let name = match model.starts_with(&format!("{provider}/")) { + true => model.to_owned(), + false => format!("{provider}/{model}"), + }; + plan_cost(deployment, self.models.get(&name)?, usage) + }) + .or_else(|| self.cost(deployment, usage)) + } +} + +/// The plan `fields` price, or `None` when they hold no usable prices. +fn plan(name: &str, fields: &Map) -> Option { + let parsed = ParsedRates::parse(fields); + if parsed.standard == Rates::EMPTY && parsed.thresholds.is_empty() { + return None; + } + parsed + .compile() + .inspect_err(|error| tracing::warn!(deployment = %name, ?error, "ignoring invalid prices")) + .ok() +} + +fn plan_cost(deployment: &str, plan: &PricingPlan, usage: &Usage) -> Option { + let one_hour = usage.cache_write_1h_tokens.min(usage.cache_write_tokens); + let split = (one_hour > 0).then_some(()); + let request = Request { + usage: litellm_cost::Usage { + prompt_tokens: usage.prompt_tokens(), + completion_tokens: usage.output_tokens, + cache_read_tokens: usage.cache_read_tokens, + cache_write_tokens: usage.cache_write_tokens, + cache_write_5m_tokens: split.map(|()| usage.cache_write_tokens - one_hour), + cache_write_1h_tokens: split.map(|()| one_hour), + prompt_convention: PromptConvention::IncludesCache, + }, + service_tier: match usage.service_tier { + ServiceTier::Priority => CostTier::Priority, + ServiceTier::Flex => CostTier::Flex, + ServiceTier::Standard | ServiceTier::Batch => CostTier::Standard, + }, + threshold_policy: ThresholdPolicy::Exclusive, + region_multiplier: None, + billed_at_utc_minute: None, + }; + match plan.calculate(&request) { + Ok(cost) => Some(cost.total()), + Err(error) => { + tracing::debug!(deployment, ?error, "no cost for response"); + None } } } /// A deployment's cost map fields: the cost map entry for its model, overlaid by the prices its /// config sets. The entry is found by `base_model`, then the model name, then the name without -/// its provider route. +/// its provider route. An Azure deployment's names are keyed under `azure/` first, and its bare +/// deployment name is never looked up, since it names a resource of the user's, not a model. pub fn deployment_fields( catalog: Option<&Catalog>, model: &str, @@ -97,14 +151,10 @@ pub fn deployment_fields( ) -> Map { let mut fields = catalog .and_then(|catalog| { - let candidates = base_model.into_iter().chain([ - model, - model.split_once('/').map_or(model, |(_, rest)| rest), - model.rsplit_once('/').map_or(model, |(_, rest)| rest), - ]); - candidates + candidates(model, base_model) + .into_iter() .filter(|name| !name.is_empty()) - .find_map(|name| catalog.lookup(name)) + .find_map(|name| catalog.lookup(&name)) }) .map(|found| found.entry.fields().clone()) .unwrap_or_default(); @@ -118,6 +168,30 @@ pub fn deployment_fields( fields } +fn candidates(model: &str, base_model: Option<&str>) -> Vec { + if model.starts_with(AZURE_ROUTE) { + return base_model + .into_iter() + .flat_map(|base| match base.starts_with(AZURE_ROUTE) { + true => vec![base.to_owned()], + false => vec![format!("{AZURE_ROUTE}{base}"), base.to_owned()], + }) + .chain([model.to_owned()]) + .collect(); + } + base_model + .into_iter() + .chain([ + model, + model.split_once('/').map_or(model, |(_, rest)| rest), + model.rsplit_once('/').map_or(model, |(_, rest)| rest), + ]) + .map(str::to_owned) + .collect() +} + +const AZURE_ROUTE: &str = "azure/"; + #[derive(Clone, Copy)] enum Slot { Input, @@ -382,4 +456,81 @@ mod tests { assert_eq!(fields["output_cost_per_token"], json!(2e-6)); assert!(!fields.contains_key("api_key")); } + + fn azure_catalog() -> Catalog { + Catalog::parse( + br#"{ + "gpt-4o": {"input_cost_per_token": 1e-6, "output_cost_per_token": 1e-6}, + "azure/gpt-4o": {"input_cost_per_token": 2e-6, "output_cost_per_token": 2e-6}, + "gpt-4o-mini": {"input_cost_per_token": 3e-6, "output_cost_per_token": 3e-6}, + "azure/my-dep": {"input_cost_per_token": 4e-6, "output_cost_per_token": 4e-6}, + "chat-dep": {"input_cost_per_token": 5e-6, "output_cost_per_token": 5e-6} + }"#, + Default::default(), + ) + .unwrap() + } + + #[rstest] + #[case::base_model_under_azure("azure/chat-dep", Some("gpt-4o"), Some(2e-6))] + #[case::base_model_already_routed("azure/chat-dep", Some("azure/gpt-4o"), Some(2e-6))] + #[case::bare_base_model("azure/chat-dep", Some("gpt-4o-mini"), Some(3e-6))] + #[case::deployment_under_azure("azure/my-dep", None, Some(4e-6))] + #[case::bare_deployment_is_not_a_model("azure/chat-dep", None, None)] + fn azure_deployments_are_priced_under_the_azure_route( + #[case] model: &str, + #[case] base_model: Option<&str>, + #[case] input: Option, + ) { + let fields = deployment_fields(Some(&azure_catalog()), model, base_model, &[]); + assert_eq!( + fields.get("input_cost_per_token").and_then(Value::as_f64), + input + ); + } + + fn response_priced(own: Value) -> Prices { + let mut prices = Prices::default(); + prices.insert("d".into(), own.as_object().unwrap()); + prices.insert_model( + "azure/gpt-4o", + json!({"input_cost_per_token": 2e-6, "output_cost_per_token": 0.0}) + .as_object() + .unwrap(), + ); + prices.price_by_response_model("d".into(), "azure"); + prices + } + + #[rstest] + #[case::response_model_under_the_route(Some("gpt-4o"), Some(2e-6))] + #[case::response_model_already_routed(Some("azure/gpt-4o"), Some(2e-6))] + #[case::unknown_response_model_uses_the_deployment(Some("gpt-9"), Some(7e-6))] + #[case::no_response_model_uses_the_deployment(None, Some(7e-6))] + fn response_model_prices_come_first( + #[case] response_model: Option<&str>, + #[case] rate: Option, + ) { + let prices = + response_priced(json!({"input_cost_per_token": 7e-6, "output_cost_per_token": 0.0})); + assert_eq!( + prices.response_cost("d", response_model, &usage(1, 0)), + rate + ); + } + + #[rstest] + fn other_deployments_ignore_the_response_model() { + let mut prices = response_priced(json!({})); + prices.insert( + "other".into(), + json!({"input_cost_per_token": 7e-6, "output_cost_per_token": 0.0}) + .as_object() + .unwrap(), + ); + assert_eq!( + prices.response_cost("other", Some("gpt-4o"), &usage(1, 0)), + Some(7e-6) + ); + } } diff --git a/litellm-rust/crates/gateway-hooks/src/usage.rs b/litellm-rust/crates/gateway-hooks/src/usage.rs index c4bcc01b431..a3aa04a9fe0 100644 --- a/litellm-rust/crates/gateway-hooks/src/usage.rs +++ b/litellm-rust/crates/gateway-hooks/src/usage.rs @@ -142,6 +142,8 @@ impl AnthropicUsage { } } +/// OpenAI's usage, with the fields LiteLLM adds when it translates another provider's: +/// cache writes by lifetime and server tool counts. #[derive(Deserialize)] struct OpenAiUsage { #[serde(default)] @@ -149,23 +151,38 @@ struct OpenAiUsage { #[serde(default)] completion_tokens: u64, prompt_tokens_details: Option, + server_tool_use: Option, } #[derive(Deserialize, Default)] struct OpenAiPromptDetails { cached_tokens: Option, cache_creation_tokens: Option, + cache_creation_token_details: Option, } impl OpenAiUsage { fn usage(self) -> Usage { let details = self.prompt_tokens_details.unwrap_or_default(); - Usage::from_prompt_counts( + let mut usage = Usage::from_prompt_counts( self.prompt_tokens, self.completion_tokens, details.cached_tokens.unwrap_or(0), details.cache_creation_tokens.unwrap_or(0), - ) + ); + if let Some(tokens) = details + .cache_creation_token_details + .and_then(|cache| cache.ephemeral_1h_input_tokens) + { + usage.cache_write_1h_tokens = tokens.min(usage.cache_write_tokens); + } + if let Some(requests) = self + .server_tool_use + .and_then(|tools| tools.web_search_requests) + { + usage.web_search_requests = requests; + } + usage } } @@ -197,6 +214,7 @@ pub struct StreamUsage { line: Vec, state: Line, usage: Option, + model: Option, } impl StreamUsage { @@ -206,6 +224,7 @@ impl StreamUsage { line: Vec::new(), state: Line::Reading, usage: None, + model: None, } } @@ -264,6 +283,11 @@ impl StreamUsage { self.usage } + /// The model named by the chunk that carried an OpenAI stream's usage. + pub fn model(&self) -> Option<&str> { + self.model.as_deref() + } + fn may_carry_usage(&self) -> bool { let line = &self.line; if line.len() > MAX_LINE_BYTES { @@ -304,8 +328,13 @@ impl StreamUsage { if !contains(data, b"\"usage\":{") && !contains(data, b"\"usage\": {") { return; } - if let Ok(OpenAiChunk { usage: Some(usage) }) = serde_json::from_slice(data) { + if let Ok(OpenAiChunk { + usage: Some(usage), + model, + }) = serde_json::from_slice(data) + { self.usage = Some(usage.usage()); + self.model = model.filter(|model| !model.is_empty()); } } } @@ -326,6 +355,7 @@ struct AnthropicMessage { #[derive(Deserialize)] struct OpenAiChunk { usage: Option, + model: Option, } fn contains(haystack: &[u8], needle: &[u8]) -> bool { @@ -398,7 +428,7 @@ mod tests { fn openai_usage_comes_from_the_last_chunk() { let stream = concat!( "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}],\"usage\":null}\n\n", - "data: {\"id\":\"c\",\"choices\":[],\"usage\":{\"prompt_tokens\":30,\"completion_tokens\":5,\"total_tokens\":35,\"prompt_tokens_details\":{\"cached_tokens\":10}}}\n\n", + "data: {\"id\":\"c\",\"model\":\"gpt-4o-2024-11-20\",\"choices\":[],\"usage\":{\"prompt_tokens\":30,\"completion_tokens\":5,\"total_tokens\":35,\"prompt_tokens_details\":{\"cached_tokens\":10}}}\n\n", "data: [DONE]\n\n", ); let mut meter = StreamUsage::new(StreamFormat::OpenAi); @@ -414,6 +444,7 @@ mod tests { ..Usage::default() }) ); + assert_eq!(meter.model(), Some("gpt-4o-2024-11-20")); } #[rstest] @@ -431,6 +462,18 @@ mod tests { assert_eq!(anthropic_usage().total_tokens(), 78); } + #[rstest] + fn translated_openai_usage_keeps_cache_lifetimes_and_server_tools() { + let usage = Usage::from_openai(&json!({ + "prompt_tokens": 36, + "completion_tokens": 42, + "prompt_tokens_details": {"cached_tokens": 20, "cache_creation_tokens": 6, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 2, "ephemeral_1h_input_tokens": 4}}, + "server_tool_use": {"web_search_requests": 2}, + })); + assert_eq!(usage, Some(anthropic_usage())); + } + #[rstest] fn openai_body_usage() { let usage = Usage::from_openai(&json!({ diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 118eb2cacad..6f4442dd053 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -12,6 +12,7 @@ use litellm_host_http::Sse; use litellm_inference_chat::{ ChatCompletionsRoute, openai::{OpenAiChatCall, OpenAiChatCompletions}, + stream::{ChatCompletionsStream, ChatStreamCall, ReportedUsage}, types::ChatCompletionsCall, }; use serde_json::{Map, Value}; @@ -93,6 +94,18 @@ async fn handle( None => route, }; + // Chat forwards no client headers of its own; hooks may set some. + let extra_headers = tracking + .as_ref() + .and_then(|tracking| tracking.edit_headers(Vec::new())) + .filter(|headers| !headers.is_empty()) + .map(|headers| { + headers + .into_iter() + .map(|(name, value)| (name, Value::from(value))) + .collect::>() + }); + let extra_headers = &extra_headers; let cache_headers = crate::caching::CacheHeaders::default(); let mut body = Some(body); let routed = gateway @@ -120,9 +133,10 @@ async fn handle( &deployment.model, deployment.custom_llm_provider.as_deref(), ) { - return openai(gateway, deployment, body).await; + return openai(gateway, deployment, body, extra_headers.clone()).await; } let (messages, optional_params) = split(body); + let stream = ChatCompletionsRoute::is_stream(&optional_params); let call = ChatCompletionsCall { model: deployment.model.clone(), messages, @@ -130,9 +144,12 @@ async fn handle( api_key: deployment.api_key.clone(), api_base: deployment.api_base.clone(), custom_llm_provider: deployment.custom_llm_provider.clone(), - extra_headers: None, + extra_headers: extra_headers.clone(), timeout: deployment.timeout, }; + if stream { + return translated_stream(route, call, name).await; + } Ok(litellm_host_http::serve_unary( route.machine(call, cache_options.policy), (), @@ -162,10 +179,29 @@ fn split(mut body: Map) -> (Value, Map) { (messages, body) } +/// The deployment's `litellm_params` a provider reads beyond its connection: its +/// `api_version`, `base_model` and every key the config adds. +fn provider_params(deployment: &Deployment) -> Map { + let params = &deployment.litellm_params; + let mut provider_params: Map = params + .extra + .iter() + .map(|(name, value)| (name.clone(), value.clone())) + .collect(); + if let Some(api_version) = ¶ms.api_version { + provider_params.insert("api_version".into(), api_version.clone().into()); + } + if let Some(base_model) = &deployment.base_model { + provider_params.insert("base_model".into(), base_model.clone().into()); + } + provider_params +} + async fn openai( gateway: &Gateway, deployment: &Deployment, body: Map, + extra_headers: Option>, ) -> Result { let machine = gateway.chat_completions.clone().openai_machine( OpenAiChatCall { @@ -174,8 +210,16 @@ async fn openai( api_key: deployment.api_key.clone(), api_base: deployment.api_base.clone(), custom_llm_provider: deployment.custom_llm_provider.clone(), - extra_headers: None, + extra_headers, timeout: deployment.timeout, + litellm_params: provider_params(deployment), + drop_params: gateway.drop_params() + || deployment + .litellm_params + .drop_params + .as_ref() + .is_some_and(|flag| flag.is_set()), + model_info: gateway.model_info(), }, None, ); @@ -185,3 +229,29 @@ async fn openai( ); Ok(litellm_host_http::serve(machine, (), (), encoder, None).await?) } + +/// A stream from a provider chat translates (Anthropic, Bedrock), sent as OpenAI chunks. The +/// usage it ends with rides on the response for the hooks, even when the client did not ask +/// to see it. +async fn translated_stream( + route: ChatCompletionsRoute, + call: ChatCompletionsCall, + model: &str, +) -> Result { + let usage = ReportedUsage::default(); + let machine = route.stream_machine( + ChatStreamCall { + call, + response_model: model.to_owned(), + usage: usage.clone(), + }, + None, + ); + let encoder = Sse::::new( + |body: Bytes| ([(CONTENT_TYPE, "application/json")], body), + |error| Bytes::from(Error::from(error).openai_sse_frame()), + ); + let mut response = litellm_host_http::serve(machine, (), (), encoder, None).await?; + response.extensions_mut().insert(usage); + Ok(response) +} diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 3aa4b44f9aa..a2eac437ba8 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -121,12 +121,18 @@ impl Error { } _ => self.to_string(), }; - json!({"error": { + let mut error = json!({ "message": message, "type": error_type(status), "param": null, "code": status.as_u16(), - }}) + }); + if let Self::Route(RouteError::Transport(TransportError::Http { body, .. })) = self + && let Some(fields) = azure_content_filter_fields(body) + { + error["provider_specific_fields"] = fields; + } + json!({ "error": error }) } /// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream @@ -233,6 +239,44 @@ fn openai_error_message(body: &str) -> Option { .map(str::to_owned) } +/// Python's `AzureOpenAIExceptionMapping` keeps Azure's error and its content filter results +/// on a content policy violation. Only Azure nests an `inner_error`, which tells its errors +/// apart from another provider's. +fn azure_content_filter_fields(body: &str) -> Option { + const POLICY_PHRASES: [&str; 7] = [ + "content_policy_violation", + "responsibleaipolicyviolation", + "the response was filtered due to the prompt triggering azure openai's content management", + "your task failed as a result of our safety system", + "the model produced invalid content", + "content_filter_policy", + "your request was rejected as a result of our safety system", + ]; + let parsed: Value = serde_json::from_str(body).ok()?; + let azure_error = parsed + .get("error") + .filter(|error| error.is_object()) + .unwrap_or(&parsed); + let inner_error = ["inner_error", "innererror"].into_iter().find_map(|key| { + azure_error + .get(key) + .or_else(|| parsed.get(key)) + .filter(|inner| !inner.is_null()) + })?; + let lowered = body.to_lowercase(); + let violation = azure_error.get("code").and_then(Value::as_str) + == Some("content_policy_violation") + || inner_error.get("code").and_then(Value::as_str) == Some("ResponsibleAIPolicyViolation") + || POLICY_PHRASES.iter().any(|phrase| lowered.contains(phrase)); + violation.then(|| { + json!({ + "innererror": inner_error, + "inner_error": inner_error, + "azure_error": azure_error, + }) + }) +} + fn provider_message(object: &Map) -> Option<&str> { if let Some(detail) = object.get("detail").and_then(Value::as_object) { return detail.get("message").and_then(Value::as_str); @@ -326,6 +370,36 @@ mod tests { assert_eq!(frame, format!("data: {expected}\n\n")); } + #[rstest] + #[case::content_filter( + r#"{"error":{"message":"The response was filtered due to the prompt triggering Azure OpenAI's content management policy.","type":null,"param":"prompt","code":"content_filter","status":400,"innererror":{"code":"ResponsibleAIPolicyViolation","content_filter_result":{"hate":{"filtered":true,"severity":"high"}}}}}"#, + true, + )] + #[case::inner_error_spelling( + r#"{"error":{"message":"blocked","code":"content_policy_violation","inner_error":{"code":"ResponsibleAIPolicyViolation"}}}"#, + true, + )] + #[case::other_azure_error( + r#"{"error":{"message":"bad value","code":"invalid_value","innererror":{"code":"Other"}}}"#, + false + )] + #[case::openai_policy_error( + r#"{"error":{"message":"rejected","code":"content_policy_violation","type":"invalid_request_error"}}"#, + false, + )] + fn azure_content_filter_errors_keep_the_azure_error(#[case] body: &str, #[case] kept: bool) { + let parsed: Value = serde_json::from_str(body).unwrap(); + let error = upstream(400, body).openai_body(); + let fields = &error["error"]["provider_specific_fields"]; + assert_eq!(!fields.is_null(), kept); + if kept { + assert_eq!(fields["azure_error"], parsed["error"]); + assert_eq!(fields["innererror"], fields["inner_error"]); + assert_eq!(fields["innererror"]["code"], "ResponsibleAIPolicyViolation"); + } + assert_eq!(error["error"]["message"], parsed["error"]["message"]); + } + #[rstest] #[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)] #[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index 07833d8cc93..8ecaa26880f 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -19,6 +19,7 @@ use litellm_gateway_hooks::{ Admission, DeploymentFacts, Facts, KeyFacts, Outcome, Pipeline, RequestFacts, Route, StreamFormat, StreamUsage, Usage, }; +use litellm_inference_chat::stream::ReportedUsage; use litellm_router::Attempt; use serde_json::{Map, Value}; @@ -52,7 +53,8 @@ pub(crate) struct Tracking { pipeline: Arc, route: Route, admission: Admission, - deployment: Mutex>, + /// The deployment the latest attempt went to, and when it was sent. + deployment: Mutex>, started: Instant, } @@ -154,6 +156,20 @@ impl Tracking { self.admission.edit().apply_upstream_body(body); } + /// The headers the gateway would forward, after the hooks' changes; `None` when no hook + /// changes upstream headers. + pub(crate) fn edit_headers( + &self, + forwarded: Vec<(String, String)>, + ) -> Option> { + let edit = self.admission.edit(); + edit.changes_upstream_headers().then(|| { + let mut forwarded = forwarded; + edit.apply_upstream_headers(&mut forwarded); + forwarded + }) + } + /// Notes the deployment an attempt went to; the last one is the one that served. pub(crate) fn attempt(&self, name: &str, attempt: &Attempt<'_>) { if !self.pipeline.reads(Facts::DEPLOYMENT) { @@ -167,7 +183,8 @@ impl Tracking { provider: deployment.custom_llm_provider.clone(), api_base: deployment.api_base.clone(), }; - *self.deployment.lock().unwrap_or_else(|e| e.into_inner()) = Some(facts); + *self.deployment.lock().unwrap_or_else(|e| e.into_inner()) = + Some((facts, SystemTime::now())); } /// Adds the hooks' response headers and follows the body to its end; an error is recorded @@ -180,7 +197,10 @@ impl Tracking { deployment, started, } = self; - let deployment = deployment.into_inner().unwrap_or_else(|e| e.into_inner()); + let (deployment, upstream_started_at) = deployment + .into_inner() + .unwrap_or_else(|e| e.into_inner()) + .unzip(); let mut response = match result { Ok(response) => response, Err(error) => { @@ -190,6 +210,7 @@ impl Tracking { usage: None, cost: None, time_to_first_byte: None, + upstream_started_at, ended_at: SystemTime::now(), interrupted: false, }; @@ -203,6 +224,7 @@ impl Tracking { headers.insert(name.clone(), value.clone()); } } + let reported = response.extensions().get::().cloned(); let streaming = response .headers() .get(CONTENT_TYPE) @@ -223,8 +245,10 @@ impl Tracking { route, admission: Some(admission), deployment, + upstream_started_at, status: response.status().as_u16(), meter, + reported, started, first_byte: None, ended: false, @@ -262,8 +286,11 @@ struct Finish { route: Route, admission: Option, deployment: Option, + upstream_started_at: Option, status: u16, meter: Meter, + /// The usage a translated chat stream ended with, sent to the client or not. + reported: Option, body: Option>, started: Instant, first_byte: Option, @@ -284,12 +311,32 @@ impl Finish { } } - fn usage(&self) -> Option { - match &self.meter { - Meter::Off => None, - Meter::Stream(meter) => meter.usage(), - Meter::Body(chunks) => body_usage(self.route, &chunks.concat()), - } + /// The usage the response reported, with its cost when the deployment is priced by the + /// model its response names, which only a chat completion carries. + fn usage_and_cost(&self) -> (Option, Option) { + let by_model = self.pipeline.prices().prices_by_response_model() + && self.route == Route::ChatCompletions; + let (usage, model) = match &self.meter { + Meter::Off => (None, None), + Meter::Stream(meter) => match self.reported.as_ref().and_then(reported_usage) { + Some(usage) => (Some(usage), None), + None => (meter.usage(), meter.model().map(str::to_owned)), + }, + Meter::Body(chunks) => { + let body = chunks.concat(); + let model = by_model.then(|| body_model(&body)).flatten(); + (body_usage(self.route, &body), model) + } + }; + let cost = by_model + .then_some(()) + .zip(self.deployment.as_ref().zip(usage.as_ref())) + .and_then(|((), (deployment, usage))| { + self.pipeline + .prices() + .response_cost(&deployment.id, model.as_deref(), usage) + }); + (usage, cost) } } @@ -298,12 +345,14 @@ impl Drop for Finish { let Some(admission) = self.admission.take() else { return; }; + let (usage, cost) = self.usage_and_cost(); let outcome = Outcome { deployment: self.deployment.take(), status: self.status, - usage: self.usage(), - cost: None, + usage, + cost, time_to_first_byte: self.first_byte, + upstream_started_at: self.upstream_started_at, ended_at: SystemTime::now(), interrupted: !self.ended, }; @@ -312,6 +361,10 @@ impl Drop for Finish { } } +fn reported_usage(reported: &ReportedUsage) -> Option { + Usage::from_openai(&serde_json::to_value(reported.get()?).ok()?) +} + /// The `usage` object of a complete JSON response. Only that object is parsed: it is found by /// its key, which inside a string would be escaped. fn body_usage(route: Route, body: &[u8]) -> Option { @@ -327,6 +380,17 @@ fn body_usage(route: Route, body: &[u8]) -> Option { } } +/// The `model` a complete chat completion names, found by its first key like [`body_usage`]. +fn body_model(body: &[u8]) -> Option { + const KEY: &[u8] = b"\"model\":"; + let start = memchr::memmem::find(body, KEY)? + KEY.len(); + serde_json::Deserializer::from_slice(&body[start..]) + .into_iter::() + .next()? + .ok() + .filter(|model| !model.is_empty()) +} + struct Tapped { inner: BodyDataStream, finish: Finish, @@ -378,4 +442,16 @@ mod tests { body_usage(Route::ChatCompletions, &serde_json::to_vec(&chat).unwrap()).unwrap(); assert_eq!((usage.prompt_tokens(), usage.output_tokens), (5, 6)); } + + #[rstest] + #[case::top_level( + br#"{"id":"c","model":"gpt-4o-2024-11-20","choices":[]}"#, + Some("gpt-4o-2024-11-20") + )] + #[case::spaced(br#"{"model": "gpt-4o", "choices":[]}"#, Some("gpt-4o"))] + #[case::inside_content(br#"{"choices":[{"message":{"content":"\"model\":\"x\""}}]}"#, None)] + #[case::empty(br#"{"model":"","choices":[]}"#, None)] + fn body_model_reads_the_response_model(#[case] body: &[u8], #[case] expected: Option<&str>) { + assert_eq!(body_model(body).as_deref(), expected); + } } diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 56b62d83d77..65aa2d72779 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -25,7 +25,10 @@ use litellm_inference_messages::MessagesRoute; use litellm_inference_ocr::OcrRoute; use litellm_inference_responses::ResponsesRoute; use litellm_inference_transcription::AudioTranscriptionRoute; -use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; +use litellm_llms::base_llm::{ + model_info::ModelInfo, + ocr::{handler::OcrClient, settings::OcrSettings}, +}; use litellm_secrets::source::SecretSource; pub use error::Error; @@ -36,6 +39,8 @@ pub use request::{JsonBody, JsonObject, RequestId}; pub struct Gateway { cache: Option>, hooks: OnceLock, + model_info: OnceLock>, + drop_params: bool, allow_client_side_credentials: bool, pub audio_transcription: AudioTranscriptionRoute, pub chat_completions: ChatCompletionsRoute, @@ -67,6 +72,29 @@ impl Gateway { self.hooks.get() } + /// Gives providers whose param rules read model capabilities the model cost map. Only the + /// first call takes effect. + pub fn install_model_info(&self, model_info: Arc) { + let _ = self.model_info.set(model_info); + } + + pub(crate) fn model_info(&self) -> Option> { + self.model_info.get().cloned() + } + + /// `litellm_settings.drop_params`: drops params a provider does not support instead of + /// refusing the request. + pub fn with_drop_params(self, drop_params: bool) -> Self { + Self { + drop_params, + ..self + } + } + + pub(crate) fn drop_params(&self) -> bool { + self.drop_params + } + /// `general_settings.allow_client_side_credentials`: lets a request body name the /// connection and credential params the proxy otherwise refuses. pub fn with_client_side_credentials(self, allowed: bool) -> Self { @@ -98,6 +126,8 @@ impl Gateway { Ok(Self { cache: None, hooks: OnceLock::new(), + model_info: OnceLock::new(), + drop_params: false, allow_client_side_credentials: false, audio_transcription: AudioTranscriptionRoute::new( provider.clone(), diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index 03f1f57b905..800cabe7659 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -82,6 +82,7 @@ async fn handle( None => route, }; + let upstream_headers = &upstream_headers(headers, tracking.as_ref()); let cache_headers = crate::caching::CacheHeaders::default(); let mut body = Some(body); let routed = gateway @@ -103,7 +104,7 @@ async fn handle( async move { crate::observe::record(name, attempt.deployment); request::authorize(identity, name, attempt.deployment).await?; - let call = project(attempt.deployment, body?, headers)?; + let call = project(attempt.deployment, body?, upstream_headers)?; let machine = route.machine(call, cache_options.policy); let stream = Sse::::new(Json, |error| { Bytes::from(Error::from(error).sse_frame()) @@ -123,10 +124,61 @@ async fn handle( } } +/// Client headers the provider gets: the API headers scoped to the providers that take them, +/// and any header a hook set. +#[derive(Default)] +struct UpstreamHeaders { + scoped: Option, + extra: Option>, +} + +fn upstream_headers( + headers: &HeaderMap, + tracking: Option<&crate::hooks::Tracking>, +) -> UpstreamHeaders { + let client = |name: &str| Some(headers.get(name)?.to_str().ok()?.to_owned()); + let Some(tracking) = tracking else { + return UpstreamHeaders { + scoped: scoped_headers(client), + extra: None, + }; + }; + let mut names: Vec<&str> = header_scopes() + .iter() + .flat_map(|(_, names)| names.iter().copied()) + .collect(); + names.sort_unstable(); + names.dedup(); + let forwarded = names + .iter() + .filter_map(|name| Some(((*name).to_owned(), client(name)?))) + .collect(); + let Some(sent) = tracking.edit_headers(forwarded) else { + return UpstreamHeaders { + scoped: scoped_headers(client), + extra: None, + }; + }; + let value = |name: &str| { + sent.iter() + .find(|(sent, _)| sent.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.clone()) + }; + let extra: Map = sent + .iter() + .filter(|(name, _)| !names.iter().any(|scoped| scoped.eq_ignore_ascii_case(name))) + .map(|(name, value)| (name.clone(), Value::from(value.as_str()))) + .collect(); + UpstreamHeaders { + scoped: scoped_headers(value), + extra: (!extra.is_empty()).then_some(extra), + } +} + fn project( deployment: &Deployment, body: Map, - headers: &HeaderMap, + headers: &UpstreamHeaders, ) -> Result { let body = body .into_iter() @@ -141,28 +193,31 @@ fn project( api_base: deployment.api_base.clone(), custom_llm_provider: deployment.custom_llm_provider.clone(), litellm_params: deployment.litellm_params.clone(), - extra_headers: None, - provider_specific_header: anthropic_api_headers(headers), + extra_headers: headers.extra.clone(), + provider_specific_header: headers.scoped.clone(), timeout: deployment.timeout, shaping: deployment.shaping.clone(), }) } -fn anthropic_api_headers(headers: &HeaderMap) -> Option { - let scopes = [ +fn header_scopes() -> [(&'static str, &'static [&'static str]); 2] { + [ ( ANTHROPIC_API_HEADER_PROVIDERS, ANTHROPIC_API_HEADERS.as_slice(), ), ("openrouter", OPENROUTER_API_HEADERS.as_slice()), - ]; - let forwarded: Vec = scopes + ] +} + +fn scoped_headers(get: impl Fn(&str) -> Option) -> Option { + let forwarded: Vec = header_scopes() .into_iter() .filter_map(|(provider, names)| { let extra_headers: Map = names .iter() .filter_map(|name| { - let value = headers.get(*name)?.to_str().ok()?; + let value = get(name)?; Some(((*name).to_owned(), Value::from(value))) }) .collect(); diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs index b0699ff10b5..73925399e38 100644 --- a/litellm-rust/crates/gateway-inference/tests/routes.rs +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -68,7 +68,6 @@ async fn chat_aliases_call_core_and_use_the_body_model_before_the_path( } #[rstest] -#[case::streaming(json!({"messages": [{"role": "user", "content": "hi"}], "stream": true}), 501)] #[case::missing_messages(json!({}), 400)] #[case::malformed_messages(json!({"messages": "hi"}), 400)] #[case::invalid_streaming_request(json!({"messages": [], "stream": true}), 400)] diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 4a2518780af..dfeeaaa8324 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -12,7 +12,7 @@ http-body-util = "0.1" litellm-inference.workspace = true litellm-gateway-inference.workspace = true litellm-gateway-mcp.workspace = true -tokio-util = "0.7" +tokio-util = { version = "0.7", features = ["codec"] } thiserror.workspace = true futures-util.workspace = true http.workspace = true @@ -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,9 @@ uuid.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] -base64.workspace = true +wiremock.workspace = true +litellm-framer.workspace = true +bytes.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 a2097417ce5..742a036e69d 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -20,6 +20,10 @@ pub enum Error { DatabaseUrl(String), #[error("reading a gateway setting from the secret source failed")] 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/hooks.rs b/litellm-rust/crates/gateway/src/hooks.rs index 02e9180c771..6c393816cbf 100644 --- a/litellm-rust/crates/gateway/src/hooks.rs +++ b/litellm-rust/crates/gateway/src/hooks.rs @@ -1,13 +1,14 @@ //! Builds the hooks and observers inference requests run through, and the per-deployment prices //! observers that read cost need. -use std::sync::Arc; +use std::{collections::HashSet, sync::Arc}; -use litellm_config::Config; +use litellm_config::{Config, Model}; use litellm_gateway_hooks::{ DEFAULT_QUEUE, Facts, Hook, Observer, Observers, Pipeline, Prices, deployment_fields, }; use litellm_gateway_inference::{Gateway, Hooks}; +use litellm_llms::base_llm::model_info::ModelInfo; use litellm_model_catalog::{Catalog, Provenance}; use litellm_secrets::source::SecretSource; use serde_json::{Map, Value}; @@ -29,18 +30,26 @@ pub async fn install( hooks: Vec>, observers: Vec>, ) -> Result<(), Error> { - if hooks.is_empty() && observers.is_empty() { + let serves_azure = config.model_list.iter().any(is_azure); + if hooks.is_empty() && observers.is_empty() && !serves_azure { return Ok(()); } let needs_cost = observers .iter() .any(|observer| observer.reads().contains(Facts::COST)) || hooks.iter().any(|hook| hook.reads().contains(Facts::COST)); + let catalog = match needs_cost || serves_azure { + true => load_catalog(inference).await.map(Arc::new), + false => None, + }; + if let Some(catalog) = catalog.as_ref().filter(|_| serves_azure) { + inference.install_model_info(Arc::new(CostMap::new(catalog.clone()))); + } + if hooks.is_empty() && observers.is_empty() { + return Ok(()); + } let prices = match needs_cost { - true => { - let catalog = load_catalog(inference).await; - prices(config, inference, catalog.as_ref()) - } + true => prices(config, inference, catalog.as_deref()), false => Prices::default(), }; let names: Vec<_> = hooks @@ -135,7 +144,9 @@ async fn fetch(inference: &Gateway, url: &str) -> Result, Box) -> Prices { let mut prices = Prices::default(); for (model, id) in config @@ -151,18 +162,75 @@ fn prices(config: &Config, inference: &Gateway, catalog: Option<&Catalog>) -> Pr .iter() .map(|(name, value)| (name, value.clone())), ); - let base_model = model_info.get("base_model").and_then(Value::as_str); - let fields = deployment_fields( - catalog, - &model.litellm_params.model, - base_model, - &[&model_info, ¶ms], - ); + let base_model = params + .get("base_model") + .or_else(|| model_info.get("base_model")) + .and_then(Value::as_str) + .filter(|name| !name.is_empty()); + let name = match is_azure(model) && !model.litellm_params.model.starts_with(AZURE_ROUTE) { + true => format!("{AZURE_ROUTE}{}", model.litellm_params.model), + false => model.litellm_params.model.clone(), + }; + let overrides = [&model_info, ¶ms]; + if is_azure(model) + && base_model.is_none() + && deployment_fields(None, &name, None, &overrides).is_empty() + { + prices.price_by_response_model(id.into(), AZURE_PROVIDER); + } + let fields = deployment_fields(catalog, &name, base_model, &overrides); prices.insert(id.into(), &fields); } + if let Some(catalog) = catalog.filter(|_| prices.prices_by_response_model()) { + for name in catalog + .model_names() + .filter(|name| name.starts_with(AZURE_ROUTE)) + { + if let Some(found) = catalog.lookup(name) { + prices.insert_model(name, found.entry.fields()); + } + } + } prices } +const AZURE_PROVIDER: &str = "azure"; +const AZURE_ROUTE: &str = "azure/"; + +fn is_azure(model: &Model) -> bool { + match model.litellm_params.custom_llm_provider.as_deref() { + Some(provider) => provider == AZURE_PROVIDER, + None => model.litellm_params.model.starts_with(AZURE_ROUTE), + } +} + +/// The model cost map as the capabilities Azure's param rules read. +struct CostMap { + catalog: Arc, + keys: HashSet, +} + +impl CostMap { + fn new(catalog: Arc) -> Self { + let keys = catalog + .model_names() + .filter_map(|name| catalog.lookup(name)) + .flat_map(|found| found.entry.fields().keys().cloned().collect::>()) + .collect(); + Self { catalog, keys } + } +} + +impl ModelInfo for CostMap { + fn entry(&self, model: &str) -> Option<&Map> { + self.catalog.lookup(model).map(|found| found.entry.fields()) + } + + fn declares(&self, key: &str) -> bool { + self.keys.contains(key) + } +} + fn json_object<'a, V: serde::Serialize + 'a>( entries: impl Iterator, ) -> Map { diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 57f05ec9322..82279726f27 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -9,7 +9,11 @@ pub mod hooks; mod key_status; mod keys; pub mod metrics; +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; @@ -75,7 +79,14 @@ pub fn build_inference(config: &Config) -> Result, Error> { let resources = CoreResources::new(pool); Ok(Arc::new( Gateway::new(resources, http, secrets, ModelRouter::from_config(config)?)? - .with_client_side_credentials(allows_client_side_credentials(config)), + .with_client_side_credentials(allows_client_side_credentials(config)) + .with_drop_params( + config + .litellm_settings + .drop_params + .as_ref() + .is_some_and(|flag| flag.is_set()), + ), )) } diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 3e648f5a758..71482524afc 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -128,12 +128,25 @@ async fn main() -> Result<(), Box> { ), )); } + if let Some(capacity) = + litellm_gateway::model_capacity::ModelCapacity::for_gateway(&config, coordination.clone()) + { + hooks.push(Arc::new(capacity)); + } + hooks.extend(litellm_gateway::request_hooks::from_config(&config)?); let mut observers: Vec> = Vec::new(); let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config)); let spend_tracking = spend.is_some(); if let Some((pool, settings)) = spend { observers.push(litellm_gateway::spend::SpendWriter::start(pool, settings)); } + let prometheus = litellm_gateway::prometheus::Prometheus::from_config(&config)?.map(Arc::new); + 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()?; @@ -148,7 +161,8 @@ async fn main() -> Result<(), Box> { tracing::info!(%address, models = config.model_list.len(), virtual_keys, shared_state, spend_tracking, log_level = %level, "gateway listening"); let draining = litellm_gateway::health::Draining::default(); - let metrics = litellm_gateway::metrics::Metrics::new(&config, draining.flag()); + let metrics = + litellm_gateway::metrics::Metrics::new(&config, draining.flag()).serve(prometheus); let metrics_auth = litellm_gateway::metrics_auth(&config, inference.secrets.clone()); let app = litellm_gateway::with_ops( litellm_gateway::router_with_key_cache(inference.clone(), &config, ui, mcp_router, keys), diff --git a/litellm-rust/crates/gateway/src/metrics.rs b/litellm-rust/crates/gateway/src/metrics.rs index 01f63b3d47e..a75d8ce8cd8 100644 --- a/litellm-rust/crates/gateway/src/metrics.rs +++ b/litellm-rust/crates/gateway/src/metrics.rs @@ -9,7 +9,7 @@ use std::{ collections::HashMap, fmt::Write as _, sync::{ - Arc, RwLock, + Arc, OnceLock, RwLock, atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}, }, time::{Duration, Instant}, @@ -26,6 +26,8 @@ use http_body_util::BodyExt; use litellm_config::Config; use litellm_gateway_inference::observe; +use crate::prometheus::Prometheus; + /// Upper bounds, in seconds, shared by the latency and time to first byte histograms. They reach /// 30 minutes because streams that long are normal here. const BUCKETS: [f64; 19] = [ @@ -123,6 +125,8 @@ struct Inner { failures: RwLock>, in_flight: AtomicI64, draining: Arc, + /// The `prometheus` callback's metrics, served after these. + prometheus: OnceLock>, } impl Metrics { @@ -164,9 +168,19 @@ impl Metrics { failures: RwLock::default(), in_flight: AtomicI64::new(0), draining, + prometheus: OnceLock::new(), })) } + /// Serves `prometheus`'s metrics too. The ones it shares a name with are then served from it + /// alone, with the proxy's labels. + pub fn serve(self, prometheus: Option>) -> Self { + if let Some(prometheus) = prometheus { + let _ = self.0.prometheus.set(prometheus); + } + self + } + fn series(&self, route: &str, key: SeriesKey) -> Arc { if let Some(series) = self .0 @@ -249,21 +263,24 @@ impl Metrics { }) .collect(); - out.push_str("# HELP litellm_proxy_total_requests_metric Total requests the proxy answered, by route, requested model and status code\n"); - out.push_str("# TYPE litellm_proxy_total_requests_metric counter\n"); - for ((_, _, series), labels) in series.iter().zip(&labels) { - let _ = writeln!( - out, - "litellm_proxy_total_requests_metric_total{{{labels}}} {}", - series.latency.count.load(Ordering::Relaxed) - ); - } - out.push_str("# HELP litellm_request_total_latency_metric Seconds from request start to the end of the response body\n"); - out.push_str("# TYPE litellm_request_total_latency_metric histogram\n"); - for ((_, _, series), labels) in series.iter().zip(&labels) { - series - .latency - .render(&mut out, "litellm_request_total_latency_metric", labels); + let replaced = inner.prometheus.get().is_some(); + if !replaced { + out.push_str("# HELP litellm_proxy_total_requests_metric Total requests the proxy answered, by route, requested model and status code\n"); + out.push_str("# TYPE litellm_proxy_total_requests_metric counter\n"); + for ((_, _, series), labels) in series.iter().zip(&labels) { + let _ = writeln!( + out, + "litellm_proxy_total_requests_metric_total{{{labels}}} {}", + series.latency.count.load(Ordering::Relaxed) + ); + } + out.push_str("# HELP litellm_request_total_latency_metric Seconds from request start to the end of the response body\n"); + out.push_str("# TYPE litellm_request_total_latency_metric histogram\n"); + for ((_, _, series), labels) in series.iter().zip(&labels) { + series + .latency + .render(&mut out, "litellm_request_total_latency_metric", labels); + } } out.push_str("# HELP litellm_request_time_to_first_byte_metric Seconds from request start to the first response body byte\n"); out.push_str("# TYPE litellm_request_time_to_first_byte_metric histogram\n"); @@ -275,25 +292,27 @@ impl Metrics { ); } - out.push_str("# HELP litellm_deployment_failure_responses Upstream calls that failed with 429 or 5xx, by model group and deployment api_base\n"); - out.push_str("# TYPE litellm_deployment_failure_responses counter\n"); - let mut failures: Vec<(FailureKey, u64)> = inner - .failures - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .map(|(key, count)| (*key, count.load(Ordering::Relaxed))) - .collect(); - failures.sort_by_key(|(key, _)| *key); - for (key, count) in failures { - let deployment = &inner.deployments[key.deployment as usize]; - let _ = writeln!( - out, - "litellm_deployment_failure_responses_total{{requested_model=\"{}\",api_base=\"{}\",exception_status=\"{}\"}} {count}", - escape(&inner.model_labels[deployment.model as usize]), - escape(&deployment.api_base), - key.status - ); + if !replaced { + out.push_str("# HELP litellm_deployment_failure_responses Upstream calls that failed with 429 or 5xx, by model group and deployment api_base\n"); + out.push_str("# TYPE litellm_deployment_failure_responses counter\n"); + let mut failures: Vec<(FailureKey, u64)> = inner + .failures + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .map(|(key, count)| (*key, count.load(Ordering::Relaxed))) + .collect(); + failures.sort_by_key(|(key, _)| *key); + for (key, count) in failures { + let deployment = &inner.deployments[key.deployment as usize]; + let _ = writeln!( + out, + "litellm_deployment_failure_responses_total{{requested_model=\"{}\",api_base=\"{}\",exception_status=\"{}\"}} {count}", + escape(&inner.model_labels[deployment.model as usize]), + escape(&deployment.api_base), + key.status + ); + } } out.push_str("# HELP litellm_in_flight_requests Requests whose response has not finished, streams included\n"); @@ -310,6 +329,9 @@ impl Metrics { "litellm_gateway_draining {}", u8::from(inner.draining.load(Ordering::Relaxed)) ); + if let Some(prometheus) = inner.prometheus.get() { + prometheus.render(&mut out); + } out } } @@ -327,7 +349,7 @@ fn api_base_label(api_base: &str) -> String { } } -fn escape(value: &str) -> std::borrow::Cow<'_, str> { +pub(crate) fn escape(value: &str) -> std::borrow::Cow<'_, str> { if value.contains(['\\', '"', '\n']) { value .replace('\\', "\\\\") diff --git a/litellm-rust/crates/gateway/src/model_capacity.rs b/litellm-rust/crates/gateway/src/model_capacity.rs new file mode 100644 index 00000000000..5cc98e808e2 --- /dev/null +++ b/litellm-rust/crates/gateway/src/model_capacity.rs @@ -0,0 +1,654 @@ +//! The model capacity hook, the proxy's `dynamic_rate_limiter_v3`: caps each model group at the +//! sum of its deployments' `rpm` and `tpm` across every replica, and with `priority_reservation` +//! splits that capacity between priorities once the group is busy. +//! +//! Each request adds one to `{model_saturation_check:}:requests` and checks the group's +//! tokens, all or nothing, with the proxy's keys and scripts. With priorities, the request's pool +//! `{priority_model::}` (or `:default_pool`) is held to its share of capacity +//! once the group's saturation reaches `saturation_threshold`, and only counted below it. Tokens +//! are added to both counters when the response ends. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use futures_util::future::BoxFuture; +use http::{HeaderName, HeaderValue}; +use litellm_config::Config; +use litellm_gateway_hooks::{ + Facts, Hook, KeyFacts, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets, +}; +use serde_json::Value; + +use crate::{ + coordination::Coordination, + rate_limits::{Kind, Pair, RateLimitSettings, Reserved, Store, Unavailable}, +}; + +/// The names a config's `callbacks` list turns this hook on with. +const NAMES: [&str; 2] = ["dynamic_rate_limiter_v3", "model_capacity"]; + +/// A priority's share of a group's capacity, as `priority_reservation` writes it. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Share { + Percent(f64), + Rpm(f64), + Tpm(f64), +} + +impl Share { + /// The share as a fraction of the group, as the proxy's `convert_priority_to_percent`. + fn fraction(self, capacity: &Capacity) -> f64 { + match self { + Self::Percent(value) => value, + Self::Rpm(value) => match capacity.rpm { + Some(rpm) if rpm > 0 => value / rpm as f64, + _ => value, + }, + Self::Tpm(value) => match capacity.tpm { + Some(tpm) if tpm > 0 => value / tpm as f64, + _ => value, + }, + } + } +} + +/// Which tokens count: `general_settings.token_rate_limit_type`. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum TokenCount { + #[default] + Total, + Input, + Output, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CapacitySettings { + pub window: u64, + pub fail_closed: bool, + /// `litellm_settings.priority_reservation`, in config order. + pub reservation: Option>, + /// `priority_reservation_settings.default_priority`: the share of keys without a listed + /// priority, together. + pub default_priority: f64, + /// `priority_reservation_settings.saturation_threshold`. + pub saturation_threshold: f64, + /// `priority_reservation_settings.saturation_check_cache_ttl`: how long a group's + /// saturation is reused before it is read again. + pub saturation_ttl: Duration, + pub tokens: TokenCount, +} + +impl Default for CapacitySettings { + fn default() -> Self { + let rate_limits = RateLimitSettings::default(); + Self { + window: rate_limits.window, + fail_closed: rate_limits.fail_closed, + reservation: None, + default_priority: 0.25, + saturation_threshold: 0.5, + saturation_ttl: Duration::from_secs(60), + tokens: TokenCount::Total, + } + } +} + +impl CapacitySettings { + /// `None` when the config's `callbacks` do not list the hook. + pub fn from_config(config: &Config) -> Option { + let litellm = &config.litellm_settings; + let listed = litellm.callbacks.iter().flatten().any(|callback| { + callback + .as_str() + .is_some_and(|name| NAMES.contains(&name.trim())) + }); + if !listed { + return None; + } + let rate_limits = RateLimitSettings::from_config(config); + let number = |value: &litellm_config::Value| match value { + litellm_config::Value::Number(number) => number.as_f64(), + litellm_config::Value::String(text) => text.trim().parse().ok(), + _ => None, + }; + let reservation = match litellm.additional_fields.get("priority_reservation") { + Some(litellm_config::Value::Mapping(mapping)) => Some( + mapping + .iter() + .filter_map(|(priority, share)| { + let priority = match priority { + litellm_config::Value::String(text) => text.clone(), + litellm_config::Value::Number(number) => number.to_string(), + _ => return None, + }; + let share = match share { + litellm_config::Value::Mapping(share) => { + let value = share.get("value").and_then(number).unwrap_or(1.0); + match share.get("type").and_then(litellm_config::Value::as_str) { + Some("rpm") => Share::Rpm(value), + Some("tpm") => Share::Tpm(value), + _ => Share::Percent(value), + } + } + other => Share::Percent(number(other)?), + }; + Some((priority, share)) + }) + .collect(), + ), + _ => None, + }; + let settings = litellm + .additional_fields + .get("priority_reservation_settings"); + let setting = |name: &str| { + settings + .and_then(|settings| settings.get(name)) + .and_then(number) + }; + let defaults = Self::default(); + Some(Self { + window: rate_limits.window, + fail_closed: rate_limits.fail_closed, + reservation, + default_priority: setting("default_priority").unwrap_or(defaults.default_priority), + saturation_threshold: setting("saturation_threshold") + .unwrap_or(defaults.saturation_threshold), + saturation_ttl: setting("saturation_check_cache_ttl") + .filter(|seconds| *seconds >= 0.0) + .map_or(defaults.saturation_ttl, Duration::from_secs_f64), + tokens: match config + .general_settings + .additional_fields + .get("token_rate_limit_type") + .and_then(litellm_config::Value::as_str) + { + Some("input") => TokenCount::Input, + Some("output") => TokenCount::Output, + _ => TokenCount::Total, + }, + }) + } +} + +/// A model group's requests and tokens per window: its deployments' `rpm` and `tpm`, summed. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Capacity { + pub rpm: Option, + pub tpm: Option, +} + +/// Each configured group's capacity, from `rpm`/`tpm` on the deployment, its `litellm_params` +/// or its `model_info`, in that order, as the proxy's router sums them. +pub fn capacities(config: &Config) -> HashMap { + let mut groups: HashMap = HashMap::new(); + for model in &config.model_list { + let configured = |field: &str| -> Option { + let from_value = |value: &litellm_config::Value| match value { + litellm_config::Value::Number(number) => number.as_f64(), + litellm_config::Value::String(text) => text.trim().parse().ok(), + _ => None, + }; + let spelled = |value: &litellm_config::Spelled| match value { + litellm_config::Spelled::Value(value) => Some(*value), + litellm_config::Spelled::Text(text) => text.trim().parse().ok(), + }; + let params = match field { + "tpm" => model.litellm_params.tpm.as_ref(), + _ => model.litellm_params.rpm.as_ref(), + }; + model + .additional_fields + .get(field) + .and_then(from_value) + .or_else(|| params.and_then(spelled)) + .or_else(|| model.model_info.get(field).and_then(from_value)) + }; + let group = groups.entry(model.model_name.clone()).or_default(); + let add = |total: &mut Option, value: Option| { + if let Some(value) = value { + *total = Some(total.unwrap_or(0) + value as i64); + } + }; + add(&mut group.tpm, configured("tpm")); + add(&mut group.rpm, configured("rpm")); + } + groups +} + +/// The pool a request draws on and its share of the group. +#[derive(Clone, Debug, PartialEq)] +struct Pool { + name: String, + share: f64, +} + +/// The `model_capacity` hook. +pub struct ModelCapacity { + settings: CapacitySettings, + groups: HashMap, + store: Arc, + saturation: Mutex, Instant)>>, +} + +impl ModelCapacity { + pub fn new( + settings: CapacitySettings, + groups: HashMap, + redis: Option>, + ) -> Self { + Self { + store: Arc::new(Store::new(redis, settings.window, settings.fail_closed)), + settings, + groups, + saturation: Mutex::default(), + } + } + + /// The hook for a gateway whose config lists it, else `None`. + pub fn for_gateway(config: &Config, redis: Option>) -> Option { + Some(Self::new( + CapacitySettings::from_config(config)?, + capacities(config), + redis, + )) + } + + /// The pool the proxy's `_get_priority_allocation` picks, with its normalized share. + fn pool(&self, group: &str, capacity: &Capacity, priority: Option<&str>) -> Pool { + let explicit = self.settings.reservation.as_ref().and_then(|reservation| { + let priority = priority?; + reservation + .iter() + .any(|(name, _)| name == priority) + .then(|| { + let fractions: Vec<(&str, f64)> = reservation + .iter() + .map(|(name, share)| (name.as_str(), share.fraction(capacity))) + .collect(); + let total: f64 = fractions.iter().map(|(_, fraction)| fraction).sum(); + let fraction = fractions + .iter() + .find(|(name, _)| *name == priority) + .map_or(0.0, |(_, fraction)| *fraction); + ( + priority, + if total > 1.0 { + fraction / total + } else { + fraction + }, + ) + }) + }); + match explicit { + Some((priority, share)) => Pool { + name: format!("{group}:{priority}"), + share, + }, + None => Pool { + name: format!("{group}:default_pool"), + share: self.settings.default_priority, + }, + } + } + + /// A counter's value, reused for `saturation_check_cache_ttl`. + async fn cached(&self, prefix: &str, kind: Kind) -> Option { + let key = format!("{prefix}:{}", kind.name()); + if let Some((value, read_at)) = self + .saturation + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&key) + && read_at.elapsed() < self.settings.saturation_ttl + { + return *value; + } + let value = self.store.read(prefix, kind).await; + let mut cache = self.saturation.lock().unwrap_or_else(|e| e.into_inner()); + if cache.len() >= 10_000 { + cache.clear(); + } + cache.insert(key, (value, Instant::now())); + value + } + + /// How full the group is: its busiest of requests and tokens against capacity. + async fn saturation(&self, prefix: &str, capacity: &Capacity) -> f64 { + let mut saturation = 0.0_f64; + for (kind, limit) in [(Kind::Requests, capacity.rpm), (Kind::Tokens, capacity.tpm)] { + if let Some(limit) = limit.filter(|limit| *limit > 0) + && let Some(count) = self.cached(prefix, kind).await + { + saturation = saturation.max(count as f64 / limit as f64); + } + } + saturation + } + + fn unavailable(&self, request: &RequestFacts) -> Refusal { + crate::rate_limits::unavailable(request.route) + } +} + +/// `team_metadata.priority`, else the key's `metadata.priority`. +fn priority(key: Option<&KeyFacts>) -> Option { + let key = key?; + let read = |value: &Value| match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + _ => None, + }; + key.team_metadata + .get("priority") + .and_then(read) + .or_else(|| key.metadata.get("priority").and_then(read)) +} + +fn limit_text(limit: Option) -> String { + limit.map_or_else(|| "not configured".to_owned(), |limit| limit.to_string()) +} + +/// `{:.%}` as Python formats a fraction. +fn percent(fraction: f64, places: usize) -> String { + format!("{:.*}%", places, fraction * 100.0) +} + +impl Hook for ModelCapacity { + fn name(&self) -> &'static str { + "model_capacity" + } + + fn reads(&self) -> Facts { + Facts::KEY | Facts::USAGE + } + + fn writes(&self) -> Targets { + Targets::ADMISSION | Targets::RESPONSE_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + let priority = priority(request.key.as_deref()); + let priority_header = priority.as_deref().unwrap_or("default"); + if !request.stream { + if priority_header.chars().all(|char| (char as u32) < 256) + && let Ok(value) = HeaderValue::from_str(priority_header) + { + edit.add_response_header(HeaderName::from_static("x-litellm-priority"), value); + } + edit.add_response_header( + HeaderName::from_static("x-litellm-rate-limiter-version"), + HeaderValue::from_static("v3"), + ); + } + let group = &*request.requested_model; + let Some(capacity) = self.groups.get(group).copied() else { + return Ok(None); + }; + if capacity.rpm.is_none() && capacity.tpm.is_none() { + return Ok(None); + } + let model_prefix = format!("{{model_saturation_check:{group}}}"); + let saturation = self.saturation(&model_prefix, &capacity).await; + let pool = self.pool(group, &capacity, priority.as_deref()); + let pool_prefix = format!("{{priority_model:{}}}", pool.name); + + // The group counts every request against its whole capacity; a zero limit is off. + let mut pairs = Vec::new(); + let push = |pairs: &mut Vec<(&'static str, Pair)>, + key: &'static str, + prefix: &str, + limits: [Option; 2]| { + for (kind, limit, increment) in + [(Kind::Requests, limits[0], 1), (Kind::Tokens, limits[1], 0)] + { + if let Some(limit) = limit { + pairs.push(( + key, + Pair { + prefix: prefix.to_owned(), + kind, + limit, + increment, + }, + )); + } + } + }; + push( + &mut pairs, + "model_saturation_check", + &model_prefix, + [ + capacity.rpm.filter(|rpm| *rpm != 0), + capacity.tpm.filter(|tpm| *tpm != 0), + ], + ); + let share = |limit: Option| limit.map(|limit| (limit as f64 * pool.share) as i64); + let pool_limits = [share(capacity.rpm), share(capacity.tpm)]; + let reserving = + self.settings.reservation.is_some() && pool_limits.iter().any(Option::is_some); + let enforced = reserving && saturation >= self.settings.saturation_threshold; + if enforced { + push(&mut pairs, "priority_model", &pool_prefix, pool_limits); + } + let checks: Vec = pairs.iter().map(|(_, pair)| pair.clone()).collect(); + let (reserved, shared) = self + .store + .check_and_add(&checks) + .await + .map_err(|Unavailable| self.unavailable(request))?; + let counts = match reserved { + Reserved::Counts(counts) => counts, + Reserved::Over { index, current } => { + let (key, pair) = &pairs[index]; + let remaining = pair.limit - current; + let kind = pair.kind.name(); + let priority_text = priority.as_deref().unwrap_or("None"); + let (message, saturation_header) = match *key { + "model_saturation_check" => ( + format!( + "Model capacity reached for {group}. Priority: {priority_text}, Rate limit type: {kind}, Model TPM: {}, Model RPM: {}, Remaining: {remaining}", + limit_text(capacity.tpm), + limit_text(capacity.rpm), + ), + None, + ), + _ => ( + format!( + "Priority-based rate limit exceeded. Model: {group}, Priority: {priority_text}, Rate limit type: {kind}, Model TPM: {}, Model RPM: {}, Remaining: {remaining}, Model saturation: {}", + limit_text(capacity.tpm), + limit_text(capacity.rpm), + percent(saturation, 1), + ), + Some(percent(saturation, 2)), + ), + }; + let mut headers = vec![ + ("retry-after", self.settings.window.to_string()), + ("rate_limit_type", kind.to_owned()), + ("x-litellm-priority", priority_header.to_owned()), + ]; + if let Some(saturation) = saturation_header { + headers.push(("x-litellm-saturation", saturation)); + } + return Err(crate::rate_limits::throttled( + request.route, + message, + headers, + )); + } + }; + let mut statuses: Vec<(&str, Kind, i64, i64)> = pairs + .iter() + .zip(&counts) + .map(|((key, pair), count)| (*key, pair.kind, pair.limit, pair.limit - count)) + .collect(); + // Below the threshold the pool's requests are only counted, so its share is known + // once it is enforced. + if reserving + && !enforced + && let Some(rpm) = pool_limits[0] + && let Ok(counts) = self + .store + .requests(std::slice::from_ref(&pool_prefix)) + .await + { + statuses.push(("priority_model", Kind::Requests, rpm, rpm - counts[0])); + } + if !request.stream { + for (key, kind, limit, remaining) in statuses { + let kind = kind.name(); + for (name, value) in [ + (format!("x-ratelimit-{key}-remaining-{kind}"), remaining), + (format!("x-ratelimit-{key}-limit-{kind}"), limit), + ] { + if let Ok(name) = HeaderName::from_bytes(name.as_bytes()) { + edit.add_response_header(name, HeaderValue::from(value)); + } + } + } + } + Ok(Some(Box::new(Charge { + store: self.store.clone(), + prefixes: [model_prefix, pool_prefix], + tokens: self.settings.tokens, + shared, + }) as Box)) + }) + } +} + +/// Adds a successful response's tokens to the group and its pool. +struct Charge { + store: Arc, + prefixes: [String; 2], + tokens: TokenCount, + shared: bool, +} + +impl Settle for Charge { + fn settle(self: Box, record: &Record) { + if !(200..300).contains(&record.outcome.status) { + return; + } + let Some(usage) = record.outcome.usage else { + return; + }; + let tokens = match self.tokens { + TokenCount::Total => usage.total_tokens(), + TokenCount::Input => usage.prompt_tokens(), + TokenCount::Output => usage.output_tokens, + }; + let Ok(tokens) = i64::try_from(tokens) else { + return; + }; + if tokens == 0 { + return; + } + let Charge { + store, + prefixes, + shared, + .. + } = *self; + tokio::spawn(async move { + let deltas = prefixes + .into_iter() + .map(|prefix| (prefix, tokens)) + .collect(); + store.charge(deltas, shared).await; + }); + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn capacity_hook(reservation: Option>) -> ModelCapacity { + ModelCapacity::new( + CapacitySettings { + reservation, + ..CapacitySettings::default() + }, + HashMap::new(), + None, + ) + } + + #[rstest] + fn pools_follow_the_proxy_allocation() { + let capacity = Capacity { + rpm: Some(100), + tpm: Some(1_000), + }; + let hook = capacity_hook(Some(vec![ + ("prod".into(), Share::Percent(0.6)), + ("dev".into(), Share::Rpm(80.0)), + ])); + // 0.6 + 0.8 is over-allocated, so both shrink to their proportion. + let prod = hook.pool("m", &capacity, Some("prod")); + assert_eq!(prod.name, "m:prod"); + assert!((prod.share - 0.6 / 1.4).abs() < 1e-12); + assert_eq!( + hook.pool("m", &capacity, Some("other")), + Pool { + name: "m:default_pool".into(), + share: 0.25 + } + ); + assert_eq!( + capacity_hook(None).pool("m", &capacity, Some("prod")).name, + "m:default_pool" + ); + } + + #[rstest] + #[case(0.625, 1, "62.5%")] + #[case(0.5, 2, "50.00%")] + fn formats_percentages_like_python( + #[case] fraction: f64, + #[case] places: usize, + #[case] expected: &str, + ) { + assert_eq!(percent(fraction, places), expected); + } + + #[rstest] + fn capacities_sum_deployments() { + let config = Config::from_yaml( + "model_list: + - model_name: a + litellm_params: {model: openai/x, rpm: 10, tpm: 1000} + - model_name: a + rpm: 5 + litellm_params: {model: openai/y, rpm: 99} + model_info: {tpm: 500} + - model_name: b + litellm_params: {model: openai/z} +", + ) + .unwrap(); + let groups = capacities(&config); + assert_eq!( + groups["a"], + Capacity { + rpm: Some(15), + tpm: Some(1_500) + } + ); + assert_eq!(groups["b"], Capacity::default()); + } +} diff --git a/litellm-rust/crates/gateway/src/prometheus.rs b/litellm-rust/crates/gateway/src/prometheus.rs new file mode 100644 index 00000000000..28cf7ef0dcd --- /dev/null +++ b/litellm-rust/crates/gateway/src/prometheus.rs @@ -0,0 +1,1253 @@ +//! The `prometheus` callback: the proxy's request metrics, under its metric and label names, +//! served on the gateway's `/metrics`. +//! +//! It is an observer, so nothing here runs on the request path. Each finished request updates a +//! few series per metric: the label values are hashed without allocating, and a series is created +//! once, the first time its label values are seen. The text is built only when Prometheus scrapes. + +use std::{ + collections::{HashMap, HashSet}, + fmt::Write as _, + hash::{BuildHasher, Hash, Hasher, RandomState}, + sync::{ + RwLock, + atomic::{AtomicU64, Ordering}, + }, + time::SystemTime, +}; + +use futures_util::future::BoxFuture; +use litellm_config::Config; +use litellm_gateway_hooks::{Facts, Observer, Record, ServiceTier}; +use serde::Deserialize; + +use crate::{Error, metrics::escape, spend::logged_model}; + +/// The names a config's callback lists turn this observer on with. +const NAMES: [&str; 1] = ["prometheus"]; + +/// The proxy's default latency buckets, in seconds; `+Inf` is implied. +const LATENCY_BUCKETS: [f64; 17] = [ + 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 420.0, + 600.0, +]; + +/// The Prometheus client's default buckets, which the proxy keeps for latency per output token. +const DEFAULT_BUCKETS: [f64; 14] = [ + 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.25, 0.5, 0.75, 1.0, 2.5, 5.0, 7.5, 10.0, +]; + +/// A label the proxy puts on its request metrics. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)] +#[strum(serialize_all = "snake_case")] +enum Label { + EndUser, + User, + UserEmail, + HashedApiKey, + ApiKeyAlias, + Team, + TeamAlias, + RequestedModel, + ModelGroup, + Model, + LitellmModelName, + ModelId, + ApiBase, + ApiProvider, + StatusCode, + ExceptionStatus, + ExceptionClass, + Route, + ClientIp, + UserAgent, + ServiceTier, +} + +use Label::*; + +const LLM_API_LATENCY: &[Label] = &[ + ModelGroup, + Model, + HashedApiKey, + ApiKeyAlias, + Team, + TeamAlias, + RequestedModel, + EndUser, + User, + ModelId, + ApiProvider, + ServiceTier, +]; +const REQUEST_TOTAL_LATENCY: &[Label] = &[ + ModelGroup, + EndUser, + HashedApiKey, + ApiKeyAlias, + RequestedModel, + Team, + TeamAlias, + User, + Model, + ModelId, + ApiProvider, + ServiceTier, +]; +const PROXY_TOTAL_REQUESTS: &[Label] = &[ + EndUser, + HashedApiKey, + ApiKeyAlias, + RequestedModel, + Team, + TeamAlias, + User, + StatusCode, + UserEmail, + Route, + ClientIp, + UserAgent, + ModelId, + ApiProvider, +]; +const PROXY_FAILED_REQUESTS: &[Label] = &[ + EndUser, + HashedApiKey, + ApiKeyAlias, + RequestedModel, + Team, + TeamAlias, + User, + UserEmail, + ExceptionStatus, + ExceptionClass, + Route, + ClientIp, + UserAgent, + ModelId, + ApiProvider, +]; +const LATENCY_PER_OUTPUT_TOKEN: &[Label] = &[ + ModelGroup, + LitellmModelName, + ModelId, + ApiBase, + ApiProvider, + HashedApiKey, + ApiKeyAlias, + Team, + TeamAlias, +]; +const REQUESTS: &[Label] = &[ + EndUser, + HashedApiKey, + ApiKeyAlias, + Model, + Team, + TeamAlias, + User, + UserEmail, + ClientIp, + UserAgent, + RequestedModel, + ModelId, + ApiProvider, +]; +const SPEND: &[Label] = &[ + EndUser, + HashedApiKey, + ApiKeyAlias, + Model, + Team, + TeamAlias, + User, + UserEmail, + ClientIp, + UserAgent, + RequestedModel, + ModelId, + ApiProvider, + ServiceTier, +]; +const TOKENS: &[Label] = &[ + EndUser, + HashedApiKey, + ApiKeyAlias, + Model, + Team, + TeamAlias, + User, + UserEmail, + RequestedModel, + ModelId, + ApiProvider, +]; +const DEPLOYMENT_STATE: &[Label] = &[LitellmModelName, ModelId, ApiBase, ApiProvider]; +const DEPLOYMENT_REQUESTS: &[Label] = &[ + ModelGroup, + RequestedModel, + LitellmModelName, + ModelId, + ApiBase, + ApiProvider, + HashedApiKey, + ApiKeyAlias, + Team, + TeamAlias, + ClientIp, + UserAgent, +]; +const DEPLOYMENT_FAILURES: &[Label] = &[ + ModelGroup, + RequestedModel, + LitellmModelName, + ModelId, + ApiBase, + ApiProvider, + ExceptionStatus, + ExceptionClass, + HashedApiKey, + ApiKeyAlias, + Team, + TeamAlias, + ClientIp, + UserAgent, +]; +const LLM_API_FAILED_REQUESTS: &[Label] = &[ + EndUser, + HashedApiKey, + ApiKeyAlias, + Model, + Team, + TeamAlias, + User, + ModelId, +]; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Kind { + Counter, + Gauge, + /// Latency histograms, whose buckets `prometheus_latency_buckets` can replace. + Latency, + /// A histogram on the Prometheus client's default buckets. + Histogram, +} + +/// The metrics this observer emits, in the order `/metrics` lists them. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Metric { + ProxyTotalRequests, + ProxyFailedRequests, + RequestTotalLatency, + LlmApiLatency, + LlmApiTimeToFirstToken, + Requests, + Spend, + TotalTokens, + InputTokens, + OutputTokens, + InputCachedTokens, + InputCacheCreationTokens, + LlmApiFailedRequests, + DeploymentState, + DeploymentSuccessResponses, + DeploymentFailureResponses, + DeploymentTotalRequests, + DeploymentLatencyPerOutputToken, +} + +impl Metric { + const ALL: [Self; 18] = [ + Self::ProxyTotalRequests, + Self::ProxyFailedRequests, + Self::RequestTotalLatency, + Self::LlmApiLatency, + Self::LlmApiTimeToFirstToken, + Self::Requests, + Self::Spend, + Self::TotalTokens, + Self::InputTokens, + Self::OutputTokens, + Self::InputCachedTokens, + Self::InputCacheCreationTokens, + Self::LlmApiFailedRequests, + Self::DeploymentState, + Self::DeploymentSuccessResponses, + Self::DeploymentFailureResponses, + Self::DeploymentTotalRequests, + Self::DeploymentLatencyPerOutputToken, + ]; + + fn name(self) -> &'static str { + match self { + Self::ProxyTotalRequests => "litellm_proxy_total_requests_metric", + Self::ProxyFailedRequests => "litellm_proxy_failed_requests_metric", + Self::RequestTotalLatency => "litellm_request_total_latency_metric", + Self::LlmApiLatency => "litellm_llm_api_latency_metric", + Self::LlmApiTimeToFirstToken => "litellm_llm_api_time_to_first_token_metric", + Self::Requests => "litellm_requests_metric", + Self::Spend => "litellm_spend_metric", + Self::TotalTokens => "litellm_total_tokens_metric", + Self::InputTokens => "litellm_input_tokens_metric", + Self::OutputTokens => "litellm_output_tokens_metric", + Self::InputCachedTokens => "litellm_input_cached_tokens_metric", + Self::InputCacheCreationTokens => "litellm_input_cache_creation_tokens_metric", + Self::LlmApiFailedRequests => "litellm_llm_api_failed_requests_metric", + Self::DeploymentState => "litellm_deployment_state", + Self::DeploymentSuccessResponses => "litellm_deployment_success_responses", + Self::DeploymentFailureResponses => "litellm_deployment_failure_responses", + Self::DeploymentTotalRequests => "litellm_deployment_total_requests", + Self::DeploymentLatencyPerOutputToken => "litellm_deployment_latency_per_output_token", + } + } + + fn help(self) -> &'static str { + match self { + Self::ProxyTotalRequests => { + "Total number of requests made to the proxy server - track number of client side requests" + } + Self::ProxyFailedRequests => { + "Total number of failed responses from proxy - the client did not get a success response from litellm proxy" + } + Self::RequestTotalLatency => "Total latency (seconds) for a request to LiteLLM", + Self::LlmApiLatency => "Total latency (seconds) for a models LLM API call", + Self::LlmApiTimeToFirstToken => "Time to first token for a models LLM API call", + Self::Requests => { + "deprecated - use litellm_proxy_total_requests_metric. Total number of LLM calls to litellm - track total per API Key, team, user" + } + Self::Spend => "Total spend on LLM requests", + Self::TotalTokens => "Total number of input + output tokens from LLM requests", + Self::InputTokens => "Total number of input tokens from LLM requests", + Self::OutputTokens => "Total number of output tokens from LLM requests", + Self::InputCachedTokens => { + "Total number of cached input tokens read from the provider cache" + } + Self::InputCacheCreationTokens => { + "Total number of input tokens written to the provider cache" + } + Self::LlmApiFailedRequests => "deprecated - use litellm_proxy_failed_requests_metric", + Self::DeploymentState => { + "LLM Deployment Analytics - The state of the deployment: 0 = healthy, 1 = partial outage, 2 = complete outage" + } + Self::DeploymentSuccessResponses => { + "LLM Deployment Analytics - Total number of successful LLM API calls via litellm" + } + Self::DeploymentFailureResponses => { + "LLM Deployment Analytics - Total number of failed LLM API calls for a specific LLM deploymeny. exception_status is the status of the exception from the llm api" + } + Self::DeploymentTotalRequests => { + "LLM Deployment Analytics - Total number of LLM API calls via litellm - success + failure" + } + Self::DeploymentLatencyPerOutputToken => { + "LLM Deployment Analytics - Latency per output token" + } + } + } + + fn kind(self) -> Kind { + match self { + Self::RequestTotalLatency | Self::LlmApiLatency | Self::LlmApiTimeToFirstToken => { + Kind::Latency + } + Self::DeploymentLatencyPerOutputToken => Kind::Histogram, + Self::DeploymentState => Kind::Gauge, + _ => Kind::Counter, + } + } + + fn labels(self) -> &'static [Label] { + match self { + Self::ProxyTotalRequests => PROXY_TOTAL_REQUESTS, + Self::ProxyFailedRequests => PROXY_FAILED_REQUESTS, + Self::RequestTotalLatency => REQUEST_TOTAL_LATENCY, + Self::LlmApiLatency | Self::LlmApiTimeToFirstToken => LLM_API_LATENCY, + Self::Requests => REQUESTS, + Self::Spend => SPEND, + Self::TotalTokens + | Self::InputTokens + | Self::OutputTokens + | Self::InputCachedTokens + | Self::InputCacheCreationTokens => TOKENS, + Self::LlmApiFailedRequests => LLM_API_FAILED_REQUESTS, + Self::DeploymentState => DEPLOYMENT_STATE, + Self::DeploymentSuccessResponses | Self::DeploymentTotalRequests => DEPLOYMENT_REQUESTS, + Self::DeploymentFailureResponses => DEPLOYMENT_FAILURES, + Self::DeploymentLatencyPerOutputToken => LATENCY_PER_OUTPUT_TOKEN, + } + } + + fn from_name(name: &str) -> Option { + Self::ALL.into_iter().find(|metric| metric.name() == name) + } +} + +#[derive(Deserialize)] +struct GroupConfig { + #[serde(default)] + metrics: Vec, + #[serde(default)] + include_labels: Option>, +} + +/// The series that share a hash of their label values, each with those values. +type Bucket = Vec<(Box<[Option>]>, Box)>; + +/// One metric's label names and series. +struct Family { + metric: Metric, + labels: Box<[Label]>, + buckets: Box<[f64]>, + series: RwLock>, +} + +/// A counter or gauge value, or a histogram. +enum Cell { + Value(AtomicF64), + Histogram(Histogram), +} + +#[derive(Default)] +struct AtomicF64(AtomicU64); + +impl AtomicF64 { + fn add(&self, amount: f64) { + let _ = self + .0 + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |bits| { + Some((f64::from_bits(bits) + amount).to_bits()) + }); + } + + fn set(&self, value: f64) { + self.0.store(value.to_bits(), Ordering::Relaxed); + } + + fn get(&self) -> f64 { + f64::from_bits(self.0.load(Ordering::Relaxed)) + } +} + +struct Histogram { + /// Per-bucket counts, not cumulative; the last slot is `+Inf`. + buckets: Box<[AtomicU64]>, + sum: AtomicF64, + count: AtomicU64, +} + +impl Histogram { + fn new(bounds: usize) -> Self { + Self { + buckets: (0..=bounds).map(|_| AtomicU64::new(0)).collect(), + sum: AtomicF64::default(), + count: AtomicU64::new(0), + } + } + + fn observe(&self, bounds: &[f64], value: f64) { + let index = bounds + .iter() + .position(|bound| value <= *bound) + .unwrap_or(bounds.len()); + self.buckets[index].fetch_add(1, Ordering::Relaxed); + self.sum.add(value); + self.count.fetch_add(1, Ordering::Relaxed); + } +} + +impl Family { + fn new(metric: Metric, labels: Vec