From e2691a70318c3cd36825107c644b384e5e930fc8 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:43 -0700 Subject: [PATCH 01/10] feat(gateway): hooks and observers around inference requests, with usage and cost (#47032) * feat(gateway): hooks and observers around inference requests, with usage and cost Adds the litellm-gateway-hooks crate: request facts built only when something reads them, on-path hooks that declare what they read and what they may change and run in config order, off-path observers on bounded queues, a usage tap for SSE and JSON responses, and per-deployment prices from the model cost map overlaid by config prices. The messages and chat handlers run through it when anything is installed; with nothing installed requests take the old path. The key store now also returns the key alias, user, team alias and key and team metadata, and router attempts carry the deployment id. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * fix(model-catalog): read token counts the cost map writes as whole floats The published cost map now has prompt_cache_min_tokens: 4096.0, which the typed u64 fields rejected, so the gateway dropped the whole map at startup and priced every request at zero. Counts accept any whole non-negative number, as Python reads any JSON number, and a test keeps the checked-in map parsing. --------- Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/Cargo.lock | 22 + litellm-rust/Cargo.toml | 3 + ...6bd622e7ed0102166fed8b09b452741c107d.json} | 50 +- .../crates/db/src/verification_token.rs | 3 + .../crates/db/tests/verification_token.rs | 11 +- .../crates/gateway-auth/src/authorization.rs | 1 + litellm-rust/crates/gateway-auth/src/keys.rs | 16 + .../crates/gateway-auth/src/virtual_keys.rs | 3 + litellm-rust/crates/gateway-hooks/Cargo.toml | 22 + .../crates/gateway-hooks/src/facts.rs | 293 ++++++++++++ litellm-rust/crates/gateway-hooks/src/hook.rs | 320 +++++++++++++ litellm-rust/crates/gateway-hooks/src/lib.rs | 22 + .../crates/gateway-hooks/src/observer.rs | 211 ++++++++ .../crates/gateway-hooks/src/pipeline.rs | 320 +++++++++++++ .../crates/gateway-hooks/src/pricing.rs | 376 +++++++++++++++ .../crates/gateway-hooks/src/usage.rs | 452 ++++++++++++++++++ .../crates/gateway-inference/Cargo.toml | 3 + .../gateway-inference/src/chat_completions.rs | 45 +- .../crates/gateway-inference/src/error.rs | 43 +- .../crates/gateway-inference/src/hooks.rs | 379 +++++++++++++++ .../crates/gateway-inference/src/lib.rs | 17 +- .../crates/gateway-inference/src/messages.rs | 38 +- litellm-rust/crates/gateway/Cargo.toml | 2 + litellm-rust/crates/gateway/src/hooks.rs | 217 +++++++++ litellm-rust/crates/gateway/src/key_status.rs | 50 +- litellm-rust/crates/gateway/src/lib.rs | 1 + litellm-rust/crates/gateway/src/main.rs | 5 +- litellm-rust/crates/gateway/tests/hooks.rs | 395 +++++++++++++++ .../crates/gateway/tests/shared_state.rs | 1 + .../crates/gateway/tests/virtual_keys.rs | 4 + .../crates/model-catalog/src/model_info.rs | 81 +++- .../tests/registry_validation.rs | 10 + litellm-rust/crates/router/src/lib.rs | 3 + .../docs/customer-proxy/hooks-pipeline.svg | 75 +++ 34 files changed, 3442 insertions(+), 52 deletions(-) rename litellm-rust/crates/db/.sqlx/{query-6ea28ff09b34b8f392af542369ed7d8f244e60752718b7d2df42e646bc4937be.json => query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json} (76%) create mode 100644 litellm-rust/crates/gateway-hooks/Cargo.toml create mode 100644 litellm-rust/crates/gateway-hooks/src/facts.rs create mode 100644 litellm-rust/crates/gateway-hooks/src/hook.rs create mode 100644 litellm-rust/crates/gateway-hooks/src/lib.rs create mode 100644 litellm-rust/crates/gateway-hooks/src/observer.rs create mode 100644 litellm-rust/crates/gateway-hooks/src/pipeline.rs create mode 100644 litellm-rust/crates/gateway-hooks/src/pricing.rs create mode 100644 litellm-rust/crates/gateway-hooks/src/usage.rs create mode 100644 litellm-rust/crates/gateway-inference/src/hooks.rs create mode 100644 litellm-rust/crates/gateway/src/hooks.rs create mode 100644 litellm-rust/crates/gateway/tests/hooks.rs create mode 100644 litellm-rust/docs/customer-proxy/hooks-pipeline.svg diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c47e5b78af0..9b60fb5520e 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3757,12 +3757,14 @@ dependencies = [ "litellm-config", "litellm-db", "litellm-gateway-auth", + "litellm-gateway-hooks", "litellm-gateway-inference", "litellm-gateway-mcp", "litellm-gateway-ui", "litellm-http", "litellm-inference", "litellm-llms", + "litellm-model-catalog", "litellm-router", "litellm-secrets", "litellm-tracing", @@ -3806,6 +3808,23 @@ dependencies = [ "veil", ] +[[package]] +name = "litellm-gateway-hooks" +version = "0.1.0" +dependencies = [ + "bytes", + "futures-util", + "http 1.4.2", + "litellm-cost", + "litellm-model-catalog", + "memchr", + "rstest", + "serde", + "serde_json", + "tokio", + "tracing", +] + [[package]] name = "litellm-gateway-inference" version = "0.1.0" @@ -3818,6 +3837,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-gateway-auth", + "litellm-gateway-hooks", "litellm-host", "litellm-host-http", "litellm-http", @@ -3831,12 +3851,14 @@ dependencies = [ "litellm-llms-types", "litellm-router", "litellm-secrets", + "memchr", "rstest", "serde", "serde_json", "thiserror 2.0.19", "tokio", "tower", + "uuid", "wiremock", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index f04eb115071..2690c4e373b 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -28,6 +28,9 @@ litellm-gateway-inference = { path = "crates/gateway-inference" } litellm-gateway-auth = { path = "crates/gateway-auth" } litellm-gateway-management = { path = "crates/gateway-management" } litellm-gateway-ui = { path = "crates/gateway-ui" } +litellm-gateway-hooks = { path = "crates/gateway-hooks" } +litellm-cost = { path = "crates/cost" } +litellm-model-catalog = { path = "crates/model-catalog" } litellm-coroutine = { path = "crates/coroutine" } litellm-host = { path = "crates/host" } litellm-host-http = { path = "crates/host-http" } diff --git a/litellm-rust/crates/db/.sqlx/query-6ea28ff09b34b8f392af542369ed7d8f244e60752718b7d2df42e646bc4937be.json b/litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json similarity index 76% rename from litellm-rust/crates/db/.sqlx/query-6ea28ff09b34b8f392af542369ed7d8f244e60752718b7d2df42e646bc4937be.json rename to litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json index 090951f4462..0283a5854a0 100644 --- a/litellm-rust/crates/db/.sqlx/query-6ea28ff09b34b8f392af542369ed7d8f244e60752718b7d2df42e646bc4937be.json +++ b/litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n v.max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL\n OR v.max_parallel_requests IS NOT NULL\n OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL\n OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL\n OR t.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL\n OR mb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL\n OR u.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", + "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n v.max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL\n OR v.max_parallel_requests IS NOT NULL\n OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL\n OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL\n OR t.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL\n OR mb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL\n OR u.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", "describe": { "columns": [ { @@ -180,6 +180,17 @@ }, { "ordinal": 20, + "name": "team_alias?", + "type_info": "Text", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamTable\"", + "name": "team_alias" + } + } + }, + { + "ordinal": 21, "name": "team_blocked?", "type_info": "Bool", "origin": { @@ -190,7 +201,7 @@ } }, { - "ordinal": 21, + "ordinal": 22, "name": "team_models?", "type_info": "TextArray", "origin": { @@ -201,7 +212,7 @@ } }, { - "ordinal": 22, + "ordinal": 23, "name": "team_spend?", "type_info": "Float8", "origin": { @@ -212,7 +223,7 @@ } }, { - "ordinal": 23, + "ordinal": 24, "name": "team_max_budget?", "type_info": "Float8", "origin": { @@ -223,7 +234,7 @@ } }, { - "ordinal": 24, + "ordinal": 25, "name": "team_metadata?", "type_info": "Jsonb", "origin": { @@ -234,43 +245,43 @@ } }, { - "ordinal": 25, + "ordinal": 26, "name": "team_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 26, + "ordinal": 27, "name": "team_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 27, + "ordinal": 28, "name": "team_model_aliases!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 28, + "ordinal": 29, "name": "team_object_permissions!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 29, + "ordinal": 30, "name": "team_organization!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 30, + "ordinal": 31, "name": "member_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 31, + "ordinal": 32, "name": "member_allowed_models?", "type_info": "TextArray", "origin": { @@ -281,7 +292,7 @@ } }, { - "ordinal": 32, + "ordinal": 33, "name": "user_id", "type_info": "Text", "origin": { @@ -292,7 +303,7 @@ } }, { - "ordinal": 33, + "ordinal": 34, "name": "user_row_id?", "type_info": "Text", "origin": { @@ -303,7 +314,7 @@ } }, { - "ordinal": 34, + "ordinal": 35, "name": "user_models?", "type_info": "TextArray", "origin": { @@ -314,19 +325,19 @@ } }, { - "ordinal": 35, + "ordinal": 36, "name": "user_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 36, + "ordinal": 37, "name": "user_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 37, + "ordinal": 38, "name": "user_object_permissions!", "type_info": "Bool", "origin": "Expression" @@ -363,6 +374,7 @@ true, true, true, + true, null, null, null, @@ -378,5 +390,5 @@ null ] }, - "hash": "6ea28ff09b34b8f392af542369ed7d8f244e60752718b7d2df42e646bc4937be" + "hash": "aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d" } diff --git a/litellm-rust/crates/db/src/verification_token.rs b/litellm-rust/crates/db/src/verification_token.rs index 5ab6b088bf3..aa48a881767 100644 --- a/litellm-rust/crates/db/src/verification_token.rs +++ b/litellm-rust/crates/db/src/verification_token.rs @@ -32,6 +32,7 @@ pub struct VerificationToken { #[derive(Clone, Debug, PartialEq)] pub struct Team { + pub alias: Option, pub blocked: bool, pub models: Vec, pub spend: f64, @@ -108,6 +109,7 @@ pub async fn find_verification_token( (v.agent_id IS NOT NULL) AS "key_agent!", v.team_id, t.team_id AS "team_row_id?", + t.team_alias AS "team_alias?", t.blocked AS "team_blocked?", t.models AS "team_models?", t.spend AS "team_spend?", @@ -174,6 +176,7 @@ pub async fn find_verification_token( }, team_id: row.team_id, team: row.team_row_id.map(|_| Team { + alias: row.team_alias, blocked: row.team_blocked.unwrap_or(false), models: row.team_models.unwrap_or_default(), spend: row.team_spend.unwrap_or(0.0), diff --git a/litellm-rust/crates/db/tests/verification_token.rs b/litellm-rust/crates/db/tests/verification_token.rs index 89bd29e1cc7..169429f4839 100644 --- a/litellm-rust/crates/db/tests/verification_token.rs +++ b/litellm-rust/crates/db/tests/verification_token.rs @@ -8,11 +8,11 @@ use rstest::rstest; use serde_json::json; const FIXTURES: &str = r#" -INSERT INTO "LiteLLM_TeamTable" (team_id, blocked, models, spend, max_budget, metadata, tpm_limit, - model_id, created_at, updated_at) - VALUES ('team-1', true, ARRAY['gpt-a'], 2.5, 10.0, '{"require_trace_id": true}', 100, NULL, - now(), now()), - ('team-2', false, ARRAY[]::text[], 0, NULL, '{}', NULL, NULL, now(), now()); +INSERT INTO "LiteLLM_TeamTable" (team_id, team_alias, blocked, models, spend, max_budget, metadata, + tpm_limit, model_id, created_at, updated_at) + VALUES ('team-1', 'research', true, ARRAY['gpt-a'], 2.5, 10.0, '{"require_trace_id": true}', 100, + NULL, now(), now()), + ('team-2', NULL, false, ARRAY[]::text[], 0, NULL, '{}', NULL, NULL, now(), now()); INSERT INTO "LiteLLM_UserTable" (user_id, models, rpm_limit, max_budget, policies) VALUES ('user-1', ARRAY['gpt-b'], 5, 1.0, ARRAY['p']), ('user-2', ARRAY[]::text[], NULL, 2.0, ARRAY[]::text[]); @@ -113,6 +113,7 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { restrictions: Restrictions::default(), team_id: Some("team-1".into()), team: Some(Team { + alias: Some("research".into()), blocked: true, models: vec!["gpt-a".into()], spend: 2.5, diff --git a/litellm-rust/crates/gateway-auth/src/authorization.rs b/litellm-rust/crates/gateway-auth/src/authorization.rs index 86e43b7f361..fcb218dadfe 100644 --- a/litellm-rust/crates/gateway-auth/src/authorization.rs +++ b/litellm-rust/crates/gateway-auth/src/authorization.rs @@ -59,6 +59,7 @@ pub struct KeyGrant { /// The key's team. A deployment owned by a team serves only that team's keys, as the /// proxy's `filter_team_based_models` does; a teamless key gets shared deployments only. pub team_id: Option, + pub profile: Arc, } impl KeyGrant { diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 99042aa6d81..2ce5338a674 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -1,6 +1,7 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use litellm_auth_types::SecretValue; +use serde_json::{Map, Value}; use crate::{KeyError, hash_token}; @@ -288,6 +289,16 @@ pub enum Owner { Proxy, } +/// Who a key belongs to beyond what gates it, for hooks and spend records. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct KeyProfile { + pub alias: Option, + pub user_id: Option, + pub team_alias: Option, + pub metadata: Arc>, + pub team_metadata: Arc>, +} + #[derive(Clone, Debug, Eq, PartialEq)] pub enum KeyStatus { Active { @@ -296,6 +307,7 @@ pub enum KeyStatus { allowed_routes: AllowedRoutes, /// The key's team, whose team-scoped deployments the key may use. team_id: Option, + profile: Arc, }, Blocked, Revoked, @@ -313,6 +325,7 @@ impl KeyStatus { models: models.into(), allowed_routes: AllowedRoutes::default(), team_id: None, + profile: Arc::default(), } } } @@ -325,6 +338,7 @@ pub struct VerifiedKey { pub models: KeyModels, pub allowed_routes: AllowedRoutes, pub team_id: Option, + pub profile: Arc, } pub type LookupFuture<'a> = @@ -382,6 +396,7 @@ pub async fn verify_key( models, allowed_routes, team_id, + profile, }) => { if expired(expires_at) { return Err(KeyError::Expired); @@ -392,6 +407,7 @@ pub async fn verify_key( models, allowed_routes, team_id, + profile, }) } Some(KeyStatus::Refused { diff --git a/litellm-rust/crates/gateway-auth/src/virtual_keys.rs b/litellm-rust/crates/gateway-auth/src/virtual_keys.rs index 6f7c45fb571..e44aa9b025c 100644 --- a/litellm-rust/crates/gateway-auth/src/virtual_keys.rs +++ b/litellm-rust/crates/gateway-auth/src/virtual_keys.rs @@ -288,6 +288,7 @@ impl Authenticator for VirtualKeyAuthenticator { models: key.models, routes: key.allowed_routes, team_id: key.team_id, + profile: key.profile, }), }) }) @@ -415,6 +416,7 @@ mod tests { }, allowed_routes: AllowedRoutes::default(), team_id: Some("team-1".into()), + profile: Default::default(), }, ); store.put( @@ -424,6 +426,7 @@ mod tests { models: KeyModels::default(), allowed_routes: AllowedRoutes::new(strings(&["anthropic_routes"])), team_id: None, + profile: Default::default(), }, ); store.put( diff --git a/litellm-rust/crates/gateway-hooks/Cargo.toml b/litellm-rust/crates/gateway-hooks/Cargo.toml new file mode 100644 index 00000000000..b0d6c67f15f --- /dev/null +++ b/litellm-rust/crates/gateway-hooks/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "litellm-gateway-hooks" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +bytes.workspace = true +futures-util.workspace = true +http.workspace = true +litellm-cost.workspace = true +litellm-model-catalog.workspace = true +memchr = "2" +serde.workspace = true +serde_json.workspace = true +tokio = { workspace = true, features = ["sync"] } +tracing.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio = { workspace = true, features = ["sync", "macros", "rt-multi-thread", "time"] } diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs new file mode 100644 index 00000000000..3e7e2033e13 --- /dev/null +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -0,0 +1,293 @@ +use std::{ + ops::{BitOr, BitOrAssign}, + sync::{Arc, OnceLock}, + time::{Duration, SystemTime}, +}; + +use bytes::Bytes; +use http::HeaderMap; +use serde_json::{Map, Value}; + +use crate::usage::Usage; + +/// The facts a hook reads or an observer asks for. The gateway builds only what some hook or +/// observer declares, so a fact nobody reads costs nothing. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Facts(u16); + +impl Facts { + pub const NONE: Self = Self(0); + /// Who sent the request: the key, its user, team and organization. + pub const KEY: Self = Self(1 << 0); + pub const END_USER: Self = Self(1 << 1); + pub const CLIENT_HEADERS: Self = Self(1 << 2); + /// An estimate of the request's input tokens, made on first read. + pub const INPUT_TOKENS: Self = Self(1 << 3); + /// The deployment that served the request. + pub const DEPLOYMENT: Self = Self(1 << 4); + /// Token usage read from the response as it streams. + pub const USAGE: Self = Self(1 << 5); + /// The response's cost, from usage and the deployment's prices. + pub const COST: Self = Self(1 << 6); + pub const TIMINGS: Self = Self(1 << 7); + /// Values earlier hooks attached for observers. + pub const ANNOTATIONS: Self = Self(1 << 8); + /// The raw request bytes, already held for retries. + pub const REQUEST_BODY: Self = Self(1 << 9); + /// The response text. The only fact that makes the gateway keep a stream's bytes. + pub const RESPONSE_BODY: Self = Self(1 << 10); + + pub const fn union(self, other: Self) -> Self { + Self(self.0 | other.0) + } + + pub const fn contains(self, other: Self) -> bool { + self.0 & other.0 == other.0 + } + + pub const fn is_empty(self) -> bool { + self.0 == 0 + } +} + +impl BitOr for Facts { + type Output = Self; + + fn bitor(self, other: Self) -> Self { + self.union(other) + } +} + +impl BitOrAssign for Facts { + fn bitor_assign(&mut self, other: Self) { + *self = self.union(other); + } +} + +/// The endpoint a request came in on. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Route { + Messages, + ChatCompletions, +} + +impl Route { + /// The proxy's `call_type` for this endpoint, as spend logs record it. + pub fn call_type(self) -> &'static str { + match self { + Self::Messages => "anthropic_messages", + Self::ChatCompletions => "acompletion", + } + } + + /// The endpoint the proxy's daily spend tables key on. + pub fn endpoint(self) -> &'static str { + match self { + Self::Messages => "/v1/messages", + Self::ChatCompletions => "/chat/completions", + } + } +} + +/// Who sent a request, as the key store knows it. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct KeyFacts { + /// The sha256 hex of the key, `LiteLLM_VerificationToken.token`. + pub hash: String, + pub alias: Option, + pub user_id: Option, + pub team_id: Option, + pub team_alias: Option, + pub org_id: Option, + pub metadata: Arc>, + pub team_metadata: Arc>, +} + +/// What the gateway knows about a request before it is sent, shared by every hook and observer. +#[derive(Debug)] +pub struct RequestFacts { + pub request_id: Arc, + pub route: Route, + pub key: Option>, + pub end_user: Option>, + /// The public model name the client asked for. + pub requested_model: Arc, + pub client_headers: Arc, + pub stream: bool, + /// The output budget the request set, if any. + pub max_tokens: Option, + pub started_at: SystemTime, + body: Bytes, + input_tokens: OnceLock, +} + +impl RequestFacts { + #[allow(clippy::too_many_arguments)] + pub fn new( + request_id: Arc, + route: Route, + key: Option>, + end_user: Option>, + requested_model: Arc, + client_headers: Arc, + stream: bool, + max_tokens: Option, + body: Bytes, + ) -> Self { + Self { + request_id, + route, + key, + end_user, + requested_model, + client_headers, + stream, + max_tokens, + started_at: SystemTime::now(), + body, + input_tokens: OnceLock::new(), + } + } + + /// The raw request bytes. Shared, never copied. + pub fn body(&self) -> &Bytes { + &self.body + } + + /// The proxy's cheap input estimate, a token per four bytes of request, made once. + pub fn input_tokens(&self) -> u64 { + *self + .input_tokens + .get_or_init(|| (self.body.len() as u64).div_ceil(4).max(1)) + } +} + +/// The deployment a request was served by. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct DeploymentFacts { + /// `model_info.id`, or the id the router derives for a deployment without one. + pub id: Arc, + pub model_group: Arc, + pub model: String, + pub provider: Option, + pub api_base: Option, +} + +impl DeploymentFacts { + /// The provider, from `custom_llm_provider` or the model's `provider/` prefix. + pub fn provider(&self) -> Option<&str> { + self.provider + .as_deref() + .or_else(|| self.model.split_once('/').map(|(provider, _)| provider)) + } +} + +/// Values hooks attach for observers, in the order they attached them. Later writes to the +/// same key replace earlier ones. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Annotations(Vec<(&'static str, Value)>); + +impl Annotations { + pub fn set(&mut self, key: &'static str, value: Value) { + match self.0.iter_mut().find(|(name, _)| *name == key) { + Some((_, slot)) => *slot = value, + None => self.0.push((key, value)), + } + } + + pub fn get(&self, key: &str) -> Option<&Value> { + self.0 + .iter() + .find(|(name, _)| *name == key) + .map(|(_, value)| value) + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +/// How a request ended. +#[derive(Clone, Debug, PartialEq)] +pub struct Outcome { + pub deployment: Option, + /// The HTTP status the client got. + pub status: u16, + pub usage: Option, + /// In dollars, from usage and the deployment's prices; `None` when either is missing. + pub cost: Option, + pub time_to_first_byte: Option, + pub ended_at: SystemTime, + /// The stream stopped before its end: the client went away or the upstream failed. + pub interrupted: bool, +} + +/// Everything an observer can be handed about one request. +#[derive(Debug)] +pub struct Record { + pub request: Arc, + pub outcome: Outcome, + pub annotations: Annotations, + /// Present only when an observer reads [`Facts::RESPONSE_BODY`]. + pub response_body: Option, +} + +impl Record { + pub fn key(&self) -> Option<&KeyFacts> { + self.request.key.as_deref() + } + + pub fn succeeded(&self) -> bool { + (200..300).contains(&self.outcome.status) && !self.outcome.interrupted + } + + pub fn duration(&self) -> Duration { + self.outcome + .ended_at + .duration_since(self.request.started_at) + .unwrap_or_default() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + fn facts_combine_and_contain() { + let facts = Facts::KEY | Facts::USAGE; + assert!(facts.contains(Facts::KEY)); + assert!(facts.contains(Facts::KEY | Facts::USAGE)); + assert!(!facts.contains(Facts::COST)); + assert!(Facts::NONE.is_empty()); + } + + #[rstest] + fn later_annotations_replace_earlier_ones() { + let mut annotations = Annotations::default(); + annotations.set("a", Value::from(1)); + annotations.set("b", Value::from(2)); + annotations.set("a", Value::from(3)); + assert_eq!(annotations.get("a"), Some(&Value::from(3))); + assert_eq!(annotations.get("b"), Some(&Value::from(2))); + } + + #[rstest] + #[case::prefix(None, "anthropic/claude-x", Some("anthropic"))] + #[case::explicit(Some("bedrock"), "anthropic.claude", Some("bedrock"))] + #[case::bare(None, "gpt-x", None)] + fn provider_comes_from_the_setting_or_the_model_prefix( + #[case] provider: Option<&str>, + #[case] model: &str, + #[case] expected: Option<&str>, + ) { + let deployment = DeploymentFacts { + model: model.into(), + provider: provider.map(str::to_owned), + ..DeploymentFacts::default() + }; + assert_eq!(deployment.provider(), expected); + } +} diff --git a/litellm-rust/crates/gateway-hooks/src/hook.rs b/litellm-rust/crates/gateway-hooks/src/hook.rs new file mode 100644 index 00000000000..0392cd0f88f --- /dev/null +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -0,0 +1,320 @@ +//! Hooks run on the request path, in the order the config lists them. Each declares the facts it +//! reads and the parts of the request it may change; several hooks may change the same part, and +//! the later one in the config wins. + +use std::{ + borrow::Cow, + ops::{BitOr, BitOrAssign}, + sync::Arc, +}; + +use futures_util::future::BoxFuture; +use http::{HeaderName, HeaderValue}; +use serde_json::{Map, Value}; + +use crate::facts::{Annotations, Facts, Record, RequestFacts}; + +/// The parts of a request a hook may change. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Targets(u8); + +impl Targets { + pub const NONE: Self = Self(0); + /// Turn the request away before it is sent. + pub const ADMISSION: Self = Self(1 << 0); + /// Headers sent to the provider. + pub const UPSTREAM_HEADERS: Self = Self(1 << 1); + /// Top-level fields of the body sent to the provider. + pub const UPSTREAM_BODY: Self = Self(1 << 2); + /// Headers returned to the client. + pub const RESPONSE_HEADERS: Self = Self(1 << 3); + /// Values handed to observers with the request's record. + pub const ANNOTATIONS: Self = Self(1 << 4); + + pub const fn union(self, other: Self) -> Self { + Self(self.0 | other.0) + } + + pub const fn contains(self, other: Self) -> bool { + self.0 & other.0 == other.0 + } +} + +impl BitOr for Targets { + type Output = Self; + + fn bitor(self, other: Self) -> Self { + self.union(other) + } +} + +impl BitOrAssign for Targets { + fn bitor_assign(&mut self, other: Self) { + *self = self.union(other); + } +} + +/// Matches header names, case-insensitively. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum HeaderPattern { + Exact(HeaderName), + /// Every header whose name starts with this lowercase prefix. + Prefix(String), +} + +impl HeaderPattern { + /// `x-foo-*` is a prefix pattern; anything else is one header. + pub fn parse(pattern: &str) -> Result { + match pattern.strip_suffix('*') { + Some(prefix) => Ok(Self::Prefix(prefix.to_ascii_lowercase())), + None => HeaderName::from_bytes(pattern.as_bytes()).map(Self::Exact), + } + } + + pub fn matches(&self, name: &str) -> bool { + match self { + Self::Exact(exact) => exact.as_str().eq_ignore_ascii_case(name), + Self::Prefix(prefix) => name + .get(..prefix.len()) + .is_some_and(|start| start.eq_ignore_ascii_case(prefix)), + } + } +} + +#[derive(Clone, Debug, PartialEq)] +enum HeaderOp { + Set(HeaderName, HeaderValue), + Remove(HeaderPattern), +} + +/// Why a hook turned a request away. The client gets `status`, `message` and `headers`. +#[derive(Clone, Debug, PartialEq)] +pub struct Refusal { + pub status: u16, + /// The proxy's error `type`, such as `budget_exceeded`. + pub kind: Cow<'static, str>, + pub message: String, + pub headers: Vec<(HeaderName, HeaderValue)>, +} + +impl Refusal { + pub fn new( + status: u16, + kind: impl Into>, + message: impl Into, + ) -> Self { + Self { + status, + kind: kind.into(), + message: message.into(), + headers: Vec::new(), + } + } + + pub fn with_header(mut self, name: HeaderName, value: HeaderValue) -> Self { + self.headers.push((name, value)); + self + } +} + +/// The changes hooks asked for, kept in the order they were made. A hook may only make the +/// changes its [`Hook::writes`] declares; anything else is dropped with a warning. +#[derive(Debug, Default)] +pub struct RequestEdit { + allowed: Targets, + writer: &'static str, + header_ops: Vec, + body_fields: Vec<(String, Value)>, + response_headers: Vec<(HeaderName, HeaderValue)>, + annotations: Annotations, +} + +impl RequestEdit { + pub(crate) fn scope(&mut self, writer: &'static str, allowed: Targets) { + self.writer = writer; + self.allowed = allowed; + } + + fn may(&self, target: Targets) -> bool { + let allowed = self.allowed.contains(target); + if !allowed { + tracing::warn!( + hook = self.writer, + ?target, + "hook wrote a target it does not declare; ignored" + ); + } + allowed + } + + pub fn set_upstream_header(&mut self, name: HeaderName, value: HeaderValue) { + if self.may(Targets::UPSTREAM_HEADERS) { + self.header_ops.push(HeaderOp::Set(name, value)); + } + } + + pub fn remove_upstream_headers(&mut self, pattern: HeaderPattern) { + if self.may(Targets::UPSTREAM_HEADERS) { + self.header_ops.push(HeaderOp::Remove(pattern)); + } + } + + /// 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)); + } + } + + pub fn add_response_header(&mut self, name: HeaderName, value: HeaderValue) { + if self.may(Targets::RESPONSE_HEADERS) { + self.response_headers.push((name, value)); + } + } + + pub fn annotate(&mut self, key: &'static str, value: Value) { + if self.may(Targets::ANNOTATIONS) { + self.annotations.set(key, value); + } + } + + /// Whether any change reaches the provider. + pub fn changes_upstream(&self) -> bool { + !self.header_ops.is_empty() || !self.body_fields.is_empty() + } + + /// Applies the header changes in order to the headers about to be sent. + pub fn apply_upstream_headers(&self, headers: &mut Vec<(String, String)>) { + for op in &self.header_ops { + match op { + HeaderOp::Set(name, value) => { + let Ok(value) = value.to_str() else { continue }; + headers.retain(|(existing, _)| !existing.eq_ignore_ascii_case(name.as_str())); + headers.push((name.as_str().to_owned(), value.to_owned())); + } + HeaderOp::Remove(pattern) => headers.retain(|(name, _)| !pattern.matches(name)), + } + } + } + + /// 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()); + } + } + + pub fn response_headers(&self) -> &[(HeaderName, HeaderValue)] { + &self.response_headers + } + + pub(crate) fn take_annotations(&mut self) -> Annotations { + std::mem::take(&mut self.annotations) + } +} + +/// Work a hook left to do once the response ends, such as returning an unused reservation. +pub trait Settle: Send + Sync + 'static { + fn settle(self: Box, record: &Record); +} + +/// An on-path hook. +pub trait Hook: Send + Sync + 'static { + /// The name the config uses for this hook. + fn name(&self) -> &'static str; + + /// The facts this hook reads. + fn reads(&self) -> Facts; + + /// The parts of the request this hook may change. + fn writes(&self) -> Targets; + + /// Runs before the request is routed. Returning a [`Settle`] asks to be told how the request + /// ended. + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>>; +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn headers(pairs: &[(&str, &str)]) -> Vec<(String, String)> { + pairs + .iter() + .map(|(name, value)| ((*name).to_owned(), (*value).to_owned())) + .collect() + } + + #[rstest] + #[case::exact("x-a", "X-A", true)] + #[case::exact_other("x-a", "x-ab", false)] + #[case::prefix("x-internal-*", "X-Internal-Trace", true)] + #[case::prefix_short("x-internal-*", "x-intern", false)] + fn header_patterns_ignore_case(#[case] pattern: &str, #[case] name: &str, #[case] hit: bool) { + assert_eq!(HeaderPattern::parse(pattern).unwrap().matches(name), hit); + } + + #[rstest] + fn header_changes_apply_in_order() { + let mut edit = RequestEdit::default(); + edit.scope("first", Targets::UPSTREAM_HEADERS); + edit.set_upstream_header( + HeaderName::from_static("x-a"), + HeaderValue::from_static("1"), + ); + edit.scope("second", Targets::UPSTREAM_HEADERS); + edit.remove_upstream_headers(HeaderPattern::parse("x-*").unwrap()); + edit.set_upstream_header( + HeaderName::from_static("x-b"), + HeaderValue::from_static("2"), + ); + + let mut sent = headers(&[("X-A", "client"), ("content-type", "application/json")]); + edit.apply_upstream_headers(&mut sent); + assert_eq!( + sent, + headers(&[("content-type", "application/json"), ("x-b", "2")]) + ); + } + + #[rstest] + fn later_hooks_win_on_the_same_body_field() { + let mut edit = RequestEdit::default(); + edit.scope("first", Targets::UPSTREAM_BODY); + edit.set_upstream_body_field("metadata", Value::from("a")); + edit.scope("second", Targets::UPSTREAM_BODY); + edit.set_upstream_body_field("metadata", Value::from("b")); + + let serde_json::Value::Object(mut body) = + serde_json::json!({"model": "m", "metadata": "client"}) + else { + unreachable!() + }; + edit.apply_upstream_body(&mut body); + assert_eq!( + Value::Object(body), + serde_json::json!({"model": "m", "metadata": "b"}) + ); + } + + #[rstest] + fn undeclared_writes_are_dropped() { + let mut edit = RequestEdit::default(); + edit.scope("reader", Targets::ANNOTATIONS); + edit.set_upstream_body_field("metadata", Value::Null); + edit.add_response_header( + HeaderName::from_static("x-a"), + HeaderValue::from_static("1"), + ); + edit.annotate("a", Value::from(1)); + assert!(!edit.changes_upstream()); + assert!(edit.response_headers().is_empty()); + assert_eq!(edit.take_annotations().get("a"), Some(&Value::from(1))); + } +} diff --git a/litellm-rust/crates/gateway-hooks/src/lib.rs b/litellm-rust/crates/gateway-hooks/src/lib.rs new file mode 100644 index 00000000000..d097bc940fd --- /dev/null +++ b/litellm-rust/crates/gateway-hooks/src/lib.rs @@ -0,0 +1,22 @@ +//! Hooks and observers for the gateway. +//! +//! A hook runs on the request path. It declares the facts it reads and the parts of the request +//! it may change, and runs in config order with the others. An observer sees each request after +//! its response ends, off the request path, and declares the facts it reads. The gateway builds +//! only the facts something reads, so with nothing configured a request pays nothing. + +mod facts; +mod hook; +mod observer; +mod pipeline; +mod pricing; +mod usage; + +pub use facts::{ + Annotations, DeploymentFacts, Facts, KeyFacts, Outcome, Record, RequestFacts, Route, +}; +pub use hook::{HeaderPattern, Hook, Refusal, RequestEdit, Settle, Targets}; +pub use observer::{DEFAULT_QUEUE, Observer, Observers}; +pub use pipeline::{Admission, Pipeline}; +pub use pricing::{Prices, deployment_fields}; +pub use usage::{ServiceTier, StreamFormat, StreamUsage, Usage}; diff --git a/litellm-rust/crates/gateway-hooks/src/observer.rs b/litellm-rust/crates/gateway-hooks/src/observer.rs new file mode 100644 index 00000000000..6a8b214d5c3 --- /dev/null +++ b/litellm-rust/crates/gateway-hooks/src/observer.rs @@ -0,0 +1,211 @@ +//! Observers see each request after its response ends, off the request path. Each runs on its own +//! task behind a bounded queue; when the queue is full the record is dropped and counted, so a +//! slow observer never slows a request. + +use std::sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, +}; + +use futures_util::future::BoxFuture; +use tokio::{sync::mpsc, task::JoinHandle}; + +use crate::facts::{Facts, Record}; + +/// Records an observer may have waiting before new ones are dropped. +pub const DEFAULT_QUEUE: usize = 10_000; + +/// An off-path consumer of finished requests. +pub trait Observer: Send + Sync + 'static { + /// The name the config and metrics use for this observer. + fn name(&self) -> &'static str; + + /// The facts this observer reads from each record. + fn reads(&self) -> Facts; + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()>; + + /// Writes out anything held back; called once at shutdown after the last record. + fn flush(&self) -> BoxFuture<'_, ()> { + Box::pin(async {}) + } +} + +/// `None` tells the observer's task to flush and stop. +type Message = Option>; + +struct Lane { + name: &'static str, + sender: mpsc::Sender, + dropped: Arc, +} + +/// Hands records to every observer. +pub struct Observers { + lanes: Vec, + tasks: Mutex>>, + reads: Facts, +} + +impl Observers { + /// Starts a task per observer. Must be called inside a tokio runtime. + pub fn start(observers: Vec>, queue: usize) -> Self { + let mut lanes = Vec::with_capacity(observers.len()); + let mut tasks = Vec::with_capacity(observers.len()); + let mut reads = Facts::NONE; + for observer in observers { + reads |= observer.reads(); + let (sender, mut receiver) = mpsc::channel::(queue.max(1)); + lanes.push(Lane { + name: observer.name(), + sender, + dropped: Arc::new(AtomicU64::new(0)), + }); + tasks.push(tokio::spawn(async move { + while let Some(Some(record)) = receiver.recv().await { + observer.observe(&record).await; + } + observer.flush().await; + })); + } + Self { + lanes, + tasks: Mutex::new(tasks), + reads, + } + } + + pub fn is_empty(&self) -> bool { + self.lanes.is_empty() + } + + /// The facts some observer reads. + pub fn reads(&self) -> Facts { + self.reads + } + + pub fn send(&self, record: Arc) { + for lane in &self.lanes { + if lane.sender.try_send(Some(Arc::clone(&record))).is_err() { + lane.dropped.fetch_add(1, Ordering::Relaxed); + } + } + } + + /// Records each observer has dropped because its queue was full, by observer name. + pub fn dropped(&self) -> impl Iterator + '_ { + self.lanes + .iter() + .map(|lane| (lane.name, lane.dropped.load(Ordering::Relaxed))) + } + + /// Lets every observer finish its queue, then flushes them. Records sent after this are + /// dropped. + pub async fn close(&self) { + for lane in &self.lanes { + let _ = lane.sender.send(None).await; + } + let tasks = std::mem::take(&mut *self.tasks.lock().unwrap_or_else(|e| e.into_inner())); + for task in tasks { + if let Err(error) = task.await { + tracing::error!(%error, "observer task failed"); + } + } + } +} + +#[cfg(test)] +mod tests { + use std::{ + sync::Mutex, + time::{Duration, SystemTime}, + }; + + use bytes::Bytes; + use http::HeaderMap; + + use super::*; + use crate::facts::{Annotations, Outcome, RequestFacts, Route}; + + fn record(status: u16) -> Arc { + Arc::new(Record { + request: Arc::new(RequestFacts::new( + "r".into(), + Route::Messages, + None, + None, + "m".into(), + Arc::new(HeaderMap::new()), + false, + None, + Bytes::new(), + )), + outcome: Outcome { + deployment: None, + status, + usage: None, + cost: None, + time_to_first_byte: None, + ended_at: SystemTime::now(), + interrupted: false, + }, + annotations: Annotations::default(), + response_body: None, + }) + } + + #[derive(Default)] + struct Collect { + seen: Mutex>, + flushed: Mutex, + delay: Duration, + } + + impl Observer for Collect { + fn name(&self) -> &'static str { + "collect" + } + + fn reads(&self) -> Facts { + Facts::USAGE + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + Box::pin(async move { + tokio::time::sleep(self.delay).await; + self.seen.lock().unwrap().push(record.outcome.status); + }) + } + + fn flush(&self) -> BoxFuture<'_, ()> { + Box::pin(async { *self.flushed.lock().unwrap() = true }) + } + } + + #[tokio::test] + async fn close_drains_then_flushes() { + let collect = Arc::new(Collect::default()); + let observers = Observers::start(vec![collect.clone()], 16); + assert_eq!(observers.reads(), Facts::USAGE); + observers.send(record(200)); + observers.send(record(500)); + observers.close().await; + assert_eq!(*collect.seen.lock().unwrap(), vec![200, 500]); + assert!(*collect.flushed.lock().unwrap()); + } + + #[tokio::test] + async fn a_full_queue_drops_and_counts() { + let collect = Arc::new(Collect { + delay: Duration::from_millis(200), + ..Collect::default() + }); + let observers = Observers::start(vec![collect.clone()], 1); + for _ in 0..5 { + observers.send(record(200)); + } + let dropped: u64 = observers.dropped().map(|(_, count)| count).sum(); + assert!(dropped >= 3, "dropped {dropped}"); + observers.close().await; + } +} diff --git a/litellm-rust/crates/gateway-hooks/src/pipeline.rs b/litellm-rust/crates/gateway-hooks/src/pipeline.rs new file mode 100644 index 00000000000..384babaf8d2 --- /dev/null +++ b/litellm-rust/crates/gateway-hooks/src/pipeline.rs @@ -0,0 +1,320 @@ +//! The configured hooks and observers, run for each request. + +use std::{sync::Arc, time::SystemTime}; + +use bytes::Bytes; + +use crate::{ + facts::{Facts, Outcome, Record, RequestFacts}, + hook::{Hook, Refusal, RequestEdit, Settle}, + observer::Observers, + pricing::Prices, +}; + +pub struct Pipeline { + hooks: Vec>, + observers: Observers, + prices: Arc, + reads: Facts, +} + +impl Pipeline { + /// `hooks` run in the order given, which is the config's. + pub fn new(hooks: Vec>, observers: Observers, prices: Arc) -> Self { + let mut reads = observers.reads(); + for hook in &hooks { + reads |= hook.reads(); + } + // Cost is computed from usage and the deployment that served the request. + if reads.contains(Facts::COST) { + reads |= Facts::USAGE | Facts::DEPLOYMENT; + } + Self { + hooks, + observers, + prices, + reads, + } + } + + /// No hooks and no observers: the gateway can skip building facts altogether. + pub fn is_empty(&self) -> bool { + self.hooks.is_empty() && self.observers.is_empty() + } + + /// Whether some hook or observer reads every fact in `facts`. + pub fn reads(&self, facts: Facts) -> bool { + self.reads.contains(facts) + } + + pub fn has_hooks(&self) -> bool { + !self.hooks.is_empty() + } + + pub fn observers(&self) -> &Observers { + &self.observers + } + + pub fn prices(&self) -> &Prices { + &self.prices + } + + /// Runs every hook in order. When one refuses, the hooks before it are told the request + /// ended with that refusal, and so are the observers. + pub async fn admit(&self, request: Arc) -> Result { + let mut admission = Admission { + request, + edit: RequestEdit::default(), + settles: Vec::new(), + }; + for hook in &self.hooks { + admission.edit.scope(hook.name(), hook.writes()); + match hook.run(&admission.request, &mut admission.edit).await { + Ok(Some(settle)) => admission.settles.push(settle), + Ok(None) => {} + Err(refusal) => { + let outcome = Outcome { + deployment: None, + status: refusal.status, + usage: None, + cost: None, + time_to_first_byte: None, + ended_at: SystemTime::now(), + interrupted: false, + }; + self.finish(admission, outcome, None); + return Err(refusal); + } + } + } + Ok(admission) + } + + /// Records how the request ended: prices its usage, lets hooks settle, and hands the + /// record to the observers. + pub fn finish(&self, admission: Admission, mut outcome: Outcome, response_body: Option) { + let Admission { + request, + mut edit, + settles, + } = admission; + if outcome.cost.is_none() + && let (Some(usage), Some(deployment)) = (&outcome.usage, &outcome.deployment) + { + outcome.cost = self.prices.cost(&deployment.id, usage); + } + let record = Record { + request, + outcome, + annotations: edit.take_annotations(), + response_body: response_body.filter(|_| self.reads(Facts::RESPONSE_BODY)), + }; + for settle in settles { + settle.settle(&record); + } + if !self.observers.is_empty() { + self.observers.send(Arc::new(record)); + } + } + + /// Lets observers finish what they hold. Call once the server has stopped taking requests. + pub async fn close(&self) { + self.observers.close().await; + } +} + +/// A request the hooks let through, with the changes they asked for. +pub struct Admission { + request: Arc, + edit: RequestEdit, + settles: Vec>, +} + +impl Admission { + pub fn request(&self) -> &Arc { + &self.request + } + + pub fn edit(&self) -> &RequestEdit { + &self.edit + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use futures_util::future::BoxFuture; + use http::HeaderMap; + use serde_json::{Value, json}; + + use super::*; + use crate::{ + facts::{DeploymentFacts, Route}, + hook::Targets, + observer::Observer, + usage::Usage, + }; + + fn request() -> Arc { + Arc::new(RequestFacts::new( + "r".into(), + Route::Messages, + None, + None, + "m".into(), + Arc::new(HeaderMap::new()), + false, + None, + Bytes::from_static(b"{}"), + )) + } + + fn outcome() -> Outcome { + Outcome { + deployment: Some(DeploymentFacts { + id: "d".into(), + ..DeploymentFacts::default() + }), + status: 200, + usage: Some(Usage { + input_tokens: 10, + output_tokens: 5, + ..Usage::default() + }), + cost: None, + time_to_first_byte: None, + ended_at: SystemTime::now(), + interrupted: false, + } + } + + type Log = Arc>>; + + struct Step { + name: &'static str, + refuse: bool, + log: Log, + } + + struct Settled(&'static str, Log); + + impl Settle for Settled { + fn settle(self: Box, record: &Record) { + self.1 + .lock() + .unwrap() + .push(format!("settle {} {}", self.0, record.outcome.status)); + } + } + + impl Hook for Step { + fn name(&self) -> &'static str { + self.name + } + + fn reads(&self) -> Facts { + Facts::KEY + } + + fn writes(&self) -> Targets { + Targets::ADMISSION | Targets::ANNOTATIONS + } + + fn run<'a>( + &'a self, + _request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + self.log.lock().unwrap().push(format!("run {}", self.name)); + if self.refuse { + return Err(Refusal::new(429, "rate_limit", "slow down")); + } + edit.annotate("last", Value::from(self.name)); + Ok(Some( + Box::new(Settled(self.name, self.log.clone())) as Box + )) + }) + } + } + + struct Costs(Mutex, Option)>>); + + impl Observer for Costs { + fn name(&self) -> &'static str { + "costs" + } + + fn reads(&self) -> Facts { + Facts::COST | Facts::ANNOTATIONS + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + Box::pin(async move { + self.0 + .lock() + .unwrap() + .push((record.outcome.cost, record.annotations.get("last").cloned())); + }) + } + } + + fn step(name: &'static str, refuse: bool, log: &Log) -> Arc { + Arc::new(Step { + name, + refuse, + log: log.clone(), + }) + } + + #[tokio::test] + async fn hooks_run_in_order_and_settle_with_the_cost() { + let log = Log::default(); + let costs = Arc::new(Costs(Mutex::default())); + let mut prices = Prices::default(); + prices.insert( + "d".into(), + json!({"input_cost_per_token": 1.0, "output_cost_per_token": 2.0}) + .as_object() + .unwrap(), + ); + let pipeline = Pipeline::new( + vec![step("a", false, &log), step("b", false, &log)], + Observers::start(vec![costs.clone()], 8), + Arc::new(prices), + ); + assert!(pipeline.reads(Facts::USAGE | Facts::DEPLOYMENT | Facts::KEY)); + assert!(!pipeline.reads(Facts::RESPONSE_BODY)); + + let admission = pipeline.admit(request()).await.unwrap(); + pipeline.finish(admission, outcome(), None); + pipeline.close().await; + + assert_eq!( + *log.lock().unwrap(), + ["run a", "run b", "settle a 200", "settle b 200"] + ); + assert_eq!( + *costs.0.lock().unwrap(), + [(Some(20.0), Some(Value::from("b")))] + ); + } + + #[tokio::test] + async fn a_refusal_settles_the_hooks_before_it() { + let log = Log::default(); + let pipeline = Pipeline::new( + vec![ + step("a", false, &log), + step("limit", true, &log), + step("c", false, &log), + ], + Observers::start(Vec::new(), 8), + Arc::default(), + ); + let refusal = pipeline.admit(request()).await.err().unwrap(); + assert_eq!(refusal.status, 429); + assert_eq!(*log.lock().unwrap(), ["run a", "run limit", "settle a 429"]); + } +} diff --git a/litellm-rust/crates/gateway-hooks/src/pricing.rs b/litellm-rust/crates/gateway-hooks/src/pricing.rs new file mode 100644 index 00000000000..813676c6f5b --- /dev/null +++ b/litellm-rust/crates/gateway-hooks/src/pricing.rs @@ -0,0 +1,376 @@ +//! Per-deployment prices, built once at startup from the model cost map and any prices the +//! config sets on a deployment, with the proxy's key names: `input_cost_per_token`, its +//! `_above_k_tokens` thresholds and its `_priority` and `_flex` tiers. + +use std::{collections::HashMap, sync::Arc}; + +use litellm_cost::{ + PricingPlan, PromptConvention, Rate, Rates, Request, ServiceTier as CostTier, ThresholdPolicy, + ThresholdRates, TierRates, +}; +use litellm_model_catalog::Catalog; +use serde_json::{Map, Value}; + +use crate::usage::{ServiceTier, Usage}; + +/// Prices by deployment id. +#[derive(Debug, Default)] +pub struct Prices { + plans: HashMap, PricingPlan>, +} + +impl Prices { + /// Adds a deployment's prices from its cost map fields. A deployment whose fields hold no + /// usable prices is left out, and its cost is unknown. + pub fn insert(&mut self, id: Arc, fields: &Map) { + let parsed = ParsedRates::parse(fields); + if parsed.standard == Rates::EMPTY && parsed.thresholds.is_empty() { + return; + } + match parsed.compile() { + Ok(plan) => { + self.plans.insert(id, plan); + } + Err(error) => tracing::warn!(deployment = %id, ?error, "ignoring invalid prices"), + } + } + + pub fn len(&self) -> usize { + self.plans.len() + } + + pub fn is_empty(&self) -> bool { + self.plans.is_empty() + } + + /// 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 + } + } + } +} + +/// 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. +pub fn deployment_fields( + catalog: Option<&Catalog>, + model: &str, + base_model: Option<&str>, + overrides: &[&Map], +) -> 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 + .filter(|name| !name.is_empty()) + .find_map(|name| catalog.lookup(name)) + }) + .map(|found| found.entry.fields().clone()) + .unwrap_or_default(); + for layer in overrides { + for (name, value) in layer.iter() { + if price_key(name).is_some() { + fields.insert(name.clone(), value.clone()); + } + } + } + fields +} + +#[derive(Clone, Copy)] +enum Slot { + Input, + Output, + CacheRead, + CacheWrite, + CacheWrite1h, +} + +/// Longest first, so the hour-long cache write is not read as a threshold of the plain one. +const SLOTS: [(&str, Slot); 5] = [ + ( + "cache_creation_input_token_cost_above_1hr", + Slot::CacheWrite1h, + ), + ("cache_creation_input_token_cost", Slot::CacheWrite), + ("cache_read_input_token_cost", Slot::CacheRead), + ("input_cost_per_token", Slot::Input), + ("output_cost_per_token", Slot::Output), +]; + +struct PriceKey { + slot: Slot, + above_tokens: Option, + tier: CostTier, +} + +fn price_key(name: &str) -> Option { + let (prefix, slot) = SLOTS.iter().find(|(prefix, _)| name.starts_with(prefix))?; + let mut rest = &name[prefix.len()..]; + let mut above_tokens = None; + if let Some(threshold) = rest.strip_prefix("_above_") { + let (count, after) = threshold.split_once("k_tokens")?; + let thousands: f64 = count.parse().ok()?; + above_tokens = Some((thousands * 1000.0) as u64); + rest = after; + } + let tier = match rest { + "" => CostTier::Standard, + "_priority" => CostTier::Priority, + "_flex" => CostTier::Flex, + _ => return None, + }; + Some(PriceKey { + slot: *slot, + above_tokens, + tier, + }) +} + +/// A price, from a number or, as the proxy's `float()` reads config values, a numeric string. +fn rate(value: &Value) -> Rate { + match value { + Value::Number(number) => number.as_f64().map_or(Rate::Null, Rate::Value), + Value::String(text) => text.trim().parse().map_or(Rate::Null, Rate::Value), + _ => Rate::Null, + } +} + +fn set(rates: &mut Rates, slot: Slot, value: Rate) { + *match slot { + Slot::Input => &mut rates.input, + Slot::Output => &mut rates.output, + Slot::CacheRead => &mut rates.cache_read, + Slot::CacheWrite => &mut rates.cache_write, + Slot::CacheWrite1h => &mut rates.cache_write_1h, + } = value; +} + +fn tier_rates(tiers: &mut Vec, tier: CostTier) -> &mut Rates { + let index = match tiers.iter().position(|entry| entry.tier == tier) { + Some(index) => index, + None => { + tiers.push(TierRates { + tier, + rates: Rates::EMPTY, + }); + tiers.len() - 1 + } + }; + &mut tiers[index].rates +} + +struct ParsedThreshold { + above_tokens: u64, + standard: Rates, + tiers: Vec, +} + +struct ParsedRates { + standard: Rates, + tiers: Vec, + thresholds: Vec, +} + +impl ParsedRates { + fn parse(fields: &Map) -> Self { + let mut parsed = Self { + standard: Rates::EMPTY, + tiers: Vec::new(), + thresholds: Vec::new(), + }; + for (name, value) in fields { + let Some(key) = price_key(name) else { continue }; + let (standard, tiers) = match key.above_tokens { + None => (&mut parsed.standard, &mut parsed.tiers), + Some(above_tokens) => { + let index = match parsed + .thresholds + .iter() + .position(|entry| entry.above_tokens == above_tokens) + { + Some(index) => index, + None => { + parsed.thresholds.push(ParsedThreshold { + above_tokens, + standard: Rates::EMPTY, + tiers: Vec::new(), + }); + parsed.thresholds.len() - 1 + } + }; + let entry = &mut parsed.thresholds[index]; + (&mut entry.standard, &mut entry.tiers) + } + }; + let rates = match key.tier { + CostTier::Standard => standard, + tier => tier_rates(tiers, tier), + }; + set(rates, key.slot, rate(value)); + } + // The proxy applies a threshold only when it prices input above it. + parsed.thresholds.retain(|entry| { + matches!(entry.standard.input, Rate::Value(_)) + || entry + .tiers + .iter() + .any(|tier| matches!(tier.rates.input, Rate::Value(_))) + }); + parsed + } + + fn compile(&self) -> Result { + let thresholds: Vec<_> = self + .thresholds + .iter() + .map(|entry| ThresholdRates { + above_prompt_tokens: entry.above_tokens, + standard: entry.standard, + tiers: &entry.tiers, + }) + .collect(); + litellm_cost::compile(&litellm_cost::Pricing { + standard: self.standard, + tiers: &self.tiers, + thresholds: &thresholds, + off_peak: None, + }) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn prices(fields: Value) -> Prices { + let mut prices = Prices::default(); + prices.insert("d".into(), fields.as_object().unwrap()); + prices + } + + fn usage(input: u64, output: u64) -> Usage { + Usage { + input_tokens: input, + output_tokens: output, + ..Usage::default() + } + } + + fn close(left: f64, right: f64) -> bool { + (left - right).abs() < 1e-12 + } + + #[rstest] + fn standard_rates() { + let prices = prices(json!({ + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 0.3e-6, + "cache_creation_input_token_cost": 3.75e-6, + "cache_creation_input_token_cost_above_1hr": 6e-6, + })); + let usage = Usage { + cache_read_tokens: 100, + cache_write_tokens: 30, + cache_write_1h_tokens: 10, + ..usage(1000, 200) + }; + let expected = + 1000.0 * 3e-6 + 200.0 * 15e-6 + 100.0 * 0.3e-6 + 20.0 * 3.75e-6 + 10.0 * 6e-6; + assert!(close(prices.cost("d", &usage).unwrap(), expected)); + } + + #[rstest] + #[case::below(200_000, 1e-6)] + #[case::above(200_001, 2e-6)] + fn thresholds_count_the_whole_prompt(#[case] input: u64, #[case] rate: f64) { + let prices = prices(json!({ + "input_cost_per_token": 1e-6, + "output_cost_per_token": 0.0, + "input_cost_per_token_above_200k_tokens": 2e-6, + })); + assert!(close( + prices.cost("d", &usage(input, 0)).unwrap(), + input as f64 * rate + )); + } + + #[rstest] + fn priority_tier_rates() { + let prices = prices(json!({ + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "input_cost_per_token_priority": 4e-6, + })); + let usage = Usage { + service_tier: ServiceTier::Priority, + ..usage(10, 10) + }; + assert!(close( + prices.cost("d", &usage).unwrap(), + 10.0 * 4e-6 + 10.0 * 2e-6 + )); + } + + #[rstest] + fn unknown_deployments_and_unpriced_entries_have_no_cost() { + let prices = prices(json!({"max_tokens": 10})); + assert!(prices.is_empty()); + assert_eq!(prices.cost("d", &usage(1, 1)), None); + } + + #[rstest] + fn config_prices_override_the_cost_map() { + let catalog = Catalog::parse( + br#"{"claude-x": {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6, "max_tokens": 5}}"#, + Default::default(), + ) + .unwrap(); + let config = json!({"input_cost_per_token": 9e-6, "api_key": "secret"}); + let fields = deployment_fields( + Some(&catalog), + "anthropic/claude-x", + None, + &[config.as_object().unwrap()], + ); + assert_eq!(fields["input_cost_per_token"], json!(9e-6)); + assert_eq!(fields["output_cost_per_token"], json!(2e-6)); + assert!(!fields.contains_key("api_key")); + } +} diff --git a/litellm-rust/crates/gateway-hooks/src/usage.rs b/litellm-rust/crates/gateway-hooks/src/usage.rs new file mode 100644 index 00000000000..c4bcc01b431 --- /dev/null +++ b/litellm-rust/crates/gateway-hooks/src/usage.rs @@ -0,0 +1,452 @@ +//! Token usage, read from a response without holding it: a complete body is read once, and a +//! stream is scanned line by line, keeping only the few lines that carry usage. + +use serde::Deserialize; +use serde_json::Value; + +/// A provider's service tier, which can change the price. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum ServiceTier { + #[default] + Standard, + Priority, + Flex, + Batch, +} + +impl ServiceTier { + fn parse(name: &str) -> Self { + match name { + "priority" => Self::Priority, + "flex" => Self::Flex, + "batch" => Self::Batch, + _ => Self::Standard, + } + } +} + +/// Token counts of one response. Input is split the way prices are: tokens read from the +/// prompt cache, tokens written to it, and the rest. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Usage { + /// Input tokens neither read from nor written to the prompt cache. + pub input_tokens: u64, + pub output_tokens: u64, + pub cache_read_tokens: u64, + /// Every token written to the prompt cache, at either lifetime. + pub cache_write_tokens: u64, + /// The part of `cache_write_tokens` written for an hour, billed at its own rate. + pub cache_write_1h_tokens: u64, + pub web_search_requests: u64, + pub service_tier: ServiceTier, +} + +impl Usage { + /// Prompt tokens as the proxy counts them: cache reads and writes included. + pub fn prompt_tokens(&self) -> u64 { + self.input_tokens + .saturating_add(self.cache_read_tokens) + .saturating_add(self.cache_write_tokens) + } + + pub fn total_tokens(&self) -> u64 { + self.prompt_tokens().saturating_add(self.output_tokens) + } + + /// Usage from an Anthropic Messages `usage` object. + pub fn from_anthropic(value: &Value) -> Option { + let usage = AnthropicUsage::deserialize(value).ok()?; + let mut merged = Self::default(); + usage.merge_into(&mut merged); + Some(merged) + } + + /// Usage from an OpenAI Chat Completions `usage` object, whose prompt count includes cache + /// reads. + pub fn from_openai(value: &Value) -> Option { + OpenAiUsage::deserialize(value).ok().map(OpenAiUsage::usage) + } + + /// Usage from counts in the proxy's convention: `prompt_tokens` includes `cached_tokens` + /// and `cache_creation_tokens`. + pub fn from_prompt_counts( + prompt_tokens: u64, + completion_tokens: u64, + cached_tokens: u64, + cache_creation_tokens: u64, + ) -> Self { + Self { + input_tokens: prompt_tokens + .saturating_sub(cached_tokens) + .saturating_sub(cache_creation_tokens), + output_tokens: completion_tokens, + cache_read_tokens: cached_tokens, + cache_write_tokens: cache_creation_tokens, + ..Self::default() + } + } +} + +#[derive(Deserialize, Default)] +struct AnthropicUsage { + input_tokens: Option, + output_tokens: Option, + cache_creation_input_tokens: Option, + cache_read_input_tokens: Option, + cache_creation: Option, + server_tool_use: Option, + service_tier: Option, +} + +#[derive(Deserialize, Default)] +struct CacheCreation { + ephemeral_1h_input_tokens: Option, +} + +#[derive(Deserialize, Default)] +struct ServerToolUse { + web_search_requests: Option, +} + +impl AnthropicUsage { + /// Applies the counts this object carries over `usage`; a stream's later events repeat or + /// raise the counts of earlier ones, so a present count replaces the earlier one. + fn merge_into(self, usage: &mut Usage) { + if let Some(tokens) = self.input_tokens { + usage.input_tokens = tokens; + } + if let Some(tokens) = self.output_tokens { + usage.output_tokens = tokens; + } + if let Some(tokens) = self.cache_read_input_tokens { + usage.cache_read_tokens = tokens; + } + if let Some(tokens) = self.cache_creation_input_tokens { + usage.cache_write_tokens = tokens; + } + if let Some(tokens) = self + .cache_creation + .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; + } + if let Some(tier) = self.service_tier { + usage.service_tier = ServiceTier::parse(&tier); + } + } +} + +#[derive(Deserialize)] +struct OpenAiUsage { + #[serde(default)] + prompt_tokens: u64, + #[serde(default)] + completion_tokens: u64, + prompt_tokens_details: Option, +} + +#[derive(Deserialize, Default)] +struct OpenAiPromptDetails { + cached_tokens: Option, + cache_creation_tokens: Option, +} + +impl OpenAiUsage { + fn usage(self) -> Usage { + let details = self.prompt_tokens_details.unwrap_or_default(); + Usage::from_prompt_counts( + self.prompt_tokens, + self.completion_tokens, + details.cached_tokens.unwrap_or(0), + details.cache_creation_tokens.unwrap_or(0), + ) + } +} + +/// The wire format of a stream being scanned. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum StreamFormat { + /// Anthropic Messages events: usage in `message_start` and `message_delta`. + Anthropic, + /// OpenAI Chat Completions chunks: usage in the last chunk when the caller asked for it. + OpenAi, +} + +/// The first bytes of a line, enough to tell an Anthropic usage event from content. +const DECIDE_BYTES: usize = 96; +/// The longest line kept. Usage lines are a few hundred bytes; anything longer is content. +const MAX_LINE_BYTES: usize = 64 * 1024; + +#[derive(Debug)] +enum Line { + Reading, + Skipping, +} + +/// Reads usage out of a server-sent-event stream as its chunks pass through. Only the line in +/// progress is buffered, and only while it may carry usage. +#[derive(Debug)] +pub struct StreamUsage { + format: StreamFormat, + line: Vec, + state: Line, + usage: Option, +} + +impl StreamUsage { + pub fn new(format: StreamFormat) -> Self { + Self { + format, + line: Vec::new(), + state: Line::Reading, + usage: None, + } + } + + pub fn feed(&mut self, mut chunk: &[u8]) { + while !chunk.is_empty() { + let end = memchr::memchr(b'\n', chunk); + let (part, rest) = match end { + Some(index) => (&chunk[..index], &chunk[index + 1..]), + None => (chunk, &[][..]), + }; + if let Line::Reading = self.state { + self.take(part); + } + if end.is_some() { + if let Line::Reading = self.state { + self.read_line(); + } + self.line.clear(); + self.state = Line::Reading; + } + chunk = rest; + } + } + + /// Appends part of the line in progress, deciding as early as possible whether to keep it, + /// so a long content line is never copied. + fn take(&mut self, part: &[u8]) { + let undecided = match self.format { + StreamFormat::Anthropic => DECIDE_BYTES.saturating_sub(self.line.len()), + StreamFormat::OpenAi => 0, + }; + let (head, tail) = part.split_at(undecided.min(part.len())); + for piece in [head, tail] { + if piece.is_empty() { + continue; + } + if self.line.len() + piece.len() > MAX_LINE_BYTES { + self.skip(); + return; + } + self.line.extend_from_slice(piece); + if !self.may_carry_usage() { + self.skip(); + return; + } + } + } + + fn skip(&mut self) { + self.line.clear(); + self.state = Line::Skipping; + } + + /// The usage seen so far; `None` when no event carried any. + pub fn usage(&self) -> Option { + self.usage + } + + fn may_carry_usage(&self) -> bool { + let line = &self.line; + if line.len() > MAX_LINE_BYTES { + return false; + } + match self.format { + StreamFormat::Anthropic if line.len() >= DECIDE_BYTES => { + let head = &line[..DECIDE_BYTES]; + contains(head, b"\"message_start\"") || contains(head, b"\"message_delta\"") + } + StreamFormat::Anthropic | StreamFormat::OpenAi => true, + } + } + + fn read_line(&mut self) { + let line = self.line.strip_suffix(b"\r").unwrap_or(&self.line); + let Some(data) = line.strip_prefix(b"data:") else { + return; + }; + let data = data.strip_prefix(b" ").unwrap_or(data); + match self.format { + StreamFormat::Anthropic => { + if !(contains(data, b"\"message_start\"") || contains(data, b"\"message_delta\"")) { + return; + } + let Ok(event) = serde_json::from_slice::(data) else { + return; + }; + let usage = self.usage.get_or_insert_with(Usage::default); + if let Some(start) = event.message.and_then(|message| message.usage) { + start.merge_into(usage); + } + if let Some(delta) = event.usage { + delta.merge_into(usage); + } + } + StreamFormat::OpenAi => { + if !contains(data, b"\"usage\":{") && !contains(data, b"\"usage\": {") { + return; + } + if let Ok(OpenAiChunk { usage: Some(usage) }) = serde_json::from_slice(data) { + self.usage = Some(usage.usage()); + } + } + } + } +} + +#[derive(Deserialize)] +struct AnthropicEvent { + message: Option, + usage: Option, +} + +#[derive(Deserialize)] +struct AnthropicMessage { + usage: Option, +} + +#[derive(Deserialize)] +struct OpenAiChunk { + usage: Option, +} + +fn contains(haystack: &[u8], needle: &[u8]) -> bool { + memchr::memmem::find(haystack, needle).is_some() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + const ANTHROPIC_STREAM: &str = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-x\",\"usage\":{\"input_tokens\":10,\"cache_creation_input_tokens\":6,\"cache_read_input_tokens\":20,\"cache_creation\":{\"ephemeral_5m_input_tokens\":2,\"ephemeral_1h_input_tokens\":4},\"output_tokens\":1,\"service_tier\":\"standard\"}}}\n", + "\n", + "event: content_block_delta\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"a long piece of text that is longer than the decision window, and then some more\"}}\n", + "\n", + "event: message_delta\n", + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":42,\"server_tool_use\":{\"web_search_requests\":2}}}\n", + "\n", + ); + + fn anthropic_usage() -> Usage { + Usage { + input_tokens: 10, + output_tokens: 42, + cache_read_tokens: 20, + cache_write_tokens: 6, + cache_write_1h_tokens: 4, + web_search_requests: 2, + service_tier: ServiceTier::Standard, + } + } + + #[rstest] + #[case::whole(ANTHROPIC_STREAM.len())] + #[case::bytes(1)] + #[case::small(7)] + #[case::mid(50)] + fn anthropic_stream_usage_survives_any_chunking(#[case] size: usize) { + let mut meter = StreamUsage::new(StreamFormat::Anthropic); + for chunk in ANTHROPIC_STREAM.as_bytes().chunks(size) { + meter.feed(chunk); + } + assert_eq!(meter.usage(), Some(anthropic_usage())); + } + + #[rstest] + fn content_lines_are_not_buffered() { + let mut meter = StreamUsage::new(StreamFormat::Anthropic); + let text = "x".repeat(MAX_LINE_BYTES * 2); + meter.feed(b"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\""); + meter.feed(text.as_bytes()); + assert!(meter.line.capacity() < 4 * DECIDE_BYTES); + meter.feed(b"\"}}\n\n"); + assert_eq!(meter.usage(), None); + } + + #[rstest] + fn crlf_lines_are_read() { + let mut meter = StreamUsage::new(StreamFormat::Anthropic); + meter.feed(b"data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":3}}\r\n\r\n"); + assert_eq!(meter.usage().map(|usage| usage.output_tokens), Some(3)); + } + + #[rstest] + 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: [DONE]\n\n", + ); + let mut meter = StreamUsage::new(StreamFormat::OpenAi); + for chunk in stream.as_bytes().chunks(9) { + meter.feed(chunk); + } + assert_eq!( + meter.usage(), + Some(Usage { + input_tokens: 20, + output_tokens: 5, + cache_read_tokens: 10, + ..Usage::default() + }) + ); + } + + #[rstest] + fn anthropic_body_usage() { + let usage = Usage::from_anthropic(&json!({ + "input_tokens": 10, + "output_tokens": 42, + "cache_creation_input_tokens": 6, + "cache_read_input_tokens": 20, + "cache_creation": {"ephemeral_5m_input_tokens": 2, "ephemeral_1h_input_tokens": 4}, + "server_tool_use": {"web_search_requests": 2}, + })); + assert_eq!(usage, Some(anthropic_usage())); + assert_eq!(anthropic_usage().prompt_tokens(), 36); + assert_eq!(anthropic_usage().total_tokens(), 78); + } + + #[rstest] + fn openai_body_usage() { + let usage = Usage::from_openai(&json!({ + "prompt_tokens": 100, + "completion_tokens": 7, + "total_tokens": 107, + "prompt_tokens_details": {"cached_tokens": 60}, + })); + assert_eq!( + usage, + Some(Usage { + input_tokens: 40, + output_tokens: 7, + cache_read_tokens: 60, + ..Usage::default() + }) + ); + } +} diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index d062e4e9223..e589eb75719 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -12,6 +12,7 @@ base64.workspace = true bytes.workspace = true litellm-auth.workspace = true litellm-gateway-auth.workspace = true +litellm-gateway-hooks.workspace = true litellm-inference.workspace = true litellm-inference-transcription.workspace = true litellm-inference-responses.workspace = true @@ -30,6 +31,8 @@ serde_json.workspace = true thiserror.workspace = true tokio.workspace = true futures-util.workspace = true +memchr = "2" +uuid.workspace = true [dev-dependencies] litellm-cache-memory.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 0c2cb543cb0..118eb2cacad 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -5,7 +5,7 @@ use axum::{ Json, body::Bytes, extract::{Path, State}, - http::header::CONTENT_TYPE, + http::{HeaderMap, header::CONTENT_TYPE}, response::{IntoResponse, Response}, }; use litellm_host_http::Sse; @@ -16,23 +16,27 @@ use litellm_inference_chat::{ }; use serde_json::{Map, Value}; +use litellm_gateway_hooks::Route; + use crate::{Deployment, Error, Gateway, JsonBody, request}; pub(crate) async fn create( State(gateway): State>, identity: AuthenticatedRequest, + headers: HeaderMap, body: JsonBody, ) -> Result { - handle(&gateway, &identity, body, None).await + handle(&gateway, &identity, &headers, body, None).await } pub(crate) async fn create_from_model_path( State(gateway): State>, identity: AuthenticatedRequest, Path(model): Path, + headers: HeaderMap, body: JsonBody, ) -> Result { - handle(&gateway, &identity, body, Some(&model)).await + handle(&gateway, &identity, &headers, body, Some(&model)).await } /// The body, naming the model from the request path when it names none itself. @@ -49,6 +53,7 @@ fn with_model(body: Map, model: Option<&str>) -> Map, ) -> Result { @@ -56,10 +61,28 @@ async fn handle( gateway.check_body(&body)?; let group = request::resolve_group(gateway, &body)?; let name = request::model_name(&body)?.to_owned(); - let (body, cache_options) = crate::caching::prepare(identity, body)?; + let tracking = crate::hooks::admit( + gateway, + identity, + Route::ChatCompletions, + headers, + &raw, + &body, + &name, + ) + .await?; + let (mut body, cache_options) = crate::caching::prepare(identity, body)?; + let edit = |body: &mut Map| { + if let Some(tracking) = &tracking { + tracking.edit_body(body); + } + }; + edit(&mut body); let reparse = || -> Result, Error> { let body = with_model(request::object(&raw)?, model); - Ok(crate::caching::prepare(identity, body)?.0) + let mut body = crate::caching::prepare(identity, body)?.0; + edit(&mut body); + Ok(body) }; let route = gateway.chat_completions.clone(); let route = match &gateway.cache { @@ -85,6 +108,9 @@ async fn handle( let route = route.clone(); let cache_headers = cache_headers.clone(); let name = name.as_str(); + if let Some(tracking) = &tracking { + tracking.attempt(name, &attempt); + } async move { let deployment = attempt.deployment; crate::observe::record(name, deployment); @@ -120,8 +146,13 @@ async fn handle( Error::failure, ) .await - .ok_or_else(|| Error::UnknownModel(name.clone()))??; - Ok(cache_headers.apply(request::hold(routed))) + .ok_or_else(|| Error::UnknownModel(name.clone())) + .and_then(|routed| routed) + .map(|routed| cache_headers.apply(request::hold(routed))); + match tracking { + Some(tracking) => tracking.finish(routed).await, + None => routed, + } } /// The `messages` array and the remaining optional params of a chat request. diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 7c58d6fc99e..ce4b5aed0ac 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -3,6 +3,7 @@ use axum::{ Json, response::{IntoResponse, Response}, }; +use litellm_gateway_hooks::Refusal; use litellm_http::transport::Error as TransportError; use litellm_inference::RouteError; use litellm_llms::base_llm::ocr::error::Error as OcrError; @@ -29,6 +30,15 @@ pub enum Error { BodyTooLarge, #[error("{0}")] Internal(String), + /// A hook turned the request away. + #[error("{}", .0.message)] + Refused(Refusal), +} + +impl From for Error { + fn from(refusal: Refusal) -> Self { + Self::Refused(refusal) + } } impl IntoResponse for Error { @@ -50,6 +60,9 @@ impl Error { pub fn status(&self) -> StatusCode { match self { Self::Auth(error) => error.status(), + Self::Refused(refusal) => { + StatusCode::from_u16(refusal.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR) + } Self::Unsupported(_) | Self::Route(RouteError::Unsupported(_)) | Self::Ocr(OcrError::Unsupported(_)) => StatusCode::NOT_IMPLEMENTED, @@ -77,10 +90,20 @@ impl Error { } pub fn openai_response(self) -> Response { - if let Self::Auth(error) = self { - return error.into_response(); + match self { + Self::Auth(error) => error.into_response(), + Self::Refused(refusal) => refused(refusal), + other => (other.status(), Json(other.openai_body())).into_response(), + } + } + + /// The response for a `/v1/messages` request: Anthropic's error envelope, or the proxy's + /// own shape for an error raised before the provider call. + pub fn messages_response(self, request_id: Option<&str>) -> Response { + match self { + Self::Refused(refusal) => refused(refusal), + other => (other.status(), Json(other.body(request_id))).into_response(), } - (self.status(), Json(self.openai_body())).into_response() } /// An OpenAI error event, for a Chat Completions stream that fails after its headers @@ -151,6 +174,20 @@ impl Error { } } +/// A hook's refusal, in the proxy's error shape and with the hook's headers. +fn refused(refusal: Refusal) -> Response { + let status = StatusCode::from_u16(refusal.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + let body = json!({"error": { + "message": refusal.message, + "type": refusal.kind, + "param": "None", + "code": status.as_u16().to_string(), + }}); + let mut response = (status, Json(body)).into_response(); + response.headers_mut().extend(refusal.headers); + response +} + fn error_type(status: StatusCode) -> &'static str { match status.as_u16() { 400 => "invalid_request_error", diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs new file mode 100644 index 00000000000..408810bf095 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -0,0 +1,379 @@ +//! Runs the configured hooks around an inference request and hands its record to the observers +//! once the response body ends. Without hooks or observers none of this runs. + +use std::{ + pin::Pin, + sync::{Arc, Mutex}, + task::{Context, Poll}, + time::{Duration, Instant, SystemTime}, +}; + +use axum::{ + body::{Body, BodyDataStream, Bytes, HttpBody}, + http::{HeaderMap, header::CONTENT_TYPE}, + response::Response, +}; +use futures_util::Stream; +use litellm_gateway_auth::{AuthenticatedRequest, AuthenticationMethod, Permissions}; +use litellm_gateway_hooks::{ + Admission, DeploymentFacts, Facts, KeyFacts, Outcome, Pipeline, RequestFacts, Route, + StreamFormat, StreamUsage, Usage, +}; +use litellm_router::Attempt; +use serde_json::{Map, Value}; + +use crate::{Deployment, Error, Gateway}; + +/// The user the proxy attributes master key requests to, `litellm_proxy_admin_name`. +const PROXY_ADMIN: &str = "default_user_id"; + +/// The hooks and observers every inference request runs through. +pub struct Hooks { + pipeline: Arc, + /// The hash of the master key, which the proxy records as the key of its requests. + master_key_hash: Option>, +} + +impl Hooks { + pub fn new(pipeline: Arc, master_key_hash: Option) -> Self { + Self { + pipeline, + master_key_hash: master_key_hash.map(Into::into), + } + } + + pub fn pipeline(&self) -> &Arc { + &self.pipeline + } +} + +/// One request's admission, and the deployment its last attempt went to. +pub(crate) struct Tracking { + pipeline: Arc, + route: Route, + admission: Admission, + deployment: Mutex>, + started: Instant, +} + +/// Builds the request's facts and runs the hooks. `None` when nothing is configured. +pub(crate) async fn admit( + gateway: &Gateway, + identity: &AuthenticatedRequest, + route: Route, + headers: &HeaderMap, + raw: &Bytes, + body: &Map, + name: &str, +) -> Result, Error> { + let Some(hooks) = gateway + .hooks + .get() + .filter(|hooks| !hooks.pipeline.is_empty()) + else { + return Ok(None); + }; + let pipeline = &hooks.pipeline; + let key = pipeline + .reads(Facts::KEY) + .then(|| key_facts(hooks, identity)) + .flatten(); + let client_headers = match pipeline.reads(Facts::CLIENT_HEADERS) { + true => Arc::new(headers.clone()), + false => Arc::default(), + }; + let request = RequestFacts::new( + uuid::Uuid::new_v4().to_string().into(), + route, + key, + end_user(body).map(Into::into), + name.into(), + client_headers, + body.get("stream").and_then(Value::as_bool).unwrap_or(false), + body.get("max_tokens") + .or_else(|| body.get("max_completion_tokens")) + .and_then(Value::as_u64), + raw.clone(), + ); + let admission = pipeline.admit(Arc::new(request)).await?; + Ok(Some(Tracking { + pipeline: pipeline.clone(), + route, + admission, + deployment: Mutex::new(None), + started: Instant::now(), + })) +} + +fn key_facts(hooks: &Hooks, identity: &AuthenticatedRequest) -> Option> { + let caller = identity.caller(); + let authentication = caller.authentication(); + match (&authentication.method, caller.restrictions()) { + (AuthenticationMethod::VirtualKey, Permissions::VirtualKey(grant)) => { + let profile = &grant.profile; + Some(Arc::new(KeyFacts { + hash: authentication.credential_id.clone(), + alias: profile.alias.clone(), + user_id: profile.user_id.clone(), + team_id: grant.team_id.clone(), + team_alias: profile.team_alias.clone(), + org_id: None, + metadata: profile.metadata.clone(), + team_metadata: profile.team_metadata.clone(), + })) + } + (AuthenticationMethod::MasterKey, _) => Some(Arc::new(KeyFacts { + hash: hooks + .master_key_hash + .as_deref() + .unwrap_or_default() + .to_owned(), + user_id: Some(PROXY_ADMIN.to_owned()), + ..KeyFacts::default() + })), + _ => None, + } +} + +/// The proxy's end user: the body's `user`, else `litellm_metadata.user`, else +/// `metadata.user_id`. +fn end_user(body: &Map) -> Option<&str> { + let nested = |object: &str, field: &str| body.get(object)?.get(field)?.as_str(); + body.get("user") + .and_then(Value::as_str) + .or_else(|| nested("litellm_metadata", "user")) + .or_else(|| nested("metadata", "user_id")) + .filter(|user| !user.is_empty()) +} + +impl Tracking { + /// Applies the hooks' body changes to an attempt's body. + pub(crate) fn edit_body(&self, body: &mut Map) { + self.admission.edit().apply_upstream_body(body); + } + + /// 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) { + return; + } + let deployment: &Deployment = attempt.deployment; + let facts = DeploymentFacts { + id: attempt.id.into(), + model_group: name.into(), + model: deployment.model.clone(), + provider: deployment.custom_llm_provider.clone(), + api_base: deployment.api_base.clone(), + }; + *self.deployment.lock().unwrap_or_else(|e| e.into_inner()) = Some(facts); + } + + /// Adds the hooks' response headers and follows the body to its end; an error is recorded + /// at once. + pub(crate) async fn finish(self, result: Result) -> Result { + let Self { + pipeline, + route, + admission, + deployment, + started, + } = self; + let deployment = deployment.into_inner().unwrap_or_else(|e| e.into_inner()); + let mut response = match result { + Ok(response) => response, + Err(error) => { + let outcome = Outcome { + deployment, + status: error.status().as_u16(), + usage: None, + cost: None, + time_to_first_byte: None, + ended_at: SystemTime::now(), + interrupted: false, + }; + pipeline.finish(admission, outcome, None); + return Err(error); + } + }; + if response.status().is_success() { + let headers = response.headers_mut(); + for (name, value) in admission.edit().response_headers() { + headers.insert(name.clone(), value.clone()); + } + } + let streaming = response + .headers() + .get(CONTENT_TYPE) + .is_some_and(|value| value.as_bytes().starts_with(b"text/event-stream")); + let meter = match (pipeline.reads(Facts::USAGE), streaming, route) { + (false, _, _) => Meter::Off, + (true, true, Route::Messages) => { + Meter::Stream(StreamUsage::new(StreamFormat::Anthropic)) + } + (true, true, Route::ChatCompletions) => { + Meter::Stream(StreamUsage::new(StreamFormat::OpenAi)) + } + (true, false, _) => Meter::Body(Vec::new()), + }; + let mut finish = Finish { + body: pipeline.reads(Facts::RESPONSE_BODY).then(Vec::new), + pipeline, + route, + admission: Some(admission), + deployment, + status: response.status().as_u16(), + meter, + started, + first_byte: None, + ended: false, + }; + if response.body().size_hint().exact().is_some() { + // A complete body is already in memory: read it now and keep its length header. + let (parts, body) = response.into_parts(); + let bytes = axum::body::to_bytes(body, usize::MAX) + .await + .map_err(|error| Error::Internal(error.to_string()))?; + finish.feed(&bytes); + finish.ended = true; + drop(finish); + return Ok(Response::from_parts(parts, Body::from(bytes))); + } + Ok(response.map(|body| { + Body::from_stream(Tapped { + inner: body.into_data_stream(), + finish, + }) + })) + } +} + +enum Meter { + Off, + Stream(StreamUsage), + /// A JSON body, read for its `usage` once complete. + Body(Vec), +} + +/// Records the request when the response body ends or is dropped. +struct Finish { + pipeline: Arc, + route: Route, + admission: Option, + deployment: Option, + status: u16, + meter: Meter, + body: Option>, + started: Instant, + first_byte: Option, + ended: bool, +} + +impl Finish { + fn feed(&mut self, chunk: &Bytes) { + self.first_byte + .get_or_insert_with(|| self.started.elapsed()); + match &mut self.meter { + Meter::Off => {} + Meter::Stream(meter) => meter.feed(chunk), + Meter::Body(chunks) => chunks.push(chunk.clone()), + } + if let Some(body) = &mut self.body { + body.push(chunk.clone()); + } + } + + fn usage(&self) -> Option { + match &self.meter { + Meter::Off => None, + Meter::Stream(meter) => meter.usage(), + Meter::Body(chunks) => body_usage(self.route, &chunks.concat()), + } + } +} + +impl Drop for Finish { + fn drop(&mut self) { + let Some(admission) = self.admission.take() else { + return; + }; + let outcome = Outcome { + deployment: self.deployment.take(), + status: self.status, + usage: self.usage(), + cost: None, + time_to_first_byte: self.first_byte, + ended_at: SystemTime::now(), + interrupted: !self.ended, + }; + let body = self.body.take().map(|chunks| Bytes::from(chunks.concat())); + self.pipeline.finish(admission, outcome, body); + } +} + +/// 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 { + const KEY: &[u8] = b"\"usage\":"; + let start = memchr::memmem::rfind(body, KEY)? + KEY.len(); + let value: Value = serde_json::Deserializer::from_slice(&body[start..]) + .into_iter() + .next()? + .ok()?; + match route { + Route::Messages => Usage::from_anthropic(&value), + Route::ChatCompletions => Usage::from_openai(&value), + } +} + +struct Tapped { + inner: BodyDataStream, + finish: Finish, +} + +impl Stream for Tapped { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let polled = Pin::new(&mut this.inner).poll_next(cx); + match &polled { + Poll::Ready(Some(Ok(chunk))) => this.finish.feed(chunk), + Poll::Ready(None) => this.finish.ended = true, + Poll::Ready(Some(Err(_))) | Poll::Pending => {} + } + polled + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::user(json!({"user": "u1", "metadata": {"user_id": "u2"}}), Some("u1"))] + #[case::metadata(json!({"metadata": {"user_id": "u2"}}), Some("u2"))] + #[case::litellm_metadata(json!({"litellm_metadata": {"user": "u3"}}), Some("u3"))] + #[case::none(json!({"user": ""}), None)] + fn end_user_follows_the_proxy(#[case] body: Value, #[case] expected: Option<&str>) { + assert_eq!(end_user(body.as_object().unwrap()), expected); + } + + #[rstest] + fn body_usage_reads_the_top_level_object() { + let body = json!({ + "id": "msg", + "content": [{"type": "text", "text": "\"usage\": {\"input_tokens\": 99}"}], + "usage": {"input_tokens": 3, "output_tokens": 4}, + }); + let usage = body_usage(Route::Messages, &serde_json::to_vec(&body).unwrap()).unwrap(); + assert_eq!((usage.input_tokens, usage.output_tokens), (3, 4)); + + let chat = json!({"choices": [], "usage": {"prompt_tokens": 5, "completion_tokens": 6}}); + let usage = + body_usage(Route::ChatCompletions, &serde_json::to_vec(&chat).unwrap()).unwrap(); + assert_eq!((usage.prompt_tokens(), usage.output_tokens), (5, 6)); + } +} diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 6839dadc8a7..56b62d83d77 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -8,13 +8,14 @@ mod caching; mod chat_completions; mod client_side; mod error; +mod hooks; pub mod messages; pub mod observe; mod ocr; mod request; mod responses; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use axum::{Router, routing::post}; use litellm_http::{ClientVariant, HttpClientConfig, media::UrlPolicy}; @@ -28,11 +29,13 @@ use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; use litellm_secrets::source::SecretSource; pub use error::Error; +pub use hooks::Hooks; pub use litellm_router::{Deployment, Error as RoutingError, Router as ModelRouter}; pub use request::{JsonBody, JsonObject, RequestId}; pub struct Gateway { cache: Option>, + hooks: OnceLock, allow_client_side_credentials: bool, pub audio_transcription: AudioTranscriptionRoute, pub chat_completions: ChatCompletionsRoute, @@ -53,6 +56,17 @@ impl Gateway { } } + /// Runs `/v1/messages` and chat completions requests through `hooks`. Only the first call + /// takes effect. + pub fn install_hooks(&self, hooks: Hooks) { + let _ = self.hooks.set(hooks); + } + + /// The hooks requests run through, once installed. + pub fn hooks(&self) -> Option<&Hooks> { + self.hooks.get() + } + /// `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 { @@ -83,6 +97,7 @@ impl Gateway { let auth = resources.auth.clone(); Ok(Self { cache: None, + hooks: OnceLock::new(), 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 976cbaf27ac..03f1f57b905 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -15,6 +15,8 @@ use litellm_inference_messages::{MessagesCall, messages_body, route::Messages}; use litellm_llms_types::headers::{ProviderSpecificHeader, ProviderSpecificHeaders}; use serde_json::{Map, Value}; +use litellm_gateway_hooks::Route; + use crate::{Deployment, Error, Gateway, JsonBody, RequestId, request}; /// Client headers Python forwards to Anthropic-speaking providers on every call. @@ -37,7 +39,7 @@ pub async fn create( Ok(body) => handle(&gateway, &identity, &headers, body).await, Err(error) => Err(error), }; - result.map_err(|error| (error.status(), Json(error.body(request_id.as_deref())))) + result.map_err(|error| error.messages_response(request_id.as_deref())) } async fn handle( @@ -49,9 +51,27 @@ async fn handle( gateway.check_body(&object)?; let group = request::resolve_group(gateway, &object)?; let name = request::model_name(&object)?.to_owned(); - let (body, cache_options) = crate::caching::prepare(identity, object)?; + let tracking = crate::hooks::admit( + gateway, + identity, + Route::Messages, + headers, + &raw, + &object, + &name, + ) + .await?; + let (mut body, cache_options) = crate::caching::prepare(identity, object)?; + let edit = |body: &mut Map| { + if let Some(tracking) = &tracking { + tracking.edit_body(body); + } + }; + edit(&mut body); let reparse = || -> Result, Error> { - Ok(crate::caching::prepare(identity, request::object(&raw)?)?.0) + let mut body = crate::caching::prepare(identity, request::object(&raw)?)?.0; + edit(&mut body); + Ok(body) }; let route = gateway.messages.clone(); let route = match &gateway.cache { @@ -77,6 +97,9 @@ async fn handle( let route = route.clone(); let cache_headers = cache_headers.clone(); let name = name.as_str(); + if let Some(tracking) = &tracking { + tracking.attempt(name, &attempt); + } async move { crate::observe::record(name, attempt.deployment); request::authorize(identity, name, attempt.deployment).await?; @@ -91,8 +114,13 @@ async fn handle( Error::failure, ) .await - .ok_or_else(|| Error::UnknownModel(name.clone()))??; - Ok(cache_headers.apply(request::hold(routed))) + .ok_or_else(|| Error::UnknownModel(name.clone())) + .and_then(|routed| routed) + .map(|routed| cache_headers.apply(request::hold(routed))); + match tracking { + Some(tracking) => tracking.finish(routed).await, + None => routed, + } } fn project( diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index d74866b623e..7d5ccf6069d 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -16,6 +16,8 @@ tokio-util = "0.7" thiserror.workspace = true futures-util.workspace = true litellm-gateway-auth.workspace = true +litellm-gateway-hooks.workspace = true +litellm-model-catalog.workspace = true litellm-gateway-ui.workspace = true litellm-auth-types.workspace = true litellm-config.workspace = true diff --git a/litellm-rust/crates/gateway/src/hooks.rs b/litellm-rust/crates/gateway/src/hooks.rs new file mode 100644 index 00000000000..02e9180c771 --- /dev/null +++ b/litellm-rust/crates/gateway/src/hooks.rs @@ -0,0 +1,217 @@ +//! Builds the hooks and observers inference requests run through, and the per-deployment prices +//! observers that read cost need. + +use std::sync::Arc; + +use litellm_config::Config; +use litellm_gateway_hooks::{ + DEFAULT_QUEUE, Facts, Hook, Observer, Observers, Pipeline, Prices, deployment_fields, +}; +use litellm_gateway_inference::{Gateway, Hooks}; +use litellm_model_catalog::{Catalog, Provenance}; +use litellm_secrets::source::SecretSource; +use serde_json::{Map, Value}; + +use crate::Error; + +/// A local copy of the model cost map, read instead of fetching it. +const COST_MAP_PATH: &str = "LITELLM_MODEL_COST_MAP_PATH"; +/// Where to fetch the model cost map from, as the proxy's `LITELLM_MODEL_COST_MAP_URL`. +const COST_MAP_URL: &str = "LITELLM_MODEL_COST_MAP_URL"; +const DEFAULT_COST_MAP_URL: &str = + "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"; + +/// Installs `hooks` and `observers` on `inference`, in the order given. Nothing is installed when +/// both are empty, so requests skip the pipeline altogether. +pub async fn install( + inference: &Gateway, + config: &Config, + hooks: Vec>, + observers: Vec>, +) -> Result<(), Error> { + if hooks.is_empty() && observers.is_empty() { + return Ok(()); + } + let needs_cost = observers + .iter() + .any(|observer| observer.reads().contains(Facts::COST)) + || hooks.iter().any(|hook| hook.reads().contains(Facts::COST)); + let prices = match needs_cost { + true => { + let catalog = load_catalog(inference).await; + prices(config, inference, catalog.as_ref()) + } + false => Prices::default(), + }; + let names: Vec<_> = hooks + .iter() + .map(|hook| hook.name()) + .chain(observers.iter().map(|observer| observer.name())) + .collect(); + let pipeline = Pipeline::new( + hooks, + Observers::start(observers, DEFAULT_QUEUE), + Arc::new(prices), + ); + let master_key_hash = master_key_hash(config, inference.secrets.as_ref()).await?; + tracing::info!( + ?names, + priced_deployments = pipeline.prices().len(), + "hooks installed" + ); + inference.install_hooks(Hooks::new(Arc::new(pipeline), master_key_hash)); + Ok(()) +} + +/// Lets observers write out what they hold; call after the server stops taking requests. +pub async fn close(inference: &Gateway) { + if let Some(hooks) = inference.hooks() { + hooks.pipeline().close().await; + } +} + +/// The sha256 of the master key, which the proxy records as the key of master key requests. +async fn master_key_hash( + config: &Config, + secrets: &dyn SecretSource, +) -> Result, Error> { + let Some(configured) = &config.general_settings.master_key else { + return Ok(None); + }; + let key = match configured.expose().strip_prefix("os.environ/") { + Some(name) => secrets.get_secret_str(name).await?, + None => Some(configured.clone()), + }; + Ok(key.map(|key| litellm_gateway_auth::hash_token(key.expose()))) +} + +/// The model cost map, from a local copy or the proxy's URL. Without it only deployments that +/// set their own prices are priced. +async fn load_catalog(inference: &Gateway) -> Option { + let (body, source) = match std::env::var(COST_MAP_PATH) { + Ok(path) => match std::fs::read(&path) { + Ok(body) => (body, path), + Err(error) => { + tracing::warn!(%path, %error, "cannot read the model cost map"); + return None; + } + }, + Err(_) => { + let url = std::env::var(COST_MAP_URL).unwrap_or_else(|_| DEFAULT_COST_MAP_URL.into()); + match fetch(inference, &url).await { + Ok(body) => (body, url), + Err(error) => { + tracing::warn!(%url, %error, "cannot fetch the model cost map"); + return None; + } + } + } + }; + let provenance = Provenance { + source: Some(source), + ..Provenance::default() + }; + match Catalog::parse(&body, provenance) { + Ok(catalog) => Some(catalog), + Err(error) => { + tracing::warn!(%error, "cannot parse the model cost map"); + None + } + } +} + +async fn fetch(inference: &Gateway, url: &str) -> Result, Box> { + let client = inference + .resources + .pool + .client(&inference.http, litellm_http::ClientVariant::Provider)?; + let response = client + .get(url) + .timeout(std::time::Duration::from_secs(10)) + .send() + .await? + .error_for_status()?; + Ok(response.bytes().await?.to_vec()) +} + +/// Every deployment's prices: the cost map entry for its model, overlaid by the prices its +/// `model_info` and `litellm_params` set, keyed by the router's deployment ids. +fn prices(config: &Config, inference: &Gateway, catalog: Option<&Catalog>) -> Prices { + let mut prices = Prices::default(); + for (model, id) in config + .model_list + .iter() + .zip(inference.models.deployment_ids()) + { + let model_info = json_object(model.model_info.iter()); + let params = json_object( + model + .litellm_params + .extra + .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], + ); + prices.insert(id.into(), &fields); + } + prices +} + +fn json_object<'a, V: serde::Serialize + 'a>( + entries: impl Iterator, +) -> Map { + entries + .filter_map(|(name, value)| Some((name.clone(), serde_json::to_value(value).ok()?))) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[tokio::test] + async fn deployment_prices_come_from_config_and_the_cost_map() { + let config = Config::from_yaml( + r#" +model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_key: sk + model_info: + id: priced-by-map + - model_name: custom + litellm_params: + model: openai/my-model + api_key: sk + input_cost_per_token: 0.5 + output_cost_per_token: 1 +"#, + ) + .unwrap(); + let inference = crate::build_inference(&config).unwrap(); + let catalog = Catalog::parse( + br#"{"claude-x": {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}}"#, + Provenance::default(), + ) + .unwrap(); + let prices = prices(&config, &inference, Some(&catalog)); + let usage = litellm_gateway_hooks::Usage { + input_tokens: 10, + output_tokens: 1, + ..Default::default() + }; + assert_eq!(prices.cost("priced-by-map", &usage), Some(0.012)); + let custom = inference.models.deployment_ids().nth(1).unwrap(); + assert_eq!(prices.cost(custom, &usage), Some(6.0)); + } +} diff --git a/litellm-rust/crates/gateway/src/key_status.rs b/litellm-rust/crates/gateway/src/key_status.rs index d68c2e6f0b3..f4f54e9ac3d 100644 --- a/litellm-rust/crates/gateway/src/key_status.rs +++ b/litellm-rust/crates/gateway/src/key_status.rs @@ -1,6 +1,8 @@ use litellm_db::{Restrictions, VerificationToken}; +use std::sync::Arc; + use litellm_gateway_auth::keys::{ - AllowedRoutes, KeyModels, KeyStatus, ModelAccess, Owner, Refusal, Restriction, + AllowedRoutes, KeyModels, KeyProfile, KeyStatus, ModelAccess, Owner, Refusal, Restriction, }; use serde_json::Value; @@ -31,14 +33,33 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K refusal, }; } + let object = |value: Value| match value { + Value::Object(fields) => Arc::new(fields), + _ => Arc::default(), + }; + let (team_models, team_alias, team_metadata) = match row.team { + Some(team) => ( + Some(ModelAccess::new(team.models)), + team.alias, + object(team.metadata), + ), + None => (None, None, Arc::default()), + }; KeyStatus::Active { expires_at: row.expires, models: KeyModels { key: ModelAccess::new(row.models), - team: row.team.map(|team| ModelAccess::new(team.models)), + team: team_models, }, allowed_routes: AllowedRoutes::new(row.allowed_routes), team_id: row.team_id, + profile: Arc::new(KeyProfile { + alias: row.key_alias, + user_id: row.user_id, + team_alias, + metadata: object(row.metadata), + team_metadata, + }), } } @@ -267,6 +288,7 @@ mod tests { fn team(blocked: bool, max_budget: Option, spend: f64) -> Option { Some(Team { + alias: None, blocked, models: strings(&["gpt-b"]), spend, @@ -285,6 +307,7 @@ mod tests { }, allowed_routes: AllowedRoutes::new(strings(&["llm_api_routes"])), team_id: team.then(|| "team-1".to_owned()), + profile: Default::default(), } } @@ -662,13 +685,32 @@ mod tests { ..row }; let personal = team_id.is_none() && user_id.is_some(); - assert_eq!(status(row.clone(), false), active(team_id.is_some())); + let expected = match active(team_id.is_some()) { + KeyStatus::Active { + expires_at, + models, + allowed_routes, + team_id, + .. + } => KeyStatus::Active { + expires_at, + models, + allowed_routes, + team_id, + profile: Arc::new(KeyProfile { + user_id: user_id.map(str::to_owned), + ..KeyProfile::default() + }), + }, + other => other, + }; + assert_eq!(status(row.clone(), false), expected); assert_eq!( status(row, true), if personal { unsupported(Restriction::PersonalKeyTeamModels, Owner::Proxy) } else { - active(team_id.is_some()) + expected } ); } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 952f7cf9e2c..55fd19bb941 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -3,6 +3,7 @@ pub mod coordination; pub mod drain; mod error; pub mod health; +pub mod hooks; mod key_status; mod keys; pub mod metrics; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 31c4a741a84..ed4f293b4de 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -108,6 +108,8 @@ async fn main() -> Result<(), Box> { ))); coordination.start(inference.clone(), keys.clone(), shutdown.clone()); } + // Built-in hooks and observers are added here as the gateway gains them. + litellm_gateway::hooks::install(&inference, &config, Vec::new(), Vec::new()).await?; let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; let address = listener.local_addr()?; // Without TCP_NODELAY, Nagle holds small SSE chunks until the client's delayed ACK (~40 ms per stream). @@ -124,7 +126,7 @@ async fn main() -> Result<(), Box> { let metrics = litellm_gateway::metrics::Metrics::new(&config, draining.flag()); let metrics_auth = litellm_gateway::metrics_auth(&config, inference.secrets.clone()); let app = litellm_gateway::with_ops( - litellm_gateway::router_with_key_cache(inference, &config, ui, mcp_router, keys), + litellm_gateway::router_with_key_cache(inference.clone(), &config, ui, mcp_router, keys), metrics, draining.clone(), metrics_auth, @@ -139,6 +141,7 @@ async fn main() -> Result<(), Box> { ) .await; shutdown.cancel(); + litellm_gateway::hooks::close(&inference).await; if let Some(mcp) = mcp { mcp.close().await; } diff --git a/litellm-rust/crates/gateway/tests/hooks.rs b/litellm-rust/crates/gateway/tests/hooks.rs new file mode 100644 index 00000000000..0988553a3b6 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/hooks.rs @@ -0,0 +1,395 @@ +//! Hooks and observers around real requests: a hook that can refuse, add a response header and +//! change the upstream body, and an observer that records what each request cost. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +use axum::{ + Router, + body::{Body, Bytes, to_bytes}, + http::{HeaderName, HeaderValue, Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use futures_util::future::BoxFuture; +use litellm_config::Config; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, +}; +use litellm_gateway_hooks::{ + Facts, Hook, Observer, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; +const MESSAGE_STREAM: &str = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-x\",\"usage\":{\"input_tokens\":10,\"cache_read_input_tokens\":20,\"output_tokens\":1}}}\n\n", + "event: content_block_delta\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n", + "event: message_delta\n", + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":42}}\n\n", + "event: message_stop\n", + "data: {\"type\":\"message_stop\"}\n\n", +); +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}}"#; + +type Bodies = Arc>>; + +/// Answers Anthropic and OpenAI calls, keeping each request body. +async fn upstream(bodies: Bodies) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(move |request: Request| { + let bodies = bodies.clone(); + async move { + let path = request.uri().path().to_owned(); + let body: Value = + serde_json::from_slice(&to_bytes(request.into_body(), 1 << 20).await.unwrap()) + .unwrap(); + let stream = body["stream"].as_bool().unwrap_or(false); + bodies.lock().unwrap().push(body); + let (content_type, payload) = match (path.ends_with("/chat/completions"), stream) { + (true, _) => ("application/json", COMPLETION), + (false, true) => ("text/event-stream", MESSAGE_STREAM), + (false, false) => ("application/json", MESSAGE), + }; + ([(CONTENT_TYPE, content_type)], payload).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 Released(Arc>>); + +impl Settle for Released { + fn settle(self: Box, record: &Record) { + self.0.lock().unwrap().push(record.outcome.status); + } +} + +/// Refuses requests carrying `x-refuse`; otherwise tags the upstream body and the response. +struct GateHook(Arc>>); + +impl Hook for GateHook { + fn name(&self) -> &'static str { + "gate" + } + + fn reads(&self) -> Facts { + Facts::CLIENT_HEADERS + } + + fn writes(&self) -> Targets { + Targets::ADMISSION | Targets::UPSTREAM_BODY | Targets::RESPONSE_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + if request.client_headers.contains_key("x-refuse") { + return Err( + Refusal::new(429, "rate_limit_exceeded", "slow down").with_header( + HeaderName::from_static("retry-after"), + HeaderValue::from_static("7"), + ), + ); + } + edit.set_upstream_body_field("metadata", json!({"user_id": "hooked"})); + edit.add_response_header( + HeaderName::from_static("x-hooked"), + HeaderValue::from_static("1"), + ); + Ok(Some(Box::new(Released(self.0.clone())) as Box)) + }) + } +} + +#[derive(Debug, PartialEq)] +struct Seen { + status: u16, + key_alias: Option, + team_id: Option, + deployment: Option, + usage: Option<(u64, u64, u64)>, + cost: Option, + interrupted: bool, +} + +struct Collect(Mutex>); + +impl Observer for Collect { + fn name(&self) -> &'static str { + "collect" + } + + fn reads(&self) -> Facts { + Facts::KEY | Facts::DEPLOYMENT | Facts::USAGE | Facts::COST + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + Box::pin(async move { + self.0.lock().unwrap().push(Seen { + status: record.outcome.status, + key_alias: record.key().and_then(|key| key.alias.clone()), + team_id: record.key().and_then(|key| key.team_id.clone()), + deployment: record + .outcome + .deployment + .as_ref() + .map(|deployment| deployment.id.to_string()), + usage: record.outcome.usage.map(|usage| { + ( + usage.input_tokens, + usage.cache_read_tokens, + usage.output_tokens, + ) + }), + cost: record.outcome.cost, + interrupted: record.outcome.interrupted, + }); + }) + } +} + +struct Harness { + app: Router, + inference: Arc, + bodies: Bodies, + settled: Arc>>, + seen: Arc, +} + +async fn harness() -> Harness { + let bodies = Bodies::default(); + let upstream = upstream(bodies.clone()).await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream + input_cost_per_token: 0.001 + output_cost_per_token: 0.002 + cache_read_input_token_cost: 0.0001 + model_info: + id: claude-1 + - model_name: gpt + litellm_params: + model: openai/gpt-x + api_base: {upstream}/v1 + api_key: sk-upstream + input_cost_per_token: 1 + output_cost_per_token: 2 + model_info: + id: gpt-1 +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let settled = Arc::new(Mutex::new(Vec::new())); + let seen = Arc::new(Collect(Mutex::new(Vec::new()))); + litellm_gateway::hooks::install( + &inference, + &config, + vec![Arc::new(GateHook(settled.clone()))], + vec![seen.clone()], + ) + .await + .unwrap(); + let app = + litellm_gateway::router_with_keys(inference.clone(), &config, None, None, Some(keys())); + Harness { + app, + inference, + bodies, + settled, + seen, + } +} + +async fn send( + app: &Router, + path: &str, + key: &str, + body: Value, + refuse: bool, +) -> (u16, axum::http::HeaderMap, Bytes) { + let mut request = Request::post(path) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {key}")); + if refuse { + request = request.header("x-refuse", "1"); + } + let response = app + .clone() + .oneshot(request.body(Body::from(body.to_string())).unwrap()) + .await + .unwrap(); + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + (status, headers, body) +} + +fn messages(stream: bool) -> Value { + json!({ + "model": "claude", + "max_tokens": 8, + "stream": stream, + "messages": [{"role": "user", "content": "hi"}], + }) +} + +fn close(left: f64, right: f64) -> bool { + (left - right).abs() < 1e-9 +} + +#[rstest] +#[case::unary(false, (3, 0, 4), 3.0 * 0.001 + 4.0 * 0.002)] +#[case::stream(true, (10, 20, 42), 10.0 * 0.001 + 20.0 * 0.0001 + 42.0 * 0.002)] +#[tokio::test] +async fn messages_are_hooked_and_priced( + #[case] stream: bool, + #[case] usage: (u64, u64, u64), + #[case] cost: f64, +) { + let harness = harness().await; + let (status, headers, body) = send( + &harness.app, + "/v1/messages", + "sk-team", + messages(stream), + false, + ) + .await; + assert_eq!(status, 200, "{body:?}"); + assert_eq!(headers["x-hooked"], "1"); + if !stream { + assert_eq!(headers["content-length"], body.len().to_string()); + } + assert_eq!( + harness.bodies.lock().unwrap()[0]["metadata"], + json!({"user_id": "hooked"}) + ); + + litellm_gateway::hooks::close(&harness.inference).await; + assert_eq!(*harness.settled.lock().unwrap(), [200]); + let seen = harness.seen.0.lock().unwrap(); + let [seen] = seen.as_slice() else { + panic!("{seen:?}") + }; + assert_eq!(seen.status, 200); + assert_eq!(seen.key_alias.as_deref(), Some("team-key")); + assert_eq!(seen.team_id.as_deref(), Some("team-1")); + assert_eq!(seen.deployment.as_deref(), Some("claude-1")); + assert_eq!(seen.usage, Some(usage)); + assert!(close(seen.cost.unwrap(), cost), "{seen:?}"); + assert!(!seen.interrupted); +} + +#[rstest] +#[tokio::test] +async fn chat_completions_are_priced() { + let harness = harness().await; + let request = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]}); + let (status, _, body) = send( + &harness.app, + "/v1/chat/completions", + "sk-master", + request, + false, + ) + .await; + assert_eq!(status, 200, "{body:?}"); + + litellm_gateway::hooks::close(&harness.inference).await; + let seen = harness.seen.0.lock().unwrap(); + assert_eq!(seen[0].deployment.as_deref(), Some("gpt-1")); + assert_eq!(seen[0].usage, Some((5, 0, 6))); + assert_eq!(seen[0].cost, Some(17.0)); + assert_eq!(seen[0].key_alias, None); +} + +#[rstest] +#[case::messages("/v1/messages")] +#[case::chat("/v1/chat/completions")] +#[tokio::test] +async fn a_refusal_stops_the_request_with_its_headers(#[case] path: &str) { + let harness = harness().await; + let (status, headers, body) = send(&harness.app, path, "sk-team", messages(false), true).await; + assert_eq!(status, 429); + assert_eq!(headers["retry-after"], "7"); + let body: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["error"]["message"], "slow down"); + assert_eq!(body["error"]["type"], "rate_limit_exceeded"); + assert!(harness.bodies.lock().unwrap().is_empty()); + + litellm_gateway::hooks::close(&harness.inference).await; + let seen = harness.seen.0.lock().unwrap(); + assert_eq!(seen[0].status, 429); + assert_eq!(seen[0].deployment, None); +} + +#[rstest] +#[tokio::test] +async fn a_stream_dropped_early_is_recorded_as_interrupted() { + let harness = harness().await; + let response = harness + .app + .clone() + .oneshot( + Request::post("/v1/messages") + .header("content-type", "application/json") + .header("authorization", "Bearer sk-team") + .body(Body::from(messages(true).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), 200); + drop(response); + + litellm_gateway::hooks::close(&harness.inference).await; + let seen = harness.seen.0.lock().unwrap(); + assert!(seen[0].interrupted); + assert_eq!(*harness.settled.lock().unwrap(), [200]); +} diff --git a/litellm-rust/crates/gateway/tests/shared_state.rs b/litellm-rust/crates/gateway/tests/shared_state.rs index 76aa77de3c0..cac3952b2a5 100644 --- a/litellm-rust/crates/gateway/tests/shared_state.rs +++ b/litellm-rust/crates/gateway/tests/shared_state.rs @@ -191,6 +191,7 @@ fn on_team() -> KeyStatus { models, allowed_routes, team_id: Some("team-1".into()), + profile: Default::default(), }, status => status, } diff --git a/litellm-rust/crates/gateway/tests/virtual_keys.rs b/litellm-rust/crates/gateway/tests/virtual_keys.rs index ea0ef00c6da..4c6d0c41957 100644 --- a/litellm-rust/crates/gateway/tests/virtual_keys.rs +++ b/litellm-rust/crates/gateway/tests/virtual_keys.rs @@ -53,6 +53,7 @@ fn on_team(team_id: &str) -> KeyStatus { }, allowed_routes: AllowedRoutes::default(), team_id: Some(team_id.to_owned()), + profile: Default::default(), } } @@ -77,6 +78,7 @@ fn keys() -> Arc { }, allowed_routes: AllowedRoutes::default(), team_id: Some("team-1".into()), + profile: Default::default(), }, ), ( @@ -86,6 +88,7 @@ fn keys() -> Arc { models: KeyModels::default(), allowed_routes: AllowedRoutes::new(vec!["openai_routes".to_owned()]), team_id: None, + profile: Default::default(), }, ), ( @@ -95,6 +98,7 @@ fn keys() -> Arc { models: KeyModels::default(), allowed_routes: AllowedRoutes::new(vec!["llm_api_routes".to_owned()]), team_id: None, + profile: Default::default(), }, ), // The key store passes a key whose stored spend is under its budget. diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index ebee2c64140..dfca07bf36b 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -283,13 +283,25 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub litellm_provider: Option, /// Maximum prompt/context tokens the model accepts. - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "whole_number", + skip_serializing_if = "Option::is_none" + )] pub max_input_tokens: Option, /// Maximum tokens the model can generate in one response. - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "whole_number", + skip_serializing_if = "Option::is_none" + )] pub max_output_tokens: Option, /// Legacy field: max output tokens if the provider specifies it, else max input tokens. - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "whole_number", + skip_serializing_if = "Option::is_none" + )] pub max_tokens: Option, /// Free-form notes about the entry (e.g. pricing derivation). #[serde(skip_serializing_if = "Option::is_none")] @@ -426,10 +438,18 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub output_dbu_cost_per_token: Option, /// Embedding dimension for embedding models. - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "whole_number", + skip_serializing_if = "Option::is_none" + )] pub output_vector_size: Option, /// Smallest prefix the provider will actually cache; absent means the provider default applies. - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "whole_number", + skip_serializing_if = "Option::is_none" + )] pub prompt_cache_min_tokens: Option, /// Provider-internal routing hints (e.g. bedrock_invocation_schema). #[serde(skip_serializing_if = "Option::is_none")] @@ -447,7 +467,11 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub regional_processing_uplift_multiplier_us: Option, /// Provider default requests-per-minute limit. - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "whole_number", + skip_serializing_if = "Option::is_none" + )] pub rpm: Option, #[serde(skip_serializing_if = "Option::is_none")] pub rules: Option>, @@ -566,7 +590,11 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub tiered_pricing: Option>, /// Provider default tokens-per-minute limit. - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "whole_number", + skip_serializing_if = "Option::is_none" + )] pub tpm: Option, #[serde(skip_serializing_if = "Option::is_none")] pub use_openai_responses_path: Option, @@ -580,3 +608,42 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub web_search_billing_unit: Option, } + +/// A count the map may write as a whole float (`4096.0`), as Python reads any JSON number. +fn whole_number<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let Some(number) = Option::::deserialize(deserializer)? else { + return Ok(None); + }; + if let Some(count) = number.as_u64() { + return Ok(Some(count)); + } + match number.as_f64() { + Some(value) if value >= 0.0 && value.fract() == 0.0 && value <= u64::MAX as f64 => { + Ok(Some(value as u64)) + } + _ => Err(serde::de::Error::custom(format!( + "expected a whole non-negative number, got {number}" + ))), + } +} + +#[cfg(test)] +mod whole_number_tests { + use super::*; + + #[test] + fn counts_written_as_whole_floats_parse() { + let info: ModelInfo = serde_json::from_value(serde_json::json!({ + "prompt_cache_min_tokens": 4096.0, + "max_tokens": 8192, + })) + .unwrap(); + assert_eq!(info.prompt_cache_min_tokens, Some(4096)); + assert_eq!(info.max_tokens, Some(8192)); + assert_eq!(info.max_input_tokens, None); + assert!(serde_json::from_value::(serde_json::json!({"rpm": 1.5})).is_err()); + } +} diff --git a/litellm-rust/crates/model-catalog/tests/registry_validation.rs b/litellm-rust/crates/model-catalog/tests/registry_validation.rs index 67589010aba..cccc63bcd7e 100644 --- a/litellm-rust/crates/model-catalog/tests/registry_validation.rs +++ b/litellm-rust/crates/model-catalog/tests/registry_validation.rs @@ -21,6 +21,16 @@ fn checked_in_registry_passes_strict_validation(repo_root: PathBuf, #[case] file validate_registry(&catalog).unwrap(); } +/// The gateway prices requests from this map, so it must at least parse, whatever the strict +/// validation above says about individual entries. +#[rstest] +#[case("model_prices_and_context_window.json")] +#[case("litellm/model_prices_and_context_window_backup.json")] +fn checked_in_registry_parses(repo_root: PathBuf, #[case] filename: &str) { + let body = std::fs::read(repo_root.join(filename)).unwrap(); + Catalog::parse(&body, Provenance::default()).unwrap(); +} + #[rstest] #[ignore] fn fallback_generalizations_are_typed(repo_root: PathBuf) { diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index b05c231a263..264555ec66b 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -54,6 +54,8 @@ impl Default for Router { /// One call to make, on `deployment`. `last` is set when no retry can follow it. pub struct Attempt<'a> { pub deployment: &'a Deployment, + /// The deployment's id: `model_info.id`, or the one the router derives without it. + pub id: &'a str, pub last: bool, } @@ -247,6 +249,7 @@ impl Router { let in_flight = InFlight::enter(&member.health); let attempt = Attempt { deployment: &member.deployment, + id: &member.id, last: tried.len() - skipped >= attempts || tried.len() >= group.len(), }; let last = attempt.last; diff --git a/litellm-rust/docs/customer-proxy/hooks-pipeline.svg b/litellm-rust/docs/customer-proxy/hooks-pipeline.svg new file mode 100644 index 00000000000..093792db70f --- /dev/null +++ b/litellm-rust/docs/customer-proxy/hooks-pipeline.svg @@ -0,0 +1,75 @@ + + + + + + + Hooks and observers around one /v1/messages or chat request + + + + Auth + key store row now also gives + alias, user, team alias, + key and team metadata + + + + Request facts (built only if read) + key, end user, client headers, model, + stream, max_tokens, raw body (shared), + input estimate (bytes / 4, made once) + + + + + Hooks, in config order + each declares reads (facts) and writes (targets): + admission, upstream headers, upstream body, + response headers, annotations + same target twice: later hook wins + undeclared write: dropped with a warning + + + + + Refusal + status, proxy error body, hook headers (retry-after) + earlier hooks settle; observers still get a record + + + + + Routing and provider call (unchanged) + body edits applied to every attempt; each attempt notes its deployment id + (model_info.id, or the router's derived id) + + + + + Response tap + stream: scans SSE lines; Anthropic decides from the first 96 bytes, + so content lines are never copied; OpenAI keeps only usage chunks + JSON: parses only the top-level "usage" object; Content-Length kept + ends on last chunk or client drop (recorded as interrupted) + + + + + Cost (litellm-cost) + prices per deployment id, built at startup: + cost map entry overlaid by config prices + _above_Nk_tokens thresholds, _priority/_flex, + cache read / 5m / 1h write rates + + + + + Settle, then observers (off the request path) + hooks that asked are told how the request ended (to release reservations) + one record (facts, deployment, status, usage, cost, timings, annotations) per request + each observer: own tokio task behind a bounded queue; full queue drops and counts; shutdown drains then flushes + nothing configured: none of this runs and requests take the old path + + + From 2d2ca95481994dd00a0ba50669240b9356efe5db Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:43 -0700 Subject: [PATCH 02/10] feat(gateway): write spend to the proxy's database from a native observer (#47041) Adds a spend writer observer that records every gateway request in the proxy's own Postgres tables with the same statements the Python proxy uses: spend logs, key/user/team/team member/end user spend, and the daily user, team and end-user activity tables. Writes are batched on the proxy's timers and are additive, so replicas and Python proxies can write side by side. The database pool is now shared between virtual keys and the writer. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/Cargo.lock | 4 + ...3420df984f9ae9460ceefa9a16e61bae0e3e2.json | 15 + ...92bf33e16fea1dbb0c32cf89a376a874927f4.json | 21 + ...d48cba84f072116aca0fcedf3566f81fa14b6.json | 14 + ...70fa5f0c55f253258f26bcae44fe2be097272.json | 23 + ...5cd10df9db5c4508b859068d800ca8be02b93.json | 40 + ...855a7324d296ae3f15ccee839a29007693f87.json | 30 + ...fd1abc967056d615940b0804915a3ebed9a2b.json | 15 + ...76fbde114f184f7cc6efe70d49d4c78e8f036.json | 30 + ...45ad5e942f4896cac72c0f002f3c544e25094.json | 15 + ...1fd9c45b29330c73c66cdbee80e259dbb391b.json | 15 + ...cb0b2cba889274ab089938b1fe011c7c94df2.json | 30 + litellm-rust/crates/db/src/lib.rs | 6 + litellm-rust/crates/db/src/spend.rs | 451 +++++++++++ litellm-rust/crates/db/tests/spend.rs | 281 +++++++ litellm-rust/crates/gateway/Cargo.toml | 4 + litellm-rust/crates/gateway/src/keys.rs | 32 +- litellm-rust/crates/gateway/src/lib.rs | 22 +- litellm-rust/crates/gateway/src/main.rs | 17 +- litellm-rust/crates/gateway/src/spend.rs | 708 ++++++++++++++++++ litellm-rust/crates/gateway/tests/spend.rs | 304 ++++++++ .../docs/customer-proxy/spend-writer.svg | 52 ++ 22 files changed, 2110 insertions(+), 19 deletions(-) create mode 100644 litellm-rust/crates/db/.sqlx/query-07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2.json create mode 100644 litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json create mode 100644 litellm-rust/crates/db/.sqlx/query-8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6.json create mode 100644 litellm-rust/crates/db/.sqlx/query-83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272.json create mode 100644 litellm-rust/crates/db/.sqlx/query-8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93.json create mode 100644 litellm-rust/crates/db/.sqlx/query-94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87.json create mode 100644 litellm-rust/crates/db/.sqlx/query-94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036.json create mode 100644 litellm-rust/crates/db/.sqlx/query-b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094.json create mode 100644 litellm-rust/crates/db/.sqlx/query-bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b.json create mode 100644 litellm-rust/crates/db/.sqlx/query-f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2.json create mode 100644 litellm-rust/crates/db/src/spend.rs create mode 100644 litellm-rust/crates/db/tests/spend.rs create mode 100644 litellm-rust/crates/gateway/src/spend.rs create mode 100644 litellm-rust/crates/gateway/tests/spend.rs create mode 100644 litellm-rust/docs/customer-proxy/spend-writer.svg diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9b60fb5520e..5da517b88b3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3752,10 +3752,12 @@ dependencies = [ "base64 0.22.1", "envy", "futures-util", + "http 1.4.2", "http-body-util", "litellm-auth-types", "litellm-config", "litellm-db", + "litellm-db-testing", "litellm-gateway-auth", "litellm-gateway-hooks", "litellm-gateway-inference", @@ -3768,10 +3770,12 @@ dependencies = [ "litellm-router", "litellm-secrets", "litellm-tracing", + "rand 0.8.7", "redis", "rstest", "serde", "serde_json", + "sqlx", "strum", "tempfile", "thiserror 2.0.19", diff --git a/litellm-rust/crates/db/.sqlx/query-07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2.json b/litellm-rust/crates/db/.sqlx/query-07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2.json new file mode 100644 index 00000000000..876eb6b645c --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)),\n locked AS (SELECT t.token FROM \"LiteLLM_VerificationToken\" t\n WHERE t.token IN (SELECT id FROM added) ORDER BY t.token FOR UPDATE)\n UPDATE \"LiteLLM_VerificationToken\" t\n SET spend = t.spend + added.cost,\n total_spend = t.total_spend + added.cost,\n last_active = (NOW() AT TIME ZONE 'UTC'),\n updated_at = (NOW() AT TIME ZONE 'UTC')\n FROM added\n WHERE t.token = added.id AND t.token IN (SELECT token FROM locked)", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "07328686221e3bfe19c205d68fe3420df984f9ae9460ceefa9a16e61bae0e3e2" +} diff --git a/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json b/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json new file mode 100644 index 00000000000..7eef5278cb5 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json @@ -0,0 +1,21 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT set_config('lock_timeout', '5000ms', true)", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "set_config", + "type_info": "Text", + "origin": "Expression" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + null + ] + }, + "hash": "5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4" +} diff --git a/litellm-rust/crates/db/.sqlx/query-8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6.json b/litellm-rust/crates/db/.sqlx/query-8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6.json new file mode 100644 index 00000000000..071a2a834ef --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_TeamMembership\" (user_id, team_id, spend, total_spend)\n SELECT member.user_id, member.team_id, member.cost, member.cost\n FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8)\n WHERE EXISTS (SELECT 1 FROM \"LiteLLM_TeamTable\" t WHERE t.team_id = member.team_id\n AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id)))\n OR EXISTS (SELECT 1 FROM \"LiteLLM_TeamMembership\" m\n WHERE m.user_id = member.user_id AND m.team_id = member.team_id)\n ON CONFLICT (user_id, team_id) DO UPDATE\n SET spend = \"LiteLLM_TeamMembership\".spend + EXCLUDED.spend,\n total_spend = \"LiteLLM_TeamMembership\".total_spend + EXCLUDED.total_spend", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Jsonb" + ] + }, + "nullable": [] + }, + "hash": "8337c5f41713afadeff72900c4cd48cba84f072116aca0fcedf3566f81fa14b6" +} diff --git a/litellm-rust/crates/db/.sqlx/query-83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272.json b/litellm-rust/crates/db/.sqlx/query-83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272.json new file mode 100644 index 00000000000..97ebe2f8fda --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS \"locked!\"", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "locked!", + "type_info": "Bool", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "83be2dd1f5892e3581ecee3d2d270fa5f0c55f253258f26bcae44fe2be097272" +} diff --git a/litellm-rust/crates/db/.sqlx/query-8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93.json b/litellm-rust/crates/db/.sqlx/query-8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93.json new file mode 100644 index 00000000000..db44ac1c3c6 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93.json @@ -0,0 +1,40 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_SpendLogs\" (\n request_id, call_type, api_key, spend, total_tokens, prompt_tokens,\n completion_tokens, \"startTime\", \"endTime\", \"completionStartTime\",\n request_duration_ms, model, model_id, model_group, custom_llm_provider, api_base,\n \"user\", metadata, cache_hit, cache_key, request_tags, team_id, organization_id,\n end_user, requester_ip_address, messages, response, proxy_server_request,\n session_id, status, litellm_call_id)\n SELECT l.request_id, l.call_type, l.api_key, l.spend, l.total_tokens, l.prompt_tokens,\n l.completion_tokens, l.start_time::timestamp(3), l.end_time::timestamp(3),\n l.completion_start_time::timestamp(3), l.request_duration_ms, l.model,\n l.model_id, l.model_group, l.custom_llm_provider, l.api_base, l.user_id,\n l.metadata, l.cache_hit, l.cache_key, '[]'::jsonb, l.team_id,\n l.organization_id, l.end_user, l.requester_ip_address, '{}'::jsonb,\n '{}'::jsonb, '{}'::jsonb, l.session_id, l.status, l.litellm_call_id\n FROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::int4[], $6::int4[],\n $7::int4[], $8::text[], $9::text[], $10::text[], $11::int4[], $12::text[],\n $13::text[], $14::text[], $15::text[], $16::text[], $17::text[],\n $18::jsonb[], $19::text[], $20::text[], $21::text[], $22::text[],\n $23::text[], $24::text[], $25::text[], $26::text[], $27::text[])\n AS l(request_id, call_type, api_key, spend, total_tokens, prompt_tokens,\n completion_tokens, start_time, end_time, completion_start_time,\n request_duration_ms, model, model_id, model_group, custom_llm_provider,\n api_base, user_id, metadata, cache_hit, cache_key, team_id, organization_id,\n end_user, requester_ip_address, session_id, status, litellm_call_id)\n ON CONFLICT (request_id) DO NOTHING", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "Float8Array", + "Int4Array", + "Int4Array", + "Int4Array", + "TextArray", + "TextArray", + "TextArray", + "Int4Array", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "JsonbArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray" + ] + }, + "nullable": [] + }, + "hash": "8abe8b069d1b8bb22f4ff49046e5cd10df9db5c4508b859068d800ca8be02b93" +} diff --git a/litellm-rust/crates/db/.sqlx/query-94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87.json b/litellm-rust/crates/db/.sqlx/query-94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87.json new file mode 100644 index 00000000000..9eb6a78c06e --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87.json @@ -0,0 +1,30 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_DailyUserSpend\" (id, user_id, date, api_key, model,\n custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group,\n prompt_tokens, completion_tokens, api_requests, successful_requests,\n failed_requests, cache_read_input_tokens, cache_creation_input_tokens,\n total_response_time_ms, timed_requests, spend, updated_at)\n SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, '', r.endpoint, r.model_group, r.prompt_tokens,\n r.completion_tokens, r.api_requests, r.successful_requests,\n r.failed_requests, r.cache_read, r.cache_creation, r.response_time,\n r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')\n FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::text[],\n $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],\n $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],\n $16::int8[], $17::float8[])\n AS r(entity_id, date, api_key, model, custom_llm_provider, endpoint,\n model_group, prompt_tokens, completion_tokens, api_requests,\n successful_requests, failed_requests, cache_read, cache_creation,\n response_time, timed_requests, spend)\n ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider\n ON CONFLICT (user_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyUserSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyUserSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyUserSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyUserSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyUserSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyUserSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyUserSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyUserSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyUserSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyUserSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "94187798c95011cae91b61417c8855a7324d296ae3f15ccee839a29007693f87" +} diff --git a/litellm-rust/crates/db/.sqlx/query-94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b.json b/litellm-rust/crates/db/.sqlx/query-94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b.json new file mode 100644 index 00000000000..f2d7ba555b9 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)),\n locked AS (SELECT t.team_id FROM \"LiteLLM_TeamTable\" t\n WHERE t.team_id IN (SELECT id FROM added) ORDER BY t.team_id FOR UPDATE)\n UPDATE \"LiteLLM_TeamTable\" t\n SET spend = t.spend + added.cost, updated_at = (NOW() AT TIME ZONE 'UTC')\n FROM added\n WHERE t.team_id = added.id AND t.team_id IN (SELECT team_id FROM locked)", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "94e8e0b4897318bd005c9d4cd50fd1abc967056d615940b0804915a3ebed9a2b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036.json b/litellm-rust/crates/db/.sqlx/query-9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036.json new file mode 100644 index 00000000000..c2d058c96ac --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036.json @@ -0,0 +1,30 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_DailyTeamSpend\" (id, team_id, date, api_key, model,\n custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group,\n prompt_tokens, completion_tokens, api_requests, successful_requests,\n failed_requests, cache_read_input_tokens, cache_creation_input_tokens,\n total_response_time_ms, timed_requests, spend, updated_at)\n SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, '', r.endpoint, r.model_group, r.prompt_tokens,\n r.completion_tokens, r.api_requests, r.successful_requests,\n r.failed_requests, r.cache_read, r.cache_creation, r.response_time,\n r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')\n FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::text[],\n $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],\n $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],\n $16::int8[], $17::float8[])\n AS r(entity_id, date, api_key, model, custom_llm_provider, endpoint,\n model_group, prompt_tokens, completion_tokens, api_requests,\n successful_requests, failed_requests, cache_read, cache_creation,\n response_time, timed_requests, spend)\n ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider\n ON CONFLICT (team_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyTeamSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyTeamSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyTeamSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyTeamSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyTeamSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyTeamSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyTeamSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyTeamSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyTeamSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyTeamSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "9e6d4c7704bfbb38ce6d12b647376fbde114f184f7cc6efe70d49d4c78e8f036" +} diff --git a/litellm-rust/crates/db/.sqlx/query-b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094.json b/litellm-rust/crates/db/.sqlx/query-b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094.json new file mode 100644 index 00000000000..90d6f426bb6 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)),\n locked AS (SELECT u.user_id FROM \"LiteLLM_UserTable\" u\n WHERE u.user_id IN (SELECT id FROM added) ORDER BY u.user_id FOR UPDATE)\n UPDATE \"LiteLLM_UserTable\" u\n SET spend = u.spend + added.cost, updated_at = (NOW() AT TIME ZONE 'UTC')\n FROM added\n WHERE u.user_id = added.id AND u.user_id IN (SELECT user_id FROM locked)", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "b0e3f762f1dd5efdd4672953f7745ad5e942f4896cac72c0f002f3c544e25094" +} diff --git a/litellm-rust/crates/db/.sqlx/query-bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b.json b/litellm-rust/crates/db/.sqlx/query-bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b.json new file mode 100644 index 00000000000..7e046d3f04f --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_EndUserTable\" (user_id, spend, blocked)\n SELECT id, cost, false FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost) ORDER BY id\n ON CONFLICT (user_id) DO UPDATE\n SET spend = \"LiteLLM_EndUserTable\".spend + EXCLUDED.spend", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "bb412ba2896341351e1f6ca96761fd9c45b29330c73c66cdbee80e259dbb391b" +} diff --git a/litellm-rust/crates/db/.sqlx/query-f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2.json b/litellm-rust/crates/db/.sqlx/query-f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2.json new file mode 100644 index 00000000000..5824982661f --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2.json @@ -0,0 +1,30 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_DailyEndUserSpend\" (id, end_user_id, date, api_key, model,\n custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group,\n prompt_tokens, completion_tokens, api_requests, successful_requests,\n failed_requests, cache_read_input_tokens, cache_creation_input_tokens,\n total_response_time_ms, timed_requests, spend, updated_at)\n SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, '', r.endpoint, r.model_group, r.prompt_tokens,\n r.completion_tokens, r.api_requests, r.successful_requests,\n r.failed_requests, r.cache_read, r.cache_creation, r.response_time,\n r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')\n FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::text[],\n $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],\n $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],\n $16::int8[], $17::float8[])\n AS r(entity_id, date, api_key, model, custom_llm_provider, endpoint,\n model_group, prompt_tokens, completion_tokens, api_requests,\n successful_requests, failed_requests, cache_read, cache_creation,\n response_time, timed_requests, spend)\n ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider\n ON CONFLICT (end_user_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyEndUserSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyEndUserSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyEndUserSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyEndUserSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyEndUserSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyEndUserSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyEndUserSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyEndUserSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyEndUserSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyEndUserSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Int8Array", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "f56e12df3f7b46178c50a58013ecb0b2cba889274ab089938b1fe011c7c94df2" +} diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index c099edb0003..39ef58241b2 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -1,11 +1,17 @@ mod error; mod failure; mod pool; +mod spend; mod verification_token; pub use error::Error; pub use failure::DbFailure; pub use pool::{PoolSettings, connect_lazy}; +pub use spend::{ + DailySpend, DailyTable, EntitySpend, SpendLog, TeamMemberSpend, add_end_user_spend, + add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, insert_spend_logs, + upsert_daily_spend, +}; pub use sqlx::PgPool; pub use verification_token::{ Restrictions, Team, TeamMembership, User, VerificationToken, find_verification_token, diff --git a/litellm-rust/crates/db/src/spend.rs b/litellm-rust/crates/db/src/spend.rs new file mode 100644 index 00000000000..67d11e4c7d3 --- /dev/null +++ b/litellm-rust/crates/db/src/spend.rs @@ -0,0 +1,451 @@ +//! The proxy's spend writes: additive spend on the key, user, team, team membership and end user +//! rows, rows in `LiteLLM_SpendLogs`, and the daily activity tables. Every statement takes a +//! batch, so a flush costs one round trip per table. + +use sqlx::{PgPool, Postgres, Transaction}; + +use crate::Error; + +/// Spend to add per id. Ids are locked in sorted order, as the proxy does, so writers on other +/// replicas cannot deadlock with this one. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct EntitySpend { + pub ids: Vec, + pub costs: Vec, +} + +impl EntitySpend { + pub fn is_empty(&self) -> bool { + self.ids.is_empty() + } +} + +impl FromIterator<(String, f64)> for EntitySpend { + fn from_iter>(entries: I) -> Self { + let mut entries: Vec<_> = entries.into_iter().collect(); + entries.sort_by(|left, right| left.0.cmp(&right.0)); + let (ids, costs) = entries.into_iter().unzip(); + Self { ids, costs } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct TeamMemberSpend { + pub team_id: String, + pub user_id: String, + pub cost: f64, +} + +/// Starts a transaction that waits at most 5 s for a row lock, the proxy's +/// `SPEND_ROLLUP_LOCK_TIMEOUT_MS`. +async fn begin(pool: &PgPool) -> Result, Error> { + let mut transaction = pool.begin().await?; + sqlx::query!("SELECT set_config('lock_timeout', '5000ms', true)") + .fetch_one(&mut *transaction) + .await?; + Ok(transaction) +} + +/// `LiteLLM_VerificationToken.spend` and `total_spend`, and `last_active` set to now. +pub async fn add_key_spend(pool: &PgPool, spend: &EntitySpend) -> Result<(), Error> { + if spend.is_empty() { + return Ok(()); + } + let mut transaction = begin(pool).await?; + sqlx::query!( + r#"WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)), + locked AS (SELECT t.token FROM "LiteLLM_VerificationToken" t + WHERE t.token IN (SELECT id FROM added) ORDER BY t.token FOR UPDATE) + UPDATE "LiteLLM_VerificationToken" t + SET spend = t.spend + added.cost, + total_spend = t.total_spend + added.cost, + last_active = (NOW() AT TIME ZONE 'UTC'), + updated_at = (NOW() AT TIME ZONE 'UTC') + FROM added + WHERE t.token = added.id AND t.token IN (SELECT token FROM locked)"#, + &spend.ids, + &spend.costs, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// `LiteLLM_UserTable.spend`. A user without a row is skipped, as the proxy skips it. +pub async fn add_user_spend(pool: &PgPool, spend: &EntitySpend) -> Result<(), Error> { + if spend.is_empty() { + return Ok(()); + } + let mut transaction = begin(pool).await?; + sqlx::query!( + r#"WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)), + locked AS (SELECT u.user_id FROM "LiteLLM_UserTable" u + WHERE u.user_id IN (SELECT id FROM added) ORDER BY u.user_id FOR UPDATE) + UPDATE "LiteLLM_UserTable" u + SET spend = u.spend + added.cost, updated_at = (NOW() AT TIME ZONE 'UTC') + FROM added + WHERE u.user_id = added.id AND u.user_id IN (SELECT user_id FROM locked)"#, + &spend.ids, + &spend.costs, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// `LiteLLM_TeamTable.spend`. +pub async fn add_team_spend(pool: &PgPool, spend: &EntitySpend) -> Result<(), Error> { + if spend.is_empty() { + return Ok(()); + } + let mut transaction = begin(pool).await?; + sqlx::query!( + r#"WITH added AS (SELECT * FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)), + locked AS (SELECT t.team_id FROM "LiteLLM_TeamTable" t + WHERE t.team_id IN (SELECT id FROM added) ORDER BY t.team_id FOR UPDATE) + UPDATE "LiteLLM_TeamTable" t + SET spend = t.spend + added.cost, updated_at = (NOW() AT TIME ZONE 'UTC') + FROM added + WHERE t.team_id = added.id AND t.team_id IN (SELECT team_id FROM locked)"#, + &spend.ids, + &spend.costs, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// `LiteLLM_EndUserTable.spend`, creating the end user on first spend as the proxy does. +pub async fn add_end_user_spend(pool: &PgPool, spend: &EntitySpend) -> Result<(), Error> { + if spend.is_empty() { + return Ok(()); + } + let mut transaction = begin(pool).await?; + sqlx::query!( + r#"INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked) + SELECT id, cost, false FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost) ORDER BY id + ON CONFLICT (user_id) DO UPDATE + SET spend = "LiteLLM_EndUserTable".spend + EXCLUDED.spend"#, + &spend.ids, + &spend.costs, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// `LiteLLM_TeamMembership.spend` and `total_spend`. A membership row is created only for a +/// user on the team's roster, under a per-team advisory lock, as the proxy writes it. +pub async fn add_team_member_spend( + pool: &PgPool, + members: &[TeamMemberSpend], +) -> Result<(), Error> { + if members.is_empty() { + return Ok(()); + } + let mut members = members.to_vec(); + members.sort_by(|left, right| { + (&left.team_id, &left.user_id).cmp(&(&right.team_id, &right.user_id)) + }); + let mut transaction = begin(pool).await?; + let mut teams: Vec<&str> = members + .iter() + .map(|member| member.team_id.as_str()) + .collect(); + teams.dedup(); + for team in teams { + sqlx::query!( + r#"SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS "locked!""#, + team + ) + .fetch_one(&mut *transaction) + .await?; + } + let rows = serde_json::Value::Array( + members + .iter() + .map(|member| { + serde_json::json!({ + "user_id": member.user_id, + "team_id": member.team_id, + "cost": member.cost, + }) + }) + .collect(), + ); + sqlx::query!( + r#"INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend) + SELECT member.user_id, member.team_id, member.cost, member.cost + FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8) + WHERE EXISTS (SELECT 1 FROM "LiteLLM_TeamTable" t WHERE t.team_id = member.team_id + AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))) + OR EXISTS (SELECT 1 FROM "LiteLLM_TeamMembership" m + WHERE m.user_id = member.user_id AND m.team_id = member.team_id) + ON CONFLICT (user_id, team_id) DO UPDATE + SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, + total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend"#, + rows, + ) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) +} + +/// One `LiteLLM_SpendLogs` row. Times are UTC, `YYYY-MM-DD HH:MM:SS.fff`. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SpendLog { + pub request_id: String, + pub call_type: String, + pub api_key: String, + pub spend: f64, + pub total_tokens: i32, + pub prompt_tokens: i32, + pub completion_tokens: i32, + pub start_time: String, + pub end_time: String, + pub completion_start_time: Option, + pub request_duration_ms: Option, + pub model: String, + pub model_id: String, + pub model_group: String, + pub custom_llm_provider: String, + pub api_base: String, + pub user: String, + pub metadata: serde_json::Value, + pub cache_hit: String, + pub cache_key: String, + pub team_id: String, + pub organization_id: String, + pub end_user: String, + pub requester_ip_address: Option, + pub session_id: Option, + pub status: String, + pub litellm_call_id: Option, +} + +/// Inserts spend logs, skipping request ids already written, as the proxy's +/// `create_many(skip_duplicates=True)` does. +pub async fn insert_spend_logs(pool: &PgPool, logs: &[SpendLog]) -> Result<(), Error> { + if logs.is_empty() { + return Ok(()); + } + macro_rules! column { + ($field:ident) => { + logs.iter() + .map(|log| log.$field.clone()) + .collect::>() + }; + } + sqlx::query!( + r#"INSERT INTO "LiteLLM_SpendLogs" ( + request_id, call_type, api_key, spend, total_tokens, prompt_tokens, + completion_tokens, "startTime", "endTime", "completionStartTime", + request_duration_ms, model, model_id, model_group, custom_llm_provider, api_base, + "user", metadata, cache_hit, cache_key, request_tags, team_id, organization_id, + end_user, requester_ip_address, messages, response, proxy_server_request, + session_id, status, litellm_call_id) + SELECT l.request_id, l.call_type, l.api_key, l.spend, l.total_tokens, l.prompt_tokens, + l.completion_tokens, l.start_time::timestamp(3), l.end_time::timestamp(3), + l.completion_start_time::timestamp(3), l.request_duration_ms, l.model, + l.model_id, l.model_group, l.custom_llm_provider, l.api_base, l.user_id, + l.metadata, l.cache_hit, l.cache_key, '[]'::jsonb, l.team_id, + l.organization_id, l.end_user, l.requester_ip_address, '{}'::jsonb, + '{}'::jsonb, '{}'::jsonb, l.session_id, l.status, l.litellm_call_id + FROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::int4[], $6::int4[], + $7::int4[], $8::text[], $9::text[], $10::text[], $11::int4[], $12::text[], + $13::text[], $14::text[], $15::text[], $16::text[], $17::text[], + $18::jsonb[], $19::text[], $20::text[], $21::text[], $22::text[], + $23::text[], $24::text[], $25::text[], $26::text[], $27::text[]) + AS l(request_id, call_type, api_key, spend, total_tokens, prompt_tokens, + completion_tokens, start_time, end_time, completion_start_time, + request_duration_ms, model, model_id, model_group, custom_llm_provider, + api_base, user_id, metadata, cache_hit, cache_key, team_id, organization_id, + end_user, requester_ip_address, session_id, status, litellm_call_id) + ON CONFLICT (request_id) DO NOTHING"#, + &column!(request_id), + &column!(call_type), + &column!(api_key), + &column!(spend), + &column!(total_tokens), + &column!(prompt_tokens), + &column!(completion_tokens), + &column!(start_time), + &column!(end_time), + &column!(completion_start_time) as &[Option], + &column!(request_duration_ms) as &[Option], + &column!(model), + &column!(model_id), + &column!(model_group), + &column!(custom_llm_provider), + &column!(api_base), + &column!(user), + &column!(metadata), + &column!(cache_hit), + &column!(cache_key), + &column!(team_id), + &column!(organization_id), + &column!(end_user), + &column!(requester_ip_address) as &[Option], + &column!(session_id) as &[Option], + &column!(status), + &column!(litellm_call_id) as &[Option], + ) + .execute(pool) + .await?; + Ok(()) +} + +/// The daily activity tables, keyed by user, team or end user. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum DailyTable { + User, + Team, + EndUser, +} + +/// One row of a daily activity table. Key columns are `''` rather than NULL, as the proxy +/// writes them, so its unique index matches. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct DailySpend { + pub entity_id: String, + /// `YYYY-MM-DD`, UTC. + pub date: String, + pub api_key: String, + pub model: String, + pub custom_llm_provider: String, + pub endpoint: String, + /// Written on insert only. + pub model_group: Option, + pub prompt_tokens: i64, + pub completion_tokens: i64, + pub api_requests: i64, + pub successful_requests: i64, + pub failed_requests: i64, + pub cache_read_input_tokens: i64, + pub cache_creation_input_tokens: i64, + pub total_response_time_ms: i64, + pub timed_requests: i64, + pub spend: f64, +} + +/// Adds `rows` to a daily table. Rows must be unique by their key columns. +pub async fn upsert_daily_spend( + pool: &PgPool, + table: DailyTable, + rows: &[DailySpend], +) -> Result<(), Error> { + if rows.is_empty() { + return Ok(()); + } + macro_rules! column { + ($field:ident) => { + rows.iter() + .map(|row| row.$field.clone()) + .collect::>() + }; + } + let entity_id = column!(entity_id); + let date = column!(date); + let api_key = column!(api_key); + let model = column!(model); + let custom_llm_provider = column!(custom_llm_provider); + let endpoint = column!(endpoint); + let model_group: Vec> = column!(model_group); + let prompt_tokens = column!(prompt_tokens); + let completion_tokens = column!(completion_tokens); + let api_requests = column!(api_requests); + let successful_requests = column!(successful_requests); + let failed_requests = column!(failed_requests); + let cache_read = column!(cache_read_input_tokens); + let cache_creation = column!(cache_creation_input_tokens); + let response_time = column!(total_response_time_ms); + let timed_requests = column!(timed_requests); + let spend = column!(spend); + let mut transaction = begin(pool).await?; + // One statement per table: a table name cannot be a parameter. + macro_rules! upsert { + ($table:literal, $entity:literal) => { + sqlx::query!( + r#"INSERT INTO "# + $table + r#" (id, "# + $entity + r#", date, api_key, model, + custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group, + prompt_tokens, completion_tokens, api_requests, successful_requests, + failed_requests, cache_read_input_tokens, cache_creation_input_tokens, + total_response_time_ms, timed_requests, spend, updated_at) + SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model, + r.custom_llm_provider, '', r.endpoint, r.model_group, r.prompt_tokens, + r.completion_tokens, r.api_requests, r.successful_requests, + r.failed_requests, r.cache_read, r.cache_creation, r.response_time, + r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC') + FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::text[], + $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[], + $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[], + $16::int8[], $17::float8[]) + AS r(entity_id, date, api_key, model, custom_llm_provider, endpoint, + model_group, prompt_tokens, completion_tokens, api_requests, + successful_requests, failed_requests, cache_read, cache_creation, + response_time, timed_requests, spend) + ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider + ON CONFLICT ("# + $entity + r#", date, api_key, model, custom_llm_provider, + mcp_namespaced_tool_name, endpoint) DO UPDATE SET + prompt_tokens = "# + $table + r#".prompt_tokens + EXCLUDED.prompt_tokens, + completion_tokens = "# + $table + r#".completion_tokens + EXCLUDED.completion_tokens, + api_requests = "# + $table + r#".api_requests + EXCLUDED.api_requests, + successful_requests = "# + $table + r#".successful_requests + EXCLUDED.successful_requests, + failed_requests = "# + $table + r#".failed_requests + EXCLUDED.failed_requests, + cache_read_input_tokens = "# + $table + r#".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens, + cache_creation_input_tokens = "# + $table + r#".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens, + total_response_time_ms = "# + $table + r#".total_response_time_ms + EXCLUDED.total_response_time_ms, + timed_requests = "# + $table + r#".timed_requests + EXCLUDED.timed_requests, + spend = "# + $table + r#".spend + EXCLUDED.spend, + updated_at = (NOW() AT TIME ZONE 'UTC')"#, + &entity_id, + &date, + &api_key, + &model, + &custom_llm_provider, + &endpoint, + &model_group as &[Option], + &prompt_tokens, + &completion_tokens, + &api_requests, + &successful_requests, + &failed_requests, + &cache_read, + &cache_creation, + &response_time, + &timed_requests, + &spend, + ) + .execute(&mut *transaction) + .await? + }; + } + match table { + DailyTable::User => upsert!(r#""LiteLLM_DailyUserSpend""#, "user_id"), + DailyTable::Team => upsert!(r#""LiteLLM_DailyTeamSpend""#, "team_id"), + DailyTable::EndUser => upsert!(r#""LiteLLM_DailyEndUserSpend""#, "end_user_id"), + }; + transaction.commit().await?; + Ok(()) +} + +impl Error { + /// Whether retrying the same write can succeed: a deadlock, a lock wait timeout, a + /// serialization failure, or a lost connection. + pub fn is_retryable(&self) -> bool { + match self { + Self::Query(sqlx::Error::Database(error)) => { + matches!(error.code().as_deref(), Some("40P01" | "55P03" | "40001")) + } + Self::Query( + sqlx::Error::Io(_) | sqlx::Error::PoolTimedOut | sqlx::Error::PoolClosed, + ) => true, + Self::Query(_) | Self::Url(_) => false, + } + } +} diff --git a/litellm-rust/crates/db/tests/spend.rs b/litellm-rust/crates/db/tests/spend.rs new file mode 100644 index 00000000000..b01ef55a68c --- /dev/null +++ b/litellm-rust/crates/db/tests/spend.rs @@ -0,0 +1,281 @@ +use litellm_db::{ + DailySpend, DailyTable, EntitySpend, SpendLog, TeamMemberSpend, add_end_user_spend, + add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, insert_spend_logs, + upsert_daily_spend, +}; +use litellm_db_testing::MigratedPostgres; +use rstest::rstest; +use serde_json::{Value, json}; +use sqlx::PgPool; + +const FIXTURES: &str = r#" +INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles, created_at, updated_at) + VALUES ('team-1', '[{"user_id": "user-1", "role": "user"}]', now(), now()); +INSERT INTO "LiteLLM_UserTable" (user_id, spend) VALUES ('user-1', 1.0); +INSERT INTO "LiteLLM_VerificationToken" (token, spend, total_spend, team_id, user_id) + VALUES ('hash-1', 2.0, 2.0, 'team-1', 'user-1'); +"#; + +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows and assertions; the writes under test use the checked macro" +)] +async fn setup(pool: &PgPool) { + sqlx::raw_sql(FIXTURES).execute(pool).await.unwrap(); +} + +#[expect( + clippy::disallowed_methods, + reason = "test assertions read rows back; the writes under test use the checked macro" +)] +async fn read(pool: &PgPool, sql: &str) -> Value { + sqlx::query_scalar::<_, Value>(sqlx::AssertSqlSafe(format!( + "SELECT to_jsonb(r) FROM ({sql}) r" + ))) + .fetch_one(pool) + .await + .unwrap() +} + +#[expect( + clippy::disallowed_methods, + reason = "test assertions count rows; the writes under test use the checked macro" +)] +async fn count(pool: &PgPool, sql: &str) -> i64 { + sqlx::query_scalar::<_, i64>(sqlx::AssertSqlSafe(sql.to_owned())) + .fetch_one(pool) + .await + .unwrap() +} + +fn spend(entries: &[(&str, f64)]) -> EntitySpend { + entries + .iter() + .map(|(id, cost)| ((*id).to_owned(), *cost)) + .collect() +} + +#[rstest] +#[tokio::test] +async fn entity_spend_adds_up_and_creates_only_what_the_proxy_creates() { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + setup(pool).await; + + for _ in 0..2 { + add_key_spend(pool, &spend(&[("hash-1", 0.25), ("missing", 1.0)])) + .await + .unwrap(); + add_user_spend(pool, &spend(&[("user-1", 0.5), ("missing", 1.0)])) + .await + .unwrap(); + add_team_spend(pool, &spend(&[("team-1", 0.125)])) + .await + .unwrap(); + add_end_user_spend(pool, &spend(&[("end-1", 0.75)])) + .await + .unwrap(); + add_team_member_spend( + pool, + &[ + TeamMemberSpend { + team_id: "team-1".into(), + user_id: "user-1".into(), + cost: 0.5, + }, + TeamMemberSpend { + team_id: "team-1".into(), + user_id: "stranger".into(), + cost: 9.0, + }, + ], + ) + .await + .unwrap(); + } + + let key = read( + pool, + r#"SELECT spend, total_spend, last_active IS NOT NULL AS active + FROM "LiteLLM_VerificationToken" WHERE token = 'hash-1'"#, + ) + .await; + assert_eq!( + key, + json!({"spend": 2.5, "total_spend": 2.5, "active": true}) + ); + let user = read( + pool, + r#"SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = 'user-1'"#, + ) + .await; + assert_eq!(user, json!({"spend": 2})); + let team = read( + pool, + r#"SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = 'team-1'"#, + ) + .await; + assert_eq!(team, json!({"spend": 0.25})); + let end_user = read( + pool, + r#"SELECT spend, blocked FROM "LiteLLM_EndUserTable" WHERE user_id = 'end-1'"#, + ) + .await; + assert_eq!(end_user, json!({"spend": 1.5, "blocked": false})); + let member = read( + pool, + r#"SELECT spend, total_spend FROM "LiteLLM_TeamMembership" + WHERE team_id = 'team-1' AND user_id = 'user-1'"#, + ) + .await; + assert_eq!(member, json!({"spend": 1, "total_spend": 1})); + assert_eq!( + count(pool, r#"SELECT count(*) FROM "LiteLLM_UserTable""#).await, + 1 + ); + assert_eq!( + count( + pool, + r#"SELECT count(*) FROM "LiteLLM_TeamMembership" WHERE user_id = 'stranger'"# + ) + .await, + 0 + ); +} + +fn log(request_id: &str) -> SpendLog { + SpendLog { + request_id: request_id.into(), + call_type: "anthropic_messages".into(), + api_key: "hash-1".into(), + spend: 0.5, + total_tokens: 30, + prompt_tokens: 10, + completion_tokens: 20, + start_time: "2026-10-17 01:02:03.456".into(), + end_time: "2026-10-17 01:02:04.000".into(), + completion_start_time: Some("2026-10-17 01:02:03.500".into()), + request_duration_ms: Some(544), + model: "claude-x".into(), + model_group: "claude".into(), + custom_llm_provider: "anthropic".into(), + metadata: json!({"user_api_key": "hash-1", "status": null}), + cache_hit: "None".into(), + cache_key: "Cache OFF".into(), + status: "success".into(), + ..SpendLog::default() + } +} + +#[rstest] +#[tokio::test] +async fn spend_logs_are_written_once() { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + insert_spend_logs(pool, &[log("a"), log("b")]) + .await + .unwrap(); + insert_spend_logs(pool, &[log("a")]).await.unwrap(); + + assert_eq!( + count(pool, r#"SELECT count(*) FROM "LiteLLM_SpendLogs""#).await, + 2 + ); + let row = read( + pool, + r#"SELECT "startTime", "completionStartTime", request_duration_ms, team_id, end_user, + metadata, request_tags, messages, response, cache_hit, status + FROM "LiteLLM_SpendLogs" WHERE request_id = 'a'"#, + ) + .await; + assert_eq!( + row, + json!({ + "startTime": "2026-10-17T01:02:03.456", + "completionStartTime": "2026-10-17T01:02:03.5", + "request_duration_ms": 544, + "team_id": "", + "end_user": "", + "metadata": {"user_api_key": "hash-1", "status": null}, + "request_tags": [], + "messages": {}, + "response": {}, + "cache_hit": "None", + "status": "success", + }) + ); +} + +fn daily(entity: &str, model_group: &str) -> DailySpend { + DailySpend { + entity_id: entity.into(), + date: "2026-10-17".into(), + api_key: "hash-1".into(), + model: "claude-x".into(), + custom_llm_provider: "anthropic".into(), + endpoint: String::new(), + model_group: Some(model_group.into()), + prompt_tokens: 10, + completion_tokens: 20, + api_requests: 1, + successful_requests: 1, + cache_read_input_tokens: 3, + total_response_time_ms: 100, + timed_requests: 1, + spend: 0.5, + ..DailySpend::default() + } +} + +#[rstest] +#[case::user(DailyTable::User, "LiteLLM_DailyUserSpend", "user_id")] +#[case::team(DailyTable::Team, "LiteLLM_DailyTeamSpend", "team_id")] +#[case::end_user(DailyTable::EndUser, "LiteLLM_DailyEndUserSpend", "end_user_id")] +#[tokio::test] +async fn daily_rows_add_up_and_keep_their_first_model_group( + #[case] table: DailyTable, + #[case] name: &str, + #[case] entity: &str, +) { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + upsert_daily_spend(pool, table, &[daily("e-1", "first"), daily("e-2", "first")]) + .await + .unwrap(); + upsert_daily_spend(pool, table, &[daily("e-1", "second")]) + .await + .unwrap(); + + let row = read( + pool, + &format!( + r#"SELECT model_group, prompt_tokens, completion_tokens, api_requests, + successful_requests, failed_requests, cache_read_input_tokens, + total_response_time_ms, timed_requests, spend, endpoint, + mcp_namespaced_tool_name + FROM "{name}" WHERE {entity} = 'e-1'"# + ), + ) + .await; + assert_eq!( + row, + json!({ + "model_group": "first", + "prompt_tokens": 20, + "completion_tokens": 40, + "api_requests": 2, + "successful_requests": 2, + "failed_requests": 0, + "cache_read_input_tokens": 6, + "total_response_time_ms": 200, + "timed_requests": 2, + "spend": 1, + "endpoint": "", + "mcp_namespaced_tool_name": "", + }) + ); + assert_eq!( + count(pool, &format!(r#"SELECT count(*) FROM "{name}""#)).await, + 2 + ); +} diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 7d5ccf6069d..d7e2e59bb53 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -15,6 +15,8 @@ litellm-gateway-mcp.workspace = true tokio-util = "0.7" thiserror.workspace = true futures-util.workspace = true +http.workspace = true +rand.workspace = true litellm-gateway-auth.workspace = true litellm-gateway-hooks.workspace = true litellm-model-catalog.workspace = true @@ -38,6 +40,8 @@ tower-sessions-moka-store = "0.15.0" [dev-dependencies] base64.workspace = true +litellm-db-testing.workspace = true +sqlx.workspace = true rstest.workspace = true tempfile.workspace = true tokio = { workspace = true, features = ["sync", "io-util"] } diff --git a/litellm-rust/crates/gateway/src/keys.rs b/litellm-rust/crates/gateway/src/keys.rs index e4a92953e84..fdf52ff90e0 100644 --- a/litellm-rust/crates/gateway/src/keys.rs +++ b/litellm-rust/crates/gateway/src/keys.rs @@ -39,24 +39,16 @@ impl KeyLookup for PostgresKeys { } } -/// The key store, when a database is configured: `general_settings.database_url` (which may be +/// The proxy's database, when one is configured: `general_settings.database_url` (which may be /// an `os.environ/` reference) or else `DATABASE_URL`, as the Python proxy resolves it. The pool /// connects on first use, so the gateway starts, and the master key works, with the database down. -pub async fn store( +pub async fn database( config: &Config, secrets: &dyn SecretSource, -) -> Result>, Error> { +) -> Result, Error> { let Some(url) = database_url(config, secrets).await? else { return Ok(None); }; - // Inference does not run global guardrails or policies, so a virtual key would get past - // them; refuse to start rather than serve keys the operator expects to be guarded. - if !config.guardrails.is_empty() - || !config.policies.is_empty() - || !config.policy_attachments.is_empty() - { - return Err(Error::GuardedVirtualKeys); - } let general = &config.general_settings; let defaults = PoolSettings::default(); let settings = PoolSettings { @@ -73,11 +65,23 @@ pub async fn store( .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) .or(defaults.idle_timeout), }; - let pool = litellm_db::connect_lazy(url.expose(), &settings)?; - Ok(Some(Arc::new(PostgresKeys { + Ok(Some(litellm_db::connect_lazy(url.expose(), &settings)?)) +} + +/// The key store over the proxy's database. +pub fn store(config: &Config, pool: PgPool) -> Result, Error> { + // Inference does not run global guardrails or policies, so a virtual key would get past + // them; refuse to start rather than serve keys the operator expects to be guarded. + if !config.guardrails.is_empty() + || !config.policies.is_empty() + || !config.policy_attachments.is_empty() + { + return Err(Error::GuardedVirtualKeys); + } + Ok(Arc::new(PostgresKeys { pool, personal_keys_use_teams: personal_keys_use_teams(config), - }))) + })) } /// `litellm_settings.personal_key_model_access_from_teams`. diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 55fd19bb941..7e5e50fcab1 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -8,6 +8,7 @@ mod key_status; mod keys; pub mod metrics; mod secrets; +pub mod spend; pub use error::Error; use std::{sync::Arc, time::Instant}; @@ -135,7 +136,26 @@ pub async fn build_keys( config: &Config, secrets: &dyn litellm_secrets::source::SecretSource, ) -> Result>, Error> { - keys::store(config, secrets).await + build_database(config, secrets) + .await? + .map(|pool| build_keys_on(config, pool)) + .transpose() +} + +/// The proxy's database, if one is configured (see [`keys::database`]). +pub async fn build_database( + config: &Config, + secrets: &dyn litellm_secrets::source::SecretSource, +) -> Result, Error> { + keys::database(config, secrets).await +} + +/// The virtual key store over an already built database pool. +pub fn build_keys_on( + config: &Config, + pool: litellm_db::PgPool, +) -> Result, Error> { + keys::store(config, pool) } pub fn router( diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index ed4f293b4de..fccfa6cfc66 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -91,8 +91,11 @@ async fn main() -> Result<(), Box> { }, ) }); - let keys = litellm_gateway::build_keys(&config, inference.secrets.as_ref()) - .await? + let database = litellm_gateway::build_database(&config, inference.secrets.as_ref()).await?; + let keys = database + .clone() + .map(|pool| litellm_gateway::build_keys_on(&config, pool)) + .transpose()? .map(|keys| litellm_gateway::key_cache(keys, &config)); let virtual_keys = keys.is_some(); let coordination = @@ -109,7 +112,13 @@ async fn main() -> Result<(), Box> { coordination.start(inference.clone(), keys.clone(), shutdown.clone()); } // Built-in hooks and observers are added here as the gateway gains them. - litellm_gateway::hooks::install(&inference, &config, Vec::new(), Vec::new()).await?; + 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)); + } + litellm_gateway::hooks::install(&inference, &config, Vec::new(), observers).await?; let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; let address = listener.local_addr()?; // Without TCP_NODELAY, Nagle holds small SSE chunks until the client's delayed ACK (~40 ms per stream). @@ -120,7 +129,7 @@ async fn main() -> Result<(), Box> { }); let shared_state = coordination.is_some(); - tracing::info!(%address, models = config.model_list.len(), virtual_keys, shared_state, log_level = %level, "gateway listening"); + 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()); diff --git a/litellm-rust/crates/gateway/src/spend.rs b/litellm-rust/crates/gateway/src/spend.rs new file mode 100644 index 00000000000..fc373307bec --- /dev/null +++ b/litellm-rust/crates/gateway/src/spend.rs @@ -0,0 +1,708 @@ +//! The spend writer: an observer that records every request in the proxy's Postgres tables, as +//! the proxy's own spend writer does, so the proxy's dashboards, spend reports and budget checks +//! see the gateway's traffic. +//! +//! Requests are added up in memory and written in batches: spend logs every couple of seconds, +//! and the key, user, team, team membership, end user and daily activity rows every +//! `proxy_batch_write_at` seconds. Every write adds to what is stored, so replicas never +//! overwrite each other. + +use std::{ + collections::HashMap, + mem, + sync::{Arc, Mutex, Weak}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use futures_util::future::BoxFuture; +use litellm_config::Config; +use litellm_db::{ + DailySpend, DailyTable, EntitySpend, PgPool, SpendLog, TeamMemberSpend, add_end_user_spend, + add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, insert_spend_logs, + upsert_daily_spend, +}; +use litellm_gateway_hooks::{Facts, Observer, Record, Route, Usage}; +use serde_json::{Map, Value, json}; + +/// How often spend logs are written, as the proxy's spend-log queue monitor polls. +const LOG_INTERVAL: Duration = Duration::from_secs(2); +/// Rows per spend-log statement, the proxy's `SPEND_LOG_WRITE_BATCH_MAX_ROWS`. +const LOG_BATCH: usize = 100; +/// Rows per daily-table statement. +const DAILY_BATCH: usize = 100; +/// Spend logs held while the database is unreachable; older ones are dropped past this. +const MAX_PENDING_LOGS: usize = 10_000; +/// Retries of a write that failed on a deadlock, a lock timeout or a lost connection. +const RETRIES: u32 = 3; + +/// The keys of the proxy's spend-log `metadata`, every one present and null when unset. +const METADATA_KEYS: [&str; 49] = [ + "actor_agent_id", + "additional_usage_values", + "agent_execution_mode", + "applied_guardrails", + "attempted_fallbacks", + "attempted_retries", + "autorouter_baseline_observation", + "autorouter_savings", + "autorouter_savings_estimate", + "azure_spillover", + "batch_failed_requests", + "batch_models", + "batch_successful_requests", + "billing_agent_id", + "cold_storage_object_key", + "compression_savings", + "cost_breakdown", + "error_information", + "eval_information", + "guardrail_information", + "internal_call_origin", + "litellm_call_id", + "litellm_gateway_injected_cache", + "litellm_overhead_time_ms", + "litellm_roi_estimator", + "max_retries", + "mcp_tool_call_metadata", + "model_map_information", + "original_model_group", + "proxy_server_request", + "requester_ip_address", + "router_metadata", + "routing_decision", + "spend_logs_metadata", + "status", + "target_agent_id", + "usage_object", + "used_client_oauth_token", + "user_agent", + "user_api_key", + "user_api_key_alias", + "user_api_key_org_id", + "user_api_key_project_alias", + "user_api_key_project_id", + "user_api_key_team_alias", + "user_api_key_team_id", + "user_api_key_user_id", + "vector_store_request_metadata", + "verified_human_user_id", +]; + +/// The annotation hooks set to add to a request's `spend_logs_metadata`. +pub const SPEND_LOGS_METADATA: &str = "spend_logs_metadata"; + +#[derive(Clone, Debug, PartialEq)] +pub struct SpendSettings { + /// `general_settings.proxy_batch_write_at`, plus the proxy's 0-5 s of start-up jitter. + pub interval: Duration, + /// `general_settings.disable_spend_logs`: everything but the spend-log rows is written. + pub disable_spend_logs: bool, + /// `litellm_settings.disable_end_user_cost_tracking`. + pub disable_end_user_tracking: bool, +} + +impl SpendSettings { + /// The writer's settings, or `None` when `general_settings.disable_spend_updates` turns + /// spend tracking off. + pub fn from_config(config: &Config) -> Option { + let general = &config.general_settings.additional_fields; + if flag(general.get("disable_spend_updates")) { + return None; + } + let seconds = match general.get("proxy_batch_write_at") { + Some(litellm_config::Value::Number(seconds)) => seconds.as_f64(), + Some(litellm_config::Value::String(seconds)) => seconds.trim().parse().ok(), + _ => std::env::var("PROXY_BATCH_WRITE_AT") + .ok() + .and_then(|seconds| seconds.trim().parse().ok()), + } + .filter(|seconds: &f64| seconds.is_finite() && *seconds > 0.0) + .unwrap_or(10.0); + let jitter = f64::from(rand::random::() % 6); + Some(Self { + interval: Duration::from_secs_f64(seconds + jitter), + disable_spend_logs: flag(general.get("disable_spend_logs")), + disable_end_user_tracking: flag( + config + .litellm_settings + .additional_fields + .get("disable_end_user_cost_tracking"), + ), + }) + } +} + +fn flag(value: Option<&litellm_config::Value>) -> bool { + match value { + Some(litellm_config::Value::Bool(flag)) => *flag, + Some(litellm_config::Value::String(text)) => text.eq_ignore_ascii_case("true"), + _ => false, + } +} + +type DailyKey = (DailyTable, String, String, String, String, String, String); + +/// Spend added up since the last write. +#[derive(Default)] +struct Pending { + keys: HashMap, + users: HashMap, + teams: HashMap, + end_users: HashMap, + members: HashMap<(String, String), f64>, + daily: HashMap, +} + +impl Pending { + fn add_daily(&mut self, table: DailyTable, row: DailySpend) { + let key = ( + table, + row.entity_id.clone(), + row.date.clone(), + row.api_key.clone(), + row.model.clone(), + row.custom_llm_provider.clone(), + row.endpoint.clone(), + ); + match self.daily.get_mut(&key) { + Some(total) => { + total.prompt_tokens += row.prompt_tokens; + total.completion_tokens += row.completion_tokens; + total.api_requests += row.api_requests; + total.successful_requests += row.successful_requests; + total.failed_requests += row.failed_requests; + total.cache_read_input_tokens += row.cache_read_input_tokens; + total.cache_creation_input_tokens += row.cache_creation_input_tokens; + total.total_response_time_ms += row.total_response_time_ms; + total.timed_requests += row.timed_requests; + total.spend += row.spend; + } + None => { + self.daily.insert(key, row); + } + } + } +} + +pub struct SpendWriter { + pool: PgPool, + settings: SpendSettings, + pending: Mutex, + logs: Mutex>, + /// Serializes flushes, so the periodic one and the one at shutdown never interleave. + flushing: tokio::sync::Mutex<()>, +} + +impl SpendWriter { + /// Starts the writer and its timers. Must be called inside a tokio runtime. + pub fn start(pool: PgPool, settings: SpendSettings) -> Arc { + let writer = Arc::new(Self { + pool, + settings, + pending: Mutex::default(), + logs: Mutex::default(), + flushing: tokio::sync::Mutex::new(()), + }); + tokio::spawn(run(Arc::downgrade(&writer))); + writer + } + + fn add(&self, record: &Record) { + let request = &record.request; + let key = record.key(); + let cost = record.outcome.cost.unwrap_or(0.0); + let end_user = request + .end_user + .as_deref() + .filter(|_| !self.settings.disable_end_user_tracking); + let hash = key.map(|key| key.hash.as_str()).unwrap_or_default(); + let user_id = key.and_then(|key| key.user_id.as_deref()); + let team_id = key.and_then(|key| key.team_id.as_deref()); + if hash.is_empty() && user_id.is_none() && team_id.is_none() && end_user.is_none() { + return; + } + let log = spend_log(record, cost, end_user); + { + let mut pending = self.pending.lock().unwrap_or_else(|e| e.into_inner()); + if !hash.is_empty() { + *pending.keys.entry(hash.to_owned()).or_default() += cost; + } + if let Some(user_id) = user_id { + *pending.users.entry(user_id.to_owned()).or_default() += cost; + } + if let Some(team_id) = team_id { + *pending.teams.entry(team_id.to_owned()).or_default() += cost; + if let Some(user_id) = user_id { + *pending + .members + .entry((team_id.to_owned(), user_id.to_owned())) + .or_default() += cost; + } + } + if let Some(end_user) = end_user { + *pending.end_users.entry(end_user.to_owned()).or_default() += cost; + } + let row = daily_row(record, &log); + pending.add_daily( + DailyTable::User, + DailySpend { + entity_id: log.user.clone(), + ..row.clone() + }, + ); + pending.add_daily( + DailyTable::Team, + DailySpend { + entity_id: log.team_id.clone(), + ..row.clone() + }, + ); + if let Some(end_user) = end_user { + pending.add_daily( + DailyTable::EndUser, + DailySpend { + entity_id: end_user.to_owned(), + ..row + }, + ); + } + } + if !self.settings.disable_spend_logs { + let mut logs = self.logs.lock().unwrap_or_else(|e| e.into_inner()); + logs.push(log); + if logs.len() > MAX_PENDING_LOGS { + let excess = logs.len() - MAX_PENDING_LOGS; + logs.drain(..excess); + tracing::warn!(dropped = excess, "spend log queue full; dropped the oldest"); + } + } + } + + /// Writes the spend logs held so far. + pub async fn flush_logs(&self) { + let _flushing = self.flushing.lock().await; + self.write_logs().await; + } + + /// Writes everything held so far. + pub async fn flush_all(&self) { + let _flushing = self.flushing.lock().await; + self.write_logs().await; + self.write_spend().await; + } + + async fn write_logs(&self) { + let logs = mem::take(&mut *self.logs.lock().unwrap_or_else(|e| e.into_inner())); + let mut failed = Vec::new(); + for batch in logs.chunks(LOG_BATCH) { + if let Err(error) = retry(|| insert_spend_logs(&self.pool, batch)).await { + tracing::error!(%error, rows = batch.len(), "writing spend logs failed"); + if error.is_retryable() { + failed.extend_from_slice(batch); + } + } + } + if !failed.is_empty() { + let mut logs = self.logs.lock().unwrap_or_else(|e| e.into_inner()); + failed.append(&mut logs); + *logs = failed; + } + } + + async fn write_spend(&self) { + let pending = mem::take(&mut *self.pending.lock().unwrap_or_else(|e| e.into_inner())); + let entity = |spend: HashMap| spend.into_iter().collect::(); + // The proxy's order: user, end user, key, team, team membership, then the daily tables. + // Entity spend that cannot be written is dropped, as the proxy drops it. + let users = entity(pending.users); + log_failure("user", retry(|| add_user_spend(&self.pool, &users)).await); + let end_users = entity(pending.end_users); + log_failure( + "end user", + retry(|| add_end_user_spend(&self.pool, &end_users)).await, + ); + let keys = entity(pending.keys); + log_failure("key", retry(|| add_key_spend(&self.pool, &keys)).await); + let teams = entity(pending.teams); + log_failure("team", retry(|| add_team_spend(&self.pool, &teams)).await); + let members: Vec<_> = pending + .members + .into_iter() + .map(|((team_id, user_id), cost)| TeamMemberSpend { + team_id, + user_id, + cost, + }) + .collect(); + log_failure( + "team member", + retry(|| add_team_member_spend(&self.pool, &members)).await, + ); + + let mut tables: HashMap> = HashMap::new(); + for ((table, ..), row) in pending.daily { + tables.entry(table).or_default().push(row); + } + let mut requeue = Vec::new(); + for table in [DailyTable::User, DailyTable::Team, DailyTable::EndUser] { + let rows = tables.remove(&table).unwrap_or_default(); + for batch in rows.chunks(DAILY_BATCH) { + if let Err(error) = retry(|| upsert_daily_spend(&self.pool, table, batch)).await { + tracing::error!(%error, ?table, rows = batch.len(), "writing daily spend failed"); + if error.is_retryable() { + requeue.extend(batch.iter().cloned().map(|row| (table, row))); + } + } + } + } + if !requeue.is_empty() { + let mut pending = self.pending.lock().unwrap_or_else(|e| e.into_inner()); + for (table, row) in requeue { + pending.add_daily(table, row); + } + } + } +} + +impl Observer for SpendWriter { + fn name(&self) -> &'static str { + "spend_writer" + } + + fn reads(&self) -> Facts { + Facts::KEY + | Facts::END_USER + | Facts::CLIENT_HEADERS + | Facts::DEPLOYMENT + | Facts::USAGE + | Facts::COST + | Facts::TIMINGS + | Facts::ANNOTATIONS + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + self.add(record); + Box::pin(async {}) + } + + fn flush(&self) -> BoxFuture<'_, ()> { + Box::pin(self.flush_all()) + } +} + +/// Writes spend logs every [`LOG_INTERVAL`] and everything else every batch interval, until the +/// writer is dropped. +async fn run(writer: Weak) { + let interval = match writer.upgrade() { + Some(writer) => writer.settings.interval, + None => return, + }; + let mut since_spend = Duration::ZERO; + loop { + tokio::time::sleep(LOG_INTERVAL).await; + let Some(writer) = writer.upgrade() else { + return; + }; + since_spend += LOG_INTERVAL; + if since_spend >= interval { + since_spend = Duration::ZERO; + writer.flush_all().await; + } else { + writer.flush_logs().await; + } + } +} + +async fn retry(mut write: F) -> Result<(), litellm_db::Error> +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempt = 0; + loop { + match write().await { + Err(error) if error.is_retryable() && attempt < RETRIES => { + let base = 1u64 << (attempt + 1); + let jitter = u64::from(rand::random::()) * base / u64::from(u16::MAX); + tokio::time::sleep(Duration::from_millis((base * 1000 + jitter * 1000) / 2)).await; + attempt += 1; + } + result => return result, + } + } +} + +fn log_failure(table: &str, result: Result<(), litellm_db::Error>) { + if let Err(error) = result { + tracing::error!(%error, table, "writing spend failed; the spend is dropped"); + } +} + +fn status(record: &Record) -> &'static str { + match (200..300).contains(&record.outcome.status) { + true => "success", + false => "failure", + } +} + +/// The model as the proxy logs it: the deployment's model without its provider route. +fn logged_model(record: &Record) -> String { + match &record.outcome.deployment { + Some(deployment) => match (deployment.provider(), deployment.model.split_once('/')) { + (Some(provider), Some((prefix, rest))) if prefix == provider => rest.to_owned(), + _ => deployment.model.clone(), + }, + None => record.request.requested_model.to_string(), + } +} + +fn spend_log(record: &Record, cost: f64, end_user: Option<&str>) -> SpendLog { + let request = &record.request; + let outcome = &record.outcome; + let key = record.key(); + let deployment = outcome.deployment.as_ref(); + let usage = outcome.usage.unwrap_or_default(); + let duration = record.duration(); + let completion_start = match (request.stream, outcome.time_to_first_byte) { + (true, Some(first_byte)) => request.started_at + first_byte, + _ => outcome.ended_at, + }; + let tokens = |count: u64| i32::try_from(count).unwrap_or(i32::MAX); + SpendLog { + request_id: request.request_id.to_string(), + call_type: request.route.call_type().into(), + api_key: key.map(|key| key.hash.clone()).unwrap_or_default(), + spend: cost, + total_tokens: tokens(usage.total_tokens()), + prompt_tokens: tokens(usage.prompt_tokens()), + completion_tokens: tokens(usage.output_tokens), + start_time: timestamp(request.started_at), + end_time: timestamp(outcome.ended_at), + completion_start_time: Some(timestamp(completion_start)), + request_duration_ms: Some(i32::try_from(duration.as_millis()).unwrap_or(i32::MAX)), + model: logged_model(record), + model_id: deployment.map(|d| d.id.to_string()).unwrap_or_default(), + model_group: deployment + .map(|d| d.model_group.to_string()) + .unwrap_or_else(|| request.requested_model.to_string()), + custom_llm_provider: deployment + .and_then(|d| d.provider()) + .unwrap_or_default() + .to_owned(), + api_base: deployment + .and_then(|d| d.api_base.clone()) + .unwrap_or_default(), + user: key.and_then(|key| key.user_id.clone()).unwrap_or_default(), + metadata: metadata(record), + cache_hit: if request.stream { "False" } else { "None" }.into(), + cache_key: "Cache OFF".into(), + team_id: key.and_then(|key| key.team_id.clone()).unwrap_or_default(), + organization_id: key.and_then(|key| key.org_id.clone()).unwrap_or_default(), + end_user: end_user.unwrap_or_default().to_owned(), + requester_ip_address: None, + session_id: Some(request.request_id.to_string()), + status: status(record).into(), + litellm_call_id: Some(request.request_id.to_string()), + } +} + +fn daily_row(record: &Record, log: &SpendLog) -> DailySpend { + let usage = record.outcome.usage.unwrap_or_default(); + let success = status(record) == "success"; + let count = |count: u64| i64::try_from(count).unwrap_or(i64::MAX); + DailySpend { + entity_id: String::new(), + date: date(record.request.started_at), + api_key: log.api_key.clone(), + model: log.model.clone(), + custom_llm_provider: log.custom_llm_provider.clone(), + // The proxy maps only chat completions to an endpoint here. + endpoint: match record.request.route { + Route::ChatCompletions => "/chat/completions".into(), + Route::Messages => String::new(), + }, + model_group: Some(log.model_group.clone()), + prompt_tokens: count(usage.prompt_tokens()), + completion_tokens: count(usage.output_tokens), + api_requests: 1, + successful_requests: i64::from(success), + failed_requests: i64::from(!success), + cache_read_input_tokens: count(usage.cache_read_tokens), + cache_creation_input_tokens: count(usage.cache_write_tokens), + total_response_time_ms: match success { + true => i64::from(log.request_duration_ms.unwrap_or(0)), + false => 0, + }, + timed_requests: i64::from(success), + spend: log.spend, + } +} + +/// The spend log's `metadata`, with every key the proxy writes. +fn metadata(record: &Record) -> 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 text = + |value: Option<&String>| value.map_or(Value::Null, |value| Value::from(value.as_str())); + let request = &record.request; + if let Some(key) = record.key() { + set("user_api_key", Value::from(key.hash.as_str())); + set("user_api_key_alias", text(key.alias.as_ref())); + set("user_api_key_team_id", text(key.team_id.as_ref())); + set("user_api_key_user_id", text(key.user_id.as_ref())); + set("user_api_key_team_alias", text(key.team_alias.as_ref())); + set("user_api_key_org_id", text(key.org_id.as_ref())); + } + set("spend_logs_metadata", spend_logs_metadata(record)); + set( + "user_agent", + request + .client_headers + .get(http::header::USER_AGENT) + .and_then(|value| value.to_str().ok()) + .map_or(Value::Null, Value::from), + ); + set("litellm_call_id", Value::from(&*request.request_id)); + set("litellm_roi_estimator", Value::Bool(false)); + if let Some(usage) = record.outcome.usage { + let (usage_object, additional) = usage_values(&usage); + set("usage_object", usage_object); + set("additional_usage_values", additional); + } + if status(record) == "failure" { + set("status", Value::from("failure")); + set( + "error_information", + json!({ + "error_code": record.outcome.status.to_string(), + "error_class": "", + "llm_provider": record + .outcome + .deployment + .as_ref() + .and_then(|deployment| deployment.provider()) + .unwrap_or_default(), + "traceback": "", + "error_message": "", + }), + ); + } + Value::Object(metadata) +} + +/// `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 { + let mut merged = Map::new(); + if let Some(key) = record.key() { + for metadata in [&key.team_metadata, &key.metadata] { + if let Some(Value::Object(fields)) = metadata.get("spend_logs_metadata") { + merged.extend(fields.clone()); + } + } + } + if let Some(Value::Object(fields)) = record.annotations.get(SPEND_LOGS_METADATA) { + merged.extend(fields.clone()); + } + match merged.is_empty() { + true => Value::Null, + false => Value::Object(merged), + } +} + +/// 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) { + let mut additional = Map::new(); + additional.insert( + "prompt_tokens_details".into(), + json!({ + "cached_tokens": usage.cache_read_tokens, + "cache_creation_tokens": usage.cache_write_tokens, + }), + ); + additional.insert( + "cache_read_input_tokens".into(), + Value::from(usage.cache_read_tokens), + ); + additional.insert( + "cache_creation_input_tokens".into(), + Value::from(usage.cache_write_tokens), + ); + if usage.web_search_requests > 0 { + additional.insert( + "server_tool_use".into(), + json!({"web_search_requests": usage.web_search_requests}), + ); + } + let mut object = additional.clone(); + object.insert("prompt_tokens".into(), Value::from(usage.prompt_tokens())); + object.insert("completion_tokens".into(), Value::from(usage.output_tokens)); + object.insert("total_tokens".into(), Value::from(usage.total_tokens())); + (Value::Object(object), Value::Object(additional)) +} + +/// Days since 1970-01-01 as a UTC calendar date. +fn civil(days: i64) -> (i64, u32, u32) { + let days = days + 719_468; + let era = days.div_euclid(146_097); + let day_of_era = days.rem_euclid(146_097); + let year_of_era = + (day_of_era - day_of_era / 1460 + day_of_era / 36_524 - day_of_era / 146_096) / 365; + let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100); + let month_index = (5 * day_of_year + 2) / 153; + let day = (day_of_year - (153 * month_index + 2) / 5 + 1) as u32; + let month = if month_index < 10 { + month_index + 3 + } else { + month_index - 9 + } as u32; + let year = year_of_era + era * 400 + i64::from(month <= 2); + (year, month, day) +} + +fn since_epoch(time: SystemTime) -> Duration { + time.duration_since(UNIX_EPOCH).unwrap_or_default() +} + +/// `YYYY-MM-DD`, UTC. +fn date(time: SystemTime) -> String { + let (year, month, day) = civil((since_epoch(time).as_secs() / 86_400) as i64); + format!("{year:04}-{month:02}-{day:02}") +} + +/// `YYYY-MM-DD HH:MM:SS.fff`, UTC, as the proxy stores `timestamp(3)` columns. +fn timestamp(time: SystemTime) -> String { + let elapsed = since_epoch(time); + let seconds = elapsed.as_secs() % 86_400; + format!( + "{} {:02}:{:02}:{:02}.{:03}", + date(time), + seconds / 3600, + seconds / 60 % 60, + seconds % 60, + elapsed.subsec_millis() + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::epoch(0, "1970-01-01 00:00:00.000")] + #[case::leap_day(1_709_208_000_123, "2024-02-29 12:00:00.123")] + #[case::year_end(1_798_761_599_999, "2026-12-31 23:59:59.999")] + fn timestamps_are_utc(#[case] millis: u64, #[case] expected: &str) { + let time = UNIX_EPOCH + Duration::from_millis(millis); + assert_eq!(timestamp(time), expected); + assert_eq!(date(time), expected[..10]); + } +} diff --git a/litellm-rust/crates/gateway/tests/spend.rs b/litellm-rust/crates/gateway/tests/spend.rs new file mode 100644 index 00000000000..ce08d645a05 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/spend.rs @@ -0,0 +1,304 @@ +//! The spend writer against a real proxy database: requests through the gateway end up in the +//! proxy's spend logs, entity spend and daily activity tables. + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_db::PgPool; +use litellm_db_testing::MigratedPostgres; +use litellm_gateway::spend::{SpendSettings, SpendWriter}; +use litellm_gateway_auth::keys::KeyHash; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"cache_read_input_tokens":2,"output_tokens":4}}"#; +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}}"#; + +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 payload = match request.uri().path().ends_with("/chat/completions") { + true => COMPLETION, + false => MESSAGE, + }; + ([(CONTENT_TYPE, "application/json")], payload).into_response() + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + format!("http://{address}") +} + +fn hash(token: &str) -> String { + KeyHash::from_token(&SecretValue::new(token)) + .as_str() + .to_owned() +} + +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows; the writes under test use the checked macros" +)] +async fn fixtures(pool: &PgPool) { + let sql = format!( + r#" +INSERT INTO "LiteLLM_TeamTable" (team_id, team_alias, members_with_roles, metadata, created_at, updated_at) + VALUES ('team-1', 'research', '[{{"user_id": "user-1", "role": "user"}}]', + '{{"spend_logs_metadata": {{"project": "team", "cost_center": "42"}}}}', now(), now()); +INSERT INTO "LiteLLM_UserTable" (user_id, spend) VALUES ('user-1', 0); +INSERT INTO "LiteLLM_VerificationToken" (token, key_alias, spend, team_id, user_id, metadata) + VALUES ('{}', 'team-key', 0, 'team-1', 'user-1', + '{{"spend_logs_metadata": {{"project": "alpha"}}}}'); +"#, + hash("sk-team") + ); + sqlx::raw_sql(sqlx::AssertSqlSafe(sql)) + .execute(pool) + .await + .unwrap(); +} + +#[expect( + clippy::disallowed_methods, + reason = "test assertions read rows back; the writes under test use the checked macros" +)] +async fn rows(pool: &PgPool, sql: &str) -> Vec { + sqlx::query_scalar::<_, Value>(sqlx::AssertSqlSafe(format!( + "SELECT to_jsonb(r) FROM ({sql}) r" + ))) + .fetch_all(pool) + .await + .unwrap() +} + +async fn send(app: &Router, path: &str, key: &str, body: Value) { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {key}")) + .header("user-agent", "spend-test/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:?}"); +} + +fn close(left: &Value, right: f64) -> bool { + (left.as_f64().unwrap() - right).abs() < 1e-9 +} + +#[rstest] +#[tokio::test] +async fn requests_are_written_as_the_proxy_writes_them() { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + fixtures(pool).await; + let upstream = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream + input_cost_per_token: 0.001 + output_cost_per_token: 0.002 + cache_read_input_token_cost: 0.0001 + model_info: + id: claude-1 + - model_name: gpt + litellm_params: + model: openai/gpt-x + api_base: {upstream}/v1 + api_key: sk-upstream + input_cost_per_token: 1 + output_cost_per_token: 2 + model_info: + id: gpt-1 +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master + database_url: '{}' +", + postgres.url() + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let database = litellm_gateway::build_database(&config, inference.secrets.as_ref()) + .await + .unwrap() + .unwrap(); + let keys = litellm_gateway::build_keys_on(&config, database.clone()).unwrap(); + let settings = SpendSettings::from_config(&config).unwrap(); + litellm_gateway::hooks::install( + &inference, + &config, + Vec::new(), + vec![SpendWriter::start(database, settings)], + ) + .await + .unwrap(); + let app = litellm_gateway::router_with_keys(inference.clone(), &config, None, None, Some(keys)); + + let message = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"user_id": "customer-7"}, + }); + send(&app, "/v1/messages", "sk-team", message.clone()).await; + send(&app, "/v1/messages", "sk-team", message).await; + let chat = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]}); + send(&app, "/v1/chat/completions", "sk-master", chat).await; + litellm_gateway::hooks::close(&inference).await; + + let message_cost = 3.0 * 0.001 + 2.0 * 0.0001 + 4.0 * 0.002; + let logs = rows( + pool, + r#"SELECT call_type, api_key, spend, prompt_tokens, completion_tokens, total_tokens, + model, model_id, model_group, custom_llm_provider, "user", team_id, end_user, + cache_hit, cache_key, status, metadata + FROM "LiteLLM_SpendLogs" ORDER BY call_type DESC, "startTime""#, + ) + .await; + assert_eq!(logs.len(), 3, "{logs:#?}"); + let [message_log, _, chat_log] = logs.as_slice() else { + unreachable!() + }; + assert!( + close(&message_log["spend"], message_cost), + "{message_log:#}" + ); + let mut message_log = message_log.clone(); + let metadata = message_log + .as_object_mut() + .unwrap() + .remove("metadata") + .unwrap(); + message_log.as_object_mut().unwrap().remove("spend"); + assert_eq!( + message_log, + json!({ + "call_type": "anthropic_messages", + "api_key": hash("sk-team"), + "prompt_tokens": 5, + "completion_tokens": 4, + "total_tokens": 9, + "model": "claude-x", + "model_id": "claude-1", + "model_group": "claude", + "custom_llm_provider": "anthropic", + "user": "user-1", + "team_id": "team-1", + "end_user": "customer-7", + "cache_hit": "None", + "cache_key": "Cache OFF", + "status": "success", + }) + ); + assert_eq!(metadata.as_object().unwrap().len(), 49, "{metadata:#}"); + assert_eq!(metadata["user_api_key_alias"], "team-key"); + assert_eq!(metadata["user_api_key_team_alias"], "research"); + assert_eq!(metadata["user_agent"], "spend-test/1.0"); + assert_eq!(metadata["status"], Value::Null); + assert_eq!( + metadata["spend_logs_metadata"], + json!({"project": "alpha", "cost_center": "42"}) + ); + assert_eq!(chat_log["call_type"], "acompletion"); + assert_eq!(chat_log["user"], "default_user_id"); + assert_eq!(chat_log["spend"], 17.0); + + let key = rows( + pool, + &format!( + r#"SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = '{}'"#, + hash("sk-team") + ), + ) + .await; + assert!(close(&key[0]["spend"], 2.0 * message_cost), "{key:?}"); + for sql in [ + r#"SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = 'user-1'"#, + r#"SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = 'team-1'"#, + r#"SELECT spend FROM "LiteLLM_TeamMembership" WHERE team_id = 'team-1'"#, + r#"SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id = 'customer-7'"#, + ] { + let row = rows(pool, sql).await; + assert!( + close(&row[0]["spend"], 2.0 * message_cost), + "{sql}: {row:?}" + ); + } + + let daily = rows( + pool, + r#"SELECT user_id, api_key, model, model_group, endpoint, api_requests, + successful_requests, prompt_tokens, completion_tokens, cache_read_input_tokens + FROM "LiteLLM_DailyUserSpend" ORDER BY user_id"#, + ) + .await; + assert_eq!( + daily, + [ + json!({ + "user_id": "default_user_id", + "api_key": hash("sk-master"), + "model": "gpt-x", + "model_group": "gpt", + "endpoint": "/chat/completions", + "api_requests": 1, + "successful_requests": 1, + "prompt_tokens": 5, + "completion_tokens": 6, + "cache_read_input_tokens": 0, + }), + json!({ + "user_id": "user-1", + "api_key": hash("sk-team"), + "model": "claude-x", + "model_group": "claude", + "endpoint": "", + "api_requests": 2, + "successful_requests": 2, + "prompt_tokens": 10, + "completion_tokens": 8, + "cache_read_input_tokens": 4, + }), + ] + ); + let teams = rows( + pool, + r#"SELECT team_id, api_requests FROM "LiteLLM_DailyTeamSpend" ORDER BY team_id"#, + ) + .await; + assert_eq!( + teams, + [ + json!({"team_id": "", "api_requests": 1}), + json!({"team_id": "team-1", "api_requests": 2}), + ] + ); + let end_users = rows( + pool, + r#"SELECT end_user_id, api_requests FROM "LiteLLM_DailyEndUserSpend""#, + ) + .await; + assert_eq!( + end_users, + [json!({"end_user_id": "customer-7", "api_requests": 2})] + ); +} diff --git a/litellm-rust/docs/customer-proxy/spend-writer.svg b/litellm-rust/docs/customer-proxy/spend-writer.svg new file mode 100644 index 00000000000..53bb0e03fdf --- /dev/null +++ b/litellm-rust/docs/customer-proxy/spend-writer.svg @@ -0,0 +1,52 @@ + + + + + + Spend writer: gateway requests land in the proxy's own tables + + + + Request record (observer queue) + key, user, team, end user, deployment, + usage, cost, timings, user agent, + spend_logs_metadata (team, key, hooks) + off the request path; never blocks it + + + + Added up in memory + spend per key, user, team, member, + end user; daily rows per entity, day, + key, model, provider, endpoint; + spend log rows queued (cap 10k) + + + + + Timers (as the proxy) + spend logs: every 2 s + everything else: proxy_batch_write_at + (default 10 s) plus 0-5 s jitter + shutdown: drain the queue, then flush + + + + + Proxy Postgres, same statements and order as the Python writer + LiteLLM_SpendLogs: batches of 100, INSERT ... ON CONFLICT DO NOTHING; columns and 49-key metadata as the proxy writes them + user, end user, key (spend, total_spend, last_active), team, team membership: spend += cost, ids sorted, lock_timeout 5 s + update-only for users, keys and teams; end users created unblocked; membership only for users on the team roster + LiteLLM_DailyUserSpend / DailyTeamSpend / DailyEndUserSpend: upsert that adds tokens, requests, spend; model_group kept from first insert + every write adds to what is stored, so gateway replicas and Python proxies can write side by side + deadlock, lock timeout, lost connection: 3 retries with backoff; then daily rows and spend logs are kept for the next flush, + entity spend is dropped (as the proxy drops it) + + + + + Read by what already reads these tables + proxy UI usage pages, /spend and /global/activity endpoints, key and team spend shown in the admin UI, + and the proxy's own budget checks on its next key lookup + + From b83353906b07553e2b490a1db9da3e0120456d69 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:44 -0700 Subject: [PATCH 03/10] feat(gateway): enforce key, team, user, member and end-user budgets on shared spend counters (#47044) Adds a built-in budgets hook that runs the proxy's budget checks, in its order and with its messages, on the proxy's Redis spend:* counters. Each request reads and seeds the counters, checks them and reserves its maximum cost in one Redis script, then settles the reservation to its real cost off the request path, so replicas and Python proxies sharing the Redis enforce one budget. Without Redis the counters are kept per replica. The key store now reads user, team member and end-user budgets, and auth no longer checks budgets against stored spend. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/Cargo.lock | 1 + ...399bc28aaeab7ee4da8e53f2ce87a52a3da9.json} | 90 +- ...f92fc7030de81c903adb593a5eca47237cfbe.json | 60 + litellm-rust/crates/db/src/end_user.rs | 40 + litellm-rust/crates/db/src/lib.rs | 2 + .../crates/db/src/verification_token.rs | 30 +- .../crates/db/tests/verification_token.rs | 76 +- litellm-rust/crates/gateway-auth/src/keys.rs | 38 + litellm-rust/crates/gateway-hooks/Cargo.toml | 1 + .../crates/gateway-hooks/src/facts.rs | 3 + litellm-rust/crates/gateway-hooks/src/hook.rs | 17 +- .../crates/gateway-hooks/src/pipeline.rs | 1 + .../crates/gateway-hooks/src/pricing.rs | 9 + .../crates/gateway-inference/src/error.rs | 2 +- .../crates/gateway-inference/src/hooks.rs | 1 + litellm-rust/crates/gateway/src/budgets.rs | 1209 +++++++++++++++++ .../crates/gateway/src/coordination.rs | 8 +- litellm-rust/crates/gateway/src/key_status.rs | 221 +-- litellm-rust/crates/gateway/src/lib.rs | 1 + litellm-rust/crates/gateway/src/main.rs | 11 +- litellm-rust/crates/gateway/tests/budgets.rs | 495 +++++++ litellm-rust/docs/customer-proxy/budgets.svg | 51 + 22 files changed, 2159 insertions(+), 208 deletions(-) rename litellm-rust/crates/db/.sqlx/{query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json => query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json} (56%) create mode 100644 litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json create mode 100644 litellm-rust/crates/db/src/end_user.rs create mode 100644 litellm-rust/crates/gateway/src/budgets.rs create mode 100644 litellm-rust/crates/gateway/tests/budgets.rs create mode 100644 litellm-rust/docs/customer-proxy/budgets.svg diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5da517b88b3..20bb696c46c 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3820,6 +3820,7 @@ dependencies = [ "futures-util", "http 1.4.2", "litellm-cost", + "litellm-gateway-auth", "litellm-model-catalog", "memchr", "rstest", diff --git a/litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json b/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json similarity index 56% rename from litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json rename to litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json index 0283a5854a0..1a9ad366cfa 100644 --- a/litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json +++ b/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n v.max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL\n OR v.max_parallel_requests IS NOT NULL\n OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL\n OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL\n OR t.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL\n OR mb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL\n OR u.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", + "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n COALESCE(v.max_budget, kb.max_budget) AS max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL\n OR v.max_parallel_requests IS NOT NULL\n OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL\n OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL\n OR t.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL\n OR mb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n tm.spend AS \"member_spend?\",\n mb.max_budget AS \"member_max_budget?\",\n CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC')\n THEN mb.temp_budget_increase END AS \"member_budget_increase?\",\n tdb.max_budget AS \"member_default_budget?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n u.spend AS \"user_spend?\",\n u.max_budget AS \"user_max_budget?\",\n (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL\n OR u.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS tdb\n ON tdb.budget_id = t.metadata ->> 'team_member_budget_id'\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", "describe": { "columns": [ { @@ -57,12 +57,7 @@ "ordinal": 5, "name": "max_budget", "type_info": "Float8", - "origin": { - "Table": { - "table": "\"LiteLLM_VerificationToken\"", - "name": "max_budget" - } - } + "origin": "Expression" }, { "ordinal": 6, @@ -293,6 +288,45 @@ }, { "ordinal": 33, + "name": "member_spend?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamMembership\"", + "name": "spend" + } + } + }, + { + "ordinal": 34, + "name": "member_max_budget?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "max_budget" + } + } + }, + { + "ordinal": 35, + "name": "member_budget_increase?", + "type_info": "Float8", + "origin": "Expression" + }, + { + "ordinal": 36, + "name": "member_default_budget?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "max_budget" + } + } + }, + { + "ordinal": 37, "name": "user_id", "type_info": "Text", "origin": { @@ -303,7 +337,7 @@ } }, { - "ordinal": 34, + "ordinal": 38, "name": "user_row_id?", "type_info": "Text", "origin": { @@ -314,7 +348,7 @@ } }, { - "ordinal": 35, + "ordinal": 39, "name": "user_models?", "type_info": "TextArray", "origin": { @@ -325,19 +359,41 @@ } }, { - "ordinal": 36, + "ordinal": 40, + "name": "user_spend?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "spend" + } + } + }, + { + "ordinal": 41, + "name": "user_max_budget?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "max_budget" + } + } + }, + { + "ordinal": 42, "name": "user_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 37, + "ordinal": 43, "name": "user_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 38, + "ordinal": 44, "name": "user_object_permissions!", "type_info": "Bool", "origin": "Expression" @@ -354,7 +410,7 @@ true, true, false, - true, + null, true, true, true, @@ -384,11 +440,17 @@ true, true, true, + null, + true, + true, + true, + true, + true, true, null, null, null ] }, - "hash": "aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d" + "hash": "08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9" } diff --git a/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json b/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json new file mode 100644 index 00000000000..c9805198db5 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json @@ -0,0 +1,60 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n (SELECT d.max_budget FROM \"LiteLLM_BudgetTable\" AS d\n WHERE d.budget_id = $2) AS \"default_max_budget?\"\n FROM (SELECT 1) AS one\n LEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\n LEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = e.budget_id", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_EndUserTable\"", + "name": "spend" + } + } + }, + { + "ordinal": 1, + "name": "blocked?", + "type_info": "Bool", + "origin": { + "Table": { + "table": "\"LiteLLM_EndUserTable\"", + "name": "blocked" + } + } + }, + { + "ordinal": 2, + "name": "max_budget?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "max_budget" + } + } + }, + { + "ordinal": 3, + "name": "default_max_budget?", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text" + ] + }, + "nullable": [ + true, + true, + true, + null + ] + }, + "hash": "a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe" +} diff --git a/litellm-rust/crates/db/src/end_user.rs b/litellm-rust/crates/db/src/end_user.rs new file mode 100644 index 00000000000..a4d98eef2aa --- /dev/null +++ b/litellm-rust/crates/db/src/end_user.rs @@ -0,0 +1,40 @@ +use sqlx::PgPool; + +use crate::Error; + +/// An end user's spend and budget, as the proxy's end-user budget checks read them. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct EndUser { + /// `LiteLLM_EndUserTable.spend`, when the end user has a row. + pub spend: Option, + pub blocked: bool, + /// The row's budget, else the default budget's. + pub max_budget: Option, +} + +/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own. +pub async fn find_end_user( + pool: &PgPool, + end_user_id: &str, + default_budget_id: Option<&str>, +) -> Result { + let row = sqlx::query!( + r#"SELECT e.spend AS "spend?", + e.blocked AS "blocked?", + b.max_budget AS "max_budget?", + (SELECT d.max_budget FROM "LiteLLM_BudgetTable" AS d + WHERE d.budget_id = $2) AS "default_max_budget?" + FROM (SELECT 1) AS one + LEFT JOIN "LiteLLM_EndUserTable" AS e ON e.user_id = $1 + LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = e.budget_id"#, + end_user_id, + default_budget_id, + ) + .fetch_one(pool) + .await?; + Ok(EndUser { + spend: row.spend, + blocked: row.blocked.unwrap_or(false), + max_budget: row.max_budget.or(row.default_max_budget), + }) +} diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 39ef58241b2..b41ee5ace78 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -1,9 +1,11 @@ +mod end_user; mod error; mod failure; mod pool; mod spend; mod verification_token; +pub use end_user::{EndUser, find_end_user}; pub use error::Error; pub use failure::DbFailure; pub use pool::{PoolSettings, connect_lazy}; diff --git a/litellm-rust/crates/db/src/verification_token.rs b/litellm-rust/crates/db/src/verification_token.rs index aa48a881767..3ff008b0cff 100644 --- a/litellm-rust/crates/db/src/verification_token.rs +++ b/litellm-rust/crates/db/src/verification_token.rs @@ -41,9 +41,15 @@ pub struct Team { pub restrictions: Restrictions, } -/// What a team membership's budget table sets beyond its dollar budget. -#[derive(Clone, Debug, Default, Eq, PartialEq)] +/// The key's team membership: its spend and budget, and what its budget table sets beyond them. +#[derive(Clone, Debug, Default, PartialEq)] pub struct TeamMembership { + /// The membership row's spend, when the key's user has a row for its team. + pub spend: Option, + /// The membership's own budget, with any temporary increase still in effect. + pub max_budget: Option, + /// The team's default member budget, `metadata.team_member_budget_id`. + pub default_max_budget: Option, /// `rpm_limit`, `tpm_limit`, `tpd_limit` or `max_parallel_requests`. pub rate_limits: bool, /// The member's own model scope; empty inherits the team's. @@ -52,6 +58,8 @@ pub struct TeamMembership { #[derive(Clone, Debug, PartialEq)] pub struct User { + pub spend: f64, + pub max_budget: Option, pub models: Vec, pub restrictions: Restrictions, } @@ -85,7 +93,7 @@ pub async fn find_verification_token( v.blocked, v.models, v.spend, - v.max_budget, + COALESCE(v.max_budget, kb.max_budget) AS max_budget, v.allowed_routes, v.key_alias, v.key_name, @@ -131,9 +139,16 @@ pub async fn find_verification_token( '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "member_rate_limits!", mb.allowed_models AS "member_allowed_models?", + tm.spend AS "member_spend?", + mb.max_budget AS "member_max_budget?", + CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC') + THEN mb.temp_budget_increase END AS "member_budget_increase?", + tdb.max_budget AS "member_default_budget?", v.user_id, u.user_id AS "user_row_id?", u.models AS "user_models?", + u.spend AS "user_spend?", + u.max_budget AS "user_max_budget?", (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL OR u.max_parallel_requests IS NOT NULL OR COALESCE(jsonb_path_exists(u.model_max_budget, @@ -147,6 +162,8 @@ pub async fn find_verification_token( LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND v.user_id = tm.user_id LEFT JOIN "LiteLLM_BudgetTable" AS mb ON tm.budget_id = mb.budget_id + LEFT JOIN "LiteLLM_BudgetTable" AS tdb + ON tdb.budget_id = t.metadata ->> 'team_member_budget_id' LEFT JOIN "LiteLLM_UserTable" AS u ON v.user_id = u.user_id WHERE v.token = $1"#, token @@ -192,11 +209,18 @@ pub async fn find_verification_token( }, }), team_membership: TeamMembership { + spend: row.member_spend, + max_budget: row + .member_max_budget + .map(|budget| budget + row.member_budget_increase.unwrap_or(0.0)), + default_max_budget: row.member_default_budget, rate_limits: row.member_rate_limits, allowed_models: row.member_allowed_models.unwrap_or_default(), }, user_id: row.user_id, user: row.user_row_id.map(|_| User { + spend: row.user_spend.unwrap_or(0.0), + max_budget: row.user_max_budget, models: row.user_models.unwrap_or_default(), restrictions: Restrictions { rate_limits: row.user_rate_limits, diff --git a/litellm-rust/crates/db/tests/verification_token.rs b/litellm-rust/crates/db/tests/verification_token.rs index 169429f4839..ffe0b4d66f6 100644 --- a/litellm-rust/crates/db/tests/verification_token.rs +++ b/litellm-rust/crates/db/tests/verification_token.rs @@ -1,7 +1,8 @@ use std::time::{Duration, SystemTime}; use litellm_db::{ - Restrictions, Team, TeamMembership, User, VerificationToken, find_verification_token, + EndUser, Restrictions, Team, TeamMembership, User, VerificationToken, find_end_user, + find_verification_token, }; use litellm_db_testing::MigratedPostgres; use rstest::rstest; @@ -21,9 +22,12 @@ INSERT INTO "LiteLLM_BudgetTable" (budget_id, max_budget, rpm_limit, model_max_b VALUES ('b-1', 5.0, NULL, NULL, ARRAY['gpt-a'], 'x', 'x'), ('b-2', NULL, 3, NULL, ARRAY[]::text[], 'x', 'x'), ('b-3', NULL, NULL, '{"gpt-a": {"tpm_limit": 7}}', ARRAY[]::text[], 'x', 'x'); -INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, budget_id) - VALUES ('user-2', 'team-2', 'b-1'), - ('user-1', 'team-1', 'b-2'); +INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, budget_id, spend) + VALUES ('user-2', 'team-2', 'b-1', 0.75), + ('user-1', 'team-1', 'b-2', 0); +INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id, blocked) + VALUES ('end-1', 1.5, 'b-1', false), + ('end-2', 0.25, NULL, true); INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, created_by, updated_by) VALUES ('org-1', 'org', 'b-1', 'x', 'x'); INSERT INTO "LiteLLM_ProjectTable" (project_id, created_by, updated_by) VALUES ('project-1', 'x', 'x'); @@ -58,19 +62,20 @@ INSERT INTO "LiteLLM_VerificationToken" '{"gpt-a": {"budget_limit": 1.0, "rpm_limit": 2}}'); "#; +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows; the lookups under test use the checked macro" +)] +async fn setup(pool: &sqlx::PgPool) { + sqlx::raw_sql(FIXTURES).execute(pool).await.unwrap(); +} + #[rstest] #[tokio::test] async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { let postgres = MigratedPostgres::start().await.unwrap(); - #[expect( - clippy::disallowed_methods, - reason = "test fixture rows; the lookup under test uses the checked macro" - )] - sqlx::raw_sql(FIXTURES) - .execute(postgres.pool()) - .await - .unwrap(); let pool = postgres.pool(); + setup(pool).await; let find = |token: &'static str| async move { find_verification_token(pool, token).await.unwrap().unwrap() }; @@ -125,11 +130,16 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { }, }), team_membership: TeamMembership { + spend: Some(0.0), + max_budget: None, + default_max_budget: None, rate_limits: true, allowed_models: vec![], }, user_id: Some("user-1".into()), user: Some(User { + spend: 0.0, + max_budget: Some(1.0), models: vec!["gpt-b".into()], restrictions: Restrictions { rate_limits: true, @@ -175,10 +185,16 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { assert_eq!( member.team_membership, TeamMembership { + spend: Some(0.75), + max_budget: Some(5.0), + default_max_budget: None, rate_limits: false, allowed_models: vec!["gpt-a".into()], } ); + // A key without its own budget takes its budget row's. + assert_eq!(find("restricted").await.max_budget, Some(5.0)); + assert_eq!(find("budgeted").await.max_budget, Some(4.0)); assert_eq!(member.team.unwrap().restrictions, Restrictions::default()); assert_eq!(member.user.unwrap().restrictions, Restrictions::default()); let orphan = find("orphan").await; @@ -193,3 +209,39 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { None ); } + +#[rstest] +#[tokio::test] +async fn reads_end_user_budgets_with_their_default() { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + setup(pool).await; + let find = + async |id: &str, default: Option<&str>| find_end_user(pool, id, default).await.unwrap(); + + assert_eq!( + find("end-1", Some("b-2")).await, + EndUser { + spend: Some(1.5), + blocked: false, + max_budget: Some(5.0) + } + ); + assert_eq!( + find("end-2", Some("b-1")).await, + EndUser { + spend: Some(0.25), + blocked: true, + max_budget: Some(5.0) + } + ); + assert_eq!(find("end-2", None).await.max_budget, None); + assert_eq!( + find("unknown", Some("b-1")).await, + EndUser { + spend: None, + blocked: false, + max_budget: Some(5.0) + } + ); +} diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 2ce5338a674..ff99572a1ee 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -297,8 +297,46 @@ pub struct KeyProfile { pub team_alias: Option, pub metadata: Arc>, pub team_metadata: Arc>, + pub budgets: KeyBudgets, } +/// Stored spend against an optional budget, as the key store last read it. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct Spend { + pub spend: f64, + pub max_budget: Option, +} + +impl Eq for Spend {} + +/// The spend and budgets of a key and of who it belongs to, read with the key and refreshed +/// with it. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct KeyBudgets { + /// The key's `max_budget`, else its budget table's. + pub key: Spend, + /// `key_name`, the key's masked form, which budget errors show while it still looks masked. + pub key_name: Option, + /// The key's team, when its row exists. + pub team: Option, + /// The key's user, when its row exists. + pub user: Option, + /// The key's user within its team, when a membership row exists. + pub team_member: Option, + /// `metadata.end_user_budget_id`: the budget of end users without their own. + pub end_user_budget_id: Option, +} + +/// A team member's spend, with the membership's own budget and the team's default one. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct MemberSpend { + pub spend: f64, + pub max_budget: Option, + pub default_max_budget: Option, +} + +impl Eq for MemberSpend {} + #[derive(Clone, Debug, Eq, PartialEq)] pub enum KeyStatus { Active { diff --git a/litellm-rust/crates/gateway-hooks/Cargo.toml b/litellm-rust/crates/gateway-hooks/Cargo.toml index b0d6c67f15f..e072cd72c7a 100644 --- a/litellm-rust/crates/gateway-hooks/Cargo.toml +++ b/litellm-rust/crates/gateway-hooks/Cargo.toml @@ -10,6 +10,7 @@ bytes.workspace = true futures-util.workspace = true http.workspace = true litellm-cost.workspace = true +litellm-gateway-auth.workspace = true litellm-model-catalog.workspace = true memchr = "2" serde.workspace = true diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs index 3e7e2033e13..437768707d0 100644 --- a/litellm-rust/crates/gateway-hooks/src/facts.rs +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -8,6 +8,8 @@ use bytes::Bytes; use http::HeaderMap; use serde_json::{Map, Value}; +use litellm_gateway_auth::keys::KeyBudgets; + use crate::usage::Usage; /// The facts a hook reads or an observer asks for. The gateway builds only what some hook or @@ -101,6 +103,7 @@ pub struct KeyFacts { pub org_id: Option, pub metadata: Arc>, pub team_metadata: Arc>, + pub budgets: KeyBudgets, } /// What the gateway knows about a request before it is sent, shared by every hook and observer. diff --git a/litellm-rust/crates/gateway-hooks/src/hook.rs b/litellm-rust/crates/gateway-hooks/src/hook.rs index 0392cd0f88f..3e57e6b12c8 100644 --- a/litellm-rust/crates/gateway-hooks/src/hook.rs +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -12,7 +12,10 @@ use futures_util::future::BoxFuture; use http::{HeaderName, HeaderValue}; use serde_json::{Map, Value}; -use crate::facts::{Annotations, Facts, Record, RequestFacts}; +use crate::{ + facts::{Annotations, Facts, Record, RequestFacts}, + pricing::Prices, +}; /// The parts of a request a hook may change. #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -93,6 +96,8 @@ pub struct Refusal { pub status: u16, /// The proxy's error `type`, such as `budget_exceeded`. pub kind: Cow<'static, str>, + /// The error's `param`: the string `"None"` for most proxy errors, `null` for some. + pub param: Option<&'static str>, pub message: String, pub headers: Vec<(HeaderName, HeaderValue)>, } @@ -106,11 +111,18 @@ impl Refusal { Self { status, kind: kind.into(), + param: Some("None"), message: message.into(), headers: Vec::new(), } } + /// Sends a `null` `param`, as the proxy's budget errors do. + pub fn without_param(mut self) -> Self { + self.param = None; + self + } + pub fn with_header(mut self, name: HeaderName, value: HeaderValue) -> Self { self.headers.push((name, value)); self @@ -229,6 +241,9 @@ pub trait Hook: Send + Sync + 'static { /// The parts of the request this hook may change. fn writes(&self) -> Targets; + /// Called once with the deployments' prices before the first request. + fn prepare(&self, _prices: &Arc) {} + /// Runs before the request is routed. Returning a [`Settle`] asks to be told how the request /// ended. fn run<'a>( diff --git a/litellm-rust/crates/gateway-hooks/src/pipeline.rs b/litellm-rust/crates/gateway-hooks/src/pipeline.rs index 384babaf8d2..04d7076a85e 100644 --- a/litellm-rust/crates/gateway-hooks/src/pipeline.rs +++ b/litellm-rust/crates/gateway-hooks/src/pipeline.rs @@ -24,6 +24,7 @@ impl Pipeline { let mut reads = observers.reads(); for hook in &hooks { reads |= hook.reads(); + hook.prepare(&prices); } // Cost is computed from usage and the deployment that served the request. if reads.contains(Facts::COST) { diff --git a/litellm-rust/crates/gateway-hooks/src/pricing.rs b/litellm-rust/crates/gateway-hooks/src/pricing.rs index 813676c6f5b..be1bca921ba 100644 --- a/litellm-rust/crates/gateway-hooks/src/pricing.rs +++ b/litellm-rust/crates/gateway-hooks/src/pricing.rs @@ -17,12 +17,16 @@ use crate::usage::{ServiceTier, Usage}; #[derive(Debug, Default)] pub struct Prices { plans: HashMap, PricingPlan>, + max_output_tokens: HashMap, u64>, } impl Prices { /// Adds a deployment's prices from its cost map fields. A deployment whose fields hold no /// usable prices is left out, and its cost is unknown. pub fn insert(&mut self, id: Arc, fields: &Map) { + 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; @@ -43,6 +47,11 @@ impl Prices { self.plans.is_empty() } + /// The most output tokens the deployment's model can return, when the cost map says. + pub fn max_output_tokens(&self, deployment: &str) -> Option { + self.max_output_tokens.get(deployment).copied() + } + /// 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)?; diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index ce4b5aed0ac..05009604fba 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -180,7 +180,7 @@ fn refused(refusal: Refusal) -> Response { let body = json!({"error": { "message": refusal.message, "type": refusal.kind, - "param": "None", + "param": refusal.param, "code": status.as_u16().to_string(), }}); let mut response = (status, Json(body)).into_response(); diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index 408810bf095..e60bcd140bc 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -120,6 +120,7 @@ fn key_facts(hooks: &Hooks, identity: &AuthenticatedRequest) -> Option Some(Arc::new(KeyFacts { diff --git a/litellm-rust/crates/gateway/src/budgets.rs b/litellm-rust/crates/gateway/src/budgets.rs new file mode 100644 index 00000000000..be04ae028be --- /dev/null +++ b/litellm-rust/crates/gateway/src/budgets.rs @@ -0,0 +1,1209 @@ +//! The budgets hook: refuses a request when its key, team, user, team member or end user is over +//! budget, with the proxy's checks, order and messages, and keeps the proxy's `spend:*` counters +//! in Redis so gateway replicas and Python proxies sharing that Redis see each other's spend. +//! +//! Each counter is read with the stored spend as its floor: a missing counter is seeded from it +//! and one below it is raised to it, as the proxy reseeds and repairs counters. A request then +//! reserves its most expensive possible cost on every counter with a budget, and settles the +//! reservation to its real cost when it ends, so concurrent requests cannot all slip under one +//! budget. Reading, checking and reserving is one Redis script, atomic per request. Without +//! Redis, or while it is unreachable, the counters live in this replica's memory. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex, OnceLock}, + time::{Duration, Instant}, +}; + +use futures_util::future::BoxFuture; +use litellm_config::Config; +use litellm_db::{EndUser, PgPool}; +use litellm_gateway_hooks::{ + Facts, Hook, Prices, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets, Usage, +}; +use redis::aio::ConnectionManager; + +use crate::coordination::Coordination; + +/// The output tokens the proxy assumes when neither the request nor the model caps them. +const DEFAULT_OUTPUT_TOKENS: u64 = 16_384; +/// The proxy's slack when a reservation exactly meets a budget. +const EPSILON: f64 = 1e-12; +/// How long an end user's budget is reused before it is read again. +const END_USER_TTL: Duration = Duration::from_secs(60); +/// End users held in memory; the cache starts over past this. +const MAX_END_USERS: usize = 10_000; +const FAIL_CLOSED_READ: &str = "{'error': 'Budget enforcement unavailable: current spend could not be verified against Redis or the database, and fail_closed_budget_enforcement is enabled, so the request was rejected to avoid exceeding the configured budget. Retry shortly.'}"; + +/// Seeds or raises each counter to its floor, runs the budget checks in the order given and +/// reserves the estimate on every counter that takes a reservation, shrinking it to what is +/// left when a counter would go over, as the proxy's reservation does. +const ADMIT: &str = r#" +local ttl = tonumber(ARGV[1]) +local estimate = tonumber(ARGV[2]) +local n = #KEYS +local current, max, mode, reserve = {}, {}, {}, {} +for i = 1, n do + local base = 3 + (i - 1) * 4 + local floor = ARGV[base + 1] + max[i] = tonumber(ARGV[base + 2]) + mode[i] = ARGV[base + 3] + reserve[i] = ARGV[base + 4] == '1' + local raw = redis.call('GET', KEYS[i]) + local value = raw and tonumber(raw) or nil + if value == nil or value < tonumber(floor) then + redis.call('SET', KEYS[i], floor, 'EX', ttl) + raw = floor + end + current[i] = raw +end +for index in string.gmatch(ARGV[3], '%d+') do + local i = tonumber(index) + local value = tonumber(current[i]) + if (mode[i] == '1' and value >= max[i]) or (mode[i] == '2' and value > max[i]) then + return {'check', i, current[i]} + end +end +local charged = {} +local applied = {} +local cost = estimate +for i = 1, n do + charged[i] = 0 + if reserve[i] and cost > 0 then + local raw = redis.call('INCRBYFLOAT', KEYS[i], cost) + redis.call('EXPIRE', KEYS[i], ttl) + charged[i] = cost + applied[#applied + 1] = i + local post = tonumber(raw) + if post > max[i] then + local remaining = max[i] - (post - cost) + if remaining <= 1e-12 then + for _, j in ipairs(applied) do + redis.call('INCRBYFLOAT', KEYS[j], -charged[j]) + end + return {'reserve', i, raw} + end + for _, j in ipairs(applied) do + redis.call('INCRBYFLOAT', KEYS[j], -(charged[j] - remaining)) + charged[j] = remaining + end + cost = remaining + end + end +end +local out = {'ok'} +for i = 1, n do + out[#out + 1] = string.format('%.17g', charged[i]) +end +return out +"#; + +/// Adds each request's settled cost: seeds a missing counter from its floor (charging the whole +/// cost, since its reservation was lost), raises one below its floor, then applies the change. +const SETTLE: &str = r#" +local ttl = tonumber(ARGV[1]) +for i, key in ipairs(KEYS) do + local base = 1 + (i - 1) * 3 + local floor, delta, fresh = ARGV[base + 1], ARGV[base + 2], ARGV[base + 3] + local raw = redis.call('GET', key) + local value = raw and tonumber(raw) or nil + if value == nil then + redis.call('SET', key, floor) + delta = fresh + elseif value < tonumber(floor) then + redis.call('SET', key, floor) + end + if tonumber(delta) ~= 0 then + redis.call('INCRBYFLOAT', key, delta) + end + redis.call('EXPIRE', key, ttl) +end +return #KEYS +"#; + +#[derive(Clone, Debug, PartialEq)] +pub struct BudgetSettings { + /// `litellm_settings.budget_exceeded_status_code`, 422 by default. + pub status: u16, + /// Off with `general_settings.disable_budget_reservation`. + pub reservations: bool, + /// `general_settings.apply_user_budget_to_team_keys`: a user's budget also covers their + /// team keys, not only their personal ones. + pub user_budget_on_team_keys: bool, + /// `litellm_settings.max_end_user_budget_id`: the budget of end users without their own. + pub end_user_budget_id: Option, + /// `general_settings.fail_closed_budget_enforcement`: refuse rather than fall back to this + /// replica's counters while Redis is unreachable. + pub fail_closed: bool, + /// `default_redis_ttl` (in `litellm_settings` or `cache_params`), 60 s by default. + pub ttl: Duration, +} + +impl BudgetSettings { + pub fn from_config(config: &Config) -> Self { + let general = &config.general_settings.additional_fields; + let litellm = &config.litellm_settings; + let ttl = litellm + .additional_fields + .get("default_redis_ttl") + .or_else(|| { + litellm + .cache_params + .as_ref() + .and_then(|params| params.get("default_redis_ttl")) + }) + .and_then(number) + .filter(|seconds| *seconds >= 1.0) + .unwrap_or(60.0); + Self { + status: litellm + .additional_fields + .get("budget_exceeded_status_code") + .and_then(number) + .and_then(|status| u16::try_from(status as i64).ok()) + .filter(|status| (400..600).contains(status)) + .unwrap_or(422), + reservations: !flag(general.get("disable_budget_reservation")), + user_budget_on_team_keys: flag(general.get("apply_user_budget_to_team_keys")), + end_user_budget_id: litellm + .additional_fields + .get("max_end_user_budget_id") + .and_then(litellm_config::Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned), + fail_closed: flag(general.get("fail_closed_budget_enforcement")), + ttl: Duration::from_secs(ttl as u64), + } + } +} + +fn number(value: &litellm_config::Value) -> Option { + match value { + litellm_config::Value::Number(number) => number.as_f64(), + litellm_config::Value::String(text) => text.trim().parse().ok(), + _ => None, + } +} + +fn flag(value: Option<&litellm_config::Value>) -> bool { + match value { + Some(litellm_config::Value::Bool(flag)) => *flag, + Some(litellm_config::Value::String(text)) => text.eq_ignore_ascii_case("true"), + _ => false, + } +} + +/// When a counter refuses at check time. +#[derive(Clone, Copy, Debug, PartialEq)] +enum Check { + None, + /// `spend >= max_budget`. + AtLeast, + /// `spend > max_budget`, for teams and end users. + Above, +} + +/// Whose counter it is, which names it in refusals. +#[derive(Clone, Debug, PartialEq)] +enum Owner { + Key { + hash: String, + descriptor: String, + }, + Team(String), + User(String), + /// The team member check the proxy runs while building the key: a positive budget. + Member { + user_id: String, + team_id: String, + }, + /// The proxy's later team member check, which also refuses on a zero budget. + ZeroBudgetMember { + user_id: String, + team_id: String, + }, + EndUser(String), +} + +impl Owner { + /// The counter, before the namespace, as the proxy names it. + fn counter(&self) -> String { + match self { + Self::Key { hash, .. } => format!("spend:key:{hash}"), + Self::Team(id) => format!("spend:team:{id}"), + Self::User(id) => format!("spend:user:{id}"), + Self::Member { user_id, team_id } | Self::ZeroBudgetMember { user_id, team_id } => { + format!("spend:team_member:{user_id}:{team_id}") + } + Self::EndUser(id) => format!("spend:end_user:{id}"), + } + } + + /// The auth-time refusal message. + fn check_message(&self, spend: f64, max_budget: f64) -> String { + let (spend, max_budget) = (python_float(spend), python_float(max_budget)); + match self { + Self::Key { descriptor, .. } => format!( + "Budget has been exceeded! Key={descriptor} Current cost: {spend}, Max budget: {max_budget}" + ), + Self::Team(id) => format!( + "Budget has been exceeded! Team={id} Current cost: {spend}, Max budget: {max_budget}" + ), + Self::User(id) => { + format!("ExceededBudget: User={id} over budget. Spend={spend}, Budget={max_budget}") + } + Self::Member { user_id, team_id } => format!( + "Budget has been exceeded! TeamMember={user_id}:{team_id} Current cost: {spend}, Max budget: {max_budget}" + ), + Self::ZeroBudgetMember { user_id, team_id } => format!( + "Budget has been exceeded! User={user_id} in Team={team_id} Current cost: {spend}, Max budget: {max_budget}" + ), + Self::EndUser(id) => format!( + "ExceededBudget: End User={id} over budget. Spend={spend}, Budget={max_budget}" + ), + } + } + + /// The reservation's refusal message, which names keys by their hash. + fn reservation_message(&self, spend: f64, max_budget: f64) -> String { + let owner = match self { + Self::Key { hash, .. } => format!("Key={hash}"), + Self::Team(id) => format!("Team={id}"), + Self::User(id) => format!("User={id}"), + Self::Member { user_id, team_id } | Self::ZeroBudgetMember { user_id, team_id } => { + format!("TeamMember={user_id}:{team_id}") + } + Self::EndUser(id) => format!("EndUser={id}"), + }; + format!( + "Budget has been exceeded! {owner} Current cost: {}, Max budget: {}", + python_float(spend), + python_float(max_budget) + ) + } +} + +/// One spend counter a request touches. +#[derive(Clone, Debug, PartialEq)] +struct Counter { + owner: Owner, + /// The stored spend, below which the counter is never read. + floor: f64, + max_budget: Option, + check: Check, + reserve: bool, +} + +/// The counters a request touches, in the proxy's reservation order, and the order its checks +/// run in. +#[derive(Debug, Default, PartialEq)] +struct Plan { + counters: Vec, + checks: Vec, +} + +/// How admission went. +#[derive(Debug, PartialEq)] +enum Admitted { + /// What was reserved on each counter. + Charged(Vec), + /// A check refused, with the counter's spend. + Over(usize, f64), + /// The reservation could not fit, with the counter's spend after it. + Reserved(usize, f64), +} + +/// Counters in this replica's memory: the proxy's in-memory spend counters, used without Redis +/// and while it is unreachable. +#[derive(Default)] +struct Local(Mutex>); + +impl Local { + fn admit(&self, keys: &[String], plan: &Plan, estimate: f64, ttl: Duration) -> Admitted { + let mut counters = self.0.lock().unwrap_or_else(|e| e.into_inner()); + let now = Instant::now(); + if counters.len() > 100_000 { + counters.retain(|_, (_, expires)| *expires > now); + } + let mut current = Vec::with_capacity(keys.len()); + for (key, counter) in keys.iter().zip(&plan.counters) { + let entry = counters + .entry(key.clone()) + .or_insert((counter.floor, now + ttl)); + if entry.1 <= now || entry.0 < counter.floor { + *entry = (counter.floor, now + ttl); + } + current.push(entry.0); + } + for &index in &plan.checks { + let counter = &plan.counters[index]; + if over(counter.check, current[index], counter.max_budget) { + return Admitted::Over(index, current[index]); + } + } + let mut charged = vec![0.0; keys.len()]; + let mut applied = Vec::new(); + let mut cost = estimate; + for (index, counter) in plan.counters.iter().enumerate() { + if !counter.reserve || cost <= 0.0 { + continue; + } + let Some(max_budget) = counter.max_budget else { + continue; + }; + let entry = counters.get_mut(&keys[index]).expect("seeded above"); + entry.0 += cost; + entry.1 = now + ttl; + let post = entry.0; + charged[index] = cost; + applied.push(index); + if post > max_budget { + let remaining = max_budget - (post - cost); + let refund = + |counters: &mut HashMap, index: usize, amount| { + if let Some(entry) = counters.get_mut(&keys[index]) { + entry.0 -= amount; + } + }; + if remaining <= EPSILON { + for &applied in &applied { + refund(&mut counters, applied, charged[applied]); + } + return Admitted::Reserved(index, post); + } + for &applied in &applied { + refund(&mut counters, applied, charged[applied] - remaining); + charged[applied] = remaining; + } + cost = remaining; + } + } + Admitted::Charged(charged) + } + + fn settle(&self, keys: &[String], changes: &[Change], ttl: Duration) { + let mut counters = self.0.lock().unwrap_or_else(|e| e.into_inner()); + let now = Instant::now(); + for (key, change) in keys.iter().zip(changes) { + let entry = counters.entry(key.clone()).or_insert((f64::NAN, now)); + let delta = if entry.0.is_nan() || entry.1 <= now { + entry.0 = change.floor; + change.fresh + } else { + entry.0 = entry.0.max(change.floor); + change.delta + }; + entry.0 += delta; + entry.1 = now + ttl; + } + } +} + +fn over(check: Check, spend: f64, max_budget: Option) -> bool { + match (check, max_budget) { + (Check::AtLeast, Some(max_budget)) => spend >= max_budget, + (Check::Above, Some(max_budget)) => spend > max_budget, + _ => false, + } +} + +/// How a counter changes when a request ends. +#[derive(Clone, Copy, Debug, PartialEq)] +struct Change { + floor: f64, + /// Added to a counter that still exists. + delta: f64, + /// Added to a counter that expired while the request ran. + fresh: f64, +} + +/// Where the counters live. +struct Counters { + redis: Option>, + local: Local, + ttl: Duration, + fail_closed: bool, +} + +/// Redis could not be used and the settings say to refuse rather than fall back. +struct Unavailable; + +impl Counters { + fn connection(&self) -> Option<(ConnectionManager, &Coordination)> { + let redis = self.redis.as_deref()?; + Some((redis.connection()?, redis)) + } + + async fn admit(&self, plan: &Plan, estimate: f64) -> Result<(Admitted, bool), Unavailable> { + let keys: Vec = plan + .counters + .iter() + .map(|counter| counter.owner.counter()) + .collect(); + if self.redis.is_some() { + match self.connection() { + Some((mut connection, redis)) => { + let keys: Vec = keys.iter().map(|key| redis.namespaced(key)).collect(); + match admit_redis(&mut connection, &keys, plan, estimate, self.ttl).await { + Ok(admitted) => return Ok((admitted, true)), + Err(error) => { + tracing::warn!(%error, "reading spend counters failed; using this replica's counters") + } + } + } + None => tracing::debug!("Redis not connected; using this replica's spend counters"), + } + if self.fail_closed { + return Err(Unavailable); + } + } + Ok((self.local.admit(&keys, plan, estimate, self.ttl), false)) + } + + async fn settle(&self, keys: Vec, changes: Vec, shared: bool) { + if shared && let Some((mut connection, redis)) = self.connection() { + let keys: Vec = keys.iter().map(|key| redis.namespaced(key)).collect(); + let mut command = redis::cmd("EVAL"); + command + .arg(SETTLE) + .arg(keys.len()) + .arg(&keys) + .arg(self.ttl.as_secs()); + for change in &changes { + command + .arg(change.floor.to_string()) + .arg(change.delta.to_string()) + .arg(change.fresh.to_string()); + } + if let Err(error) = command.query_async::(&mut connection).await { + // A counter left half-written would mislead every replica; drop it so the next + // read reseeds it from the database, as the proxy does. + tracing::warn!(%error, "updating spend counters failed; dropping them"); + let _: Result = redis::cmd("DEL") + .arg(&keys) + .query_async(&mut connection) + .await; + } + return; + } + self.local.settle(&keys, &changes, self.ttl); + } + + /// Keeps reserved counters alive while a request runs, renewing them every half TTL. + fn lease(&self, keys: Vec) -> Option> { + let (mut connection, redis) = self.connection()?; + let keys: Vec = keys.iter().map(|key| redis.namespaced(key)).collect(); + let ttl = self.ttl; + let (done, mut finished) = tokio::sync::oneshot::channel::<()>(); + tokio::spawn(async move { + let mut renew = tokio::time::interval(ttl / 2); + renew.tick().await; + loop { + tokio::select! { + _ = &mut finished => return, + _ = renew.tick() => { + let mut pipeline = redis::pipe(); + for key in &keys { + pipeline.cmd("EXPIRE").arg(key).arg(ttl.as_secs()).ignore(); + } + let _: Result<(), _> = pipeline.query_async(&mut connection).await; + } + } + } + }); + Some(done) + } +} + +async fn admit_redis( + connection: &mut ConnectionManager, + keys: &[String], + plan: &Plan, + estimate: f64, + ttl: Duration, +) -> redis::RedisResult { + let mut command = redis::cmd("EVAL"); + command + .arg(ADMIT) + .arg(keys.len()) + .arg(keys) + .arg(ttl.as_secs()) + .arg(estimate.to_string()) + .arg( + plan.checks + .iter() + .map(|index| (index + 1).to_string()) + .collect::>() + .join(","), + ); + for counter in &plan.counters { + let mode = match (counter.check, counter.max_budget) { + (_, None) | (Check::None, _) => "0", + (Check::AtLeast, Some(_)) => "1", + (Check::Above, Some(_)) => "2", + }; + command + .arg(counter.floor.to_string()) + .arg( + counter + .max_budget + .map(|max| max.to_string()) + .unwrap_or_default(), + ) + .arg(mode) + .arg(if counter.reserve && counter.max_budget.is_some() { + "1" + } else { + "0" + }); + } + let reply: Vec = command.query_async(connection).await?; + let text = |value: &redis::Value| -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(text) => Some(text.clone()), + redis::Value::Int(number) => Some(number.to_string()), + _ => None, + } + }; + let malformed = || { + redis::RedisError::from(( + redis::ErrorKind::UnexpectedReturnType, + "unexpected budget script reply", + )) + }; + let status = reply.first().and_then(text).ok_or_else(malformed)?; + let float = |value: Option<&redis::Value>| { + value + .and_then(text) + .and_then(|value| value.parse::().ok()) + .ok_or_else(malformed) + }; + let index = || { + reply + .get(1) + .and_then(text) + .and_then(|value| value.parse::().ok()) + .and_then(|index| index.checked_sub(1)) + .filter(|index| *index < keys.len()) + .ok_or_else(malformed) + }; + match status.as_str() { + "ok" => reply[1..] + .iter() + .map(|value| float(Some(value))) + .collect::>() + .map(Admitted::Charged), + "check" => Ok(Admitted::Over(index()?, float(reply.get(2))?)), + "reserve" => Ok(Admitted::Reserved(index()?, float(reply.get(2))?)), + _ => Err(malformed()), + } +} + +/// An end user and the default budget it was read with. +type EndUserKey = (String, Option); + +/// End users' budgets, read from the proxy's database and reused for a minute. +struct EndUsers { + pool: PgPool, + cache: Mutex>, +} + +impl EndUsers { + async fn get(&self, id: &str, default_budget_id: Option<&str>) -> Option { + let key = (id.to_owned(), default_budget_id.map(str::to_owned)); + let now = Instant::now(); + if let Some((end_user, read_at)) = self + .cache + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&key) + && now.duration_since(*read_at) < END_USER_TTL + { + return Some(end_user.clone()); + } + match litellm_db::find_end_user(&self.pool, id, default_budget_id).await { + Ok(end_user) => { + let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); + if cache.len() >= MAX_END_USERS { + cache.clear(); + } + cache.insert(key, (end_user.clone(), now)); + Some(end_user) + } + Err(error) => { + tracing::warn!(%error, "reading an end user's budget failed; not checking it"); + None + } + } + } +} + +/// The `budgets` hook. +pub struct Budgets { + settings: BudgetSettings, + counters: Arc, + end_users: Option, + /// Each model group's deployment ids, for pricing a request's most expensive outcome. + groups: HashMap>>, + prices: OnceLock>, +} + +impl Budgets { + /// `deployments` pairs each deployment's model group with its id. `redis`, when given, holds + /// the shared counters; `database` gives end users' budgets. + pub fn new( + settings: BudgetSettings, + deployments: impl IntoIterator)>, + redis: Option>, + database: Option, + ) -> Self { + let mut groups: HashMap>> = HashMap::new(); + for (group, id) in deployments { + groups.entry(group).or_default().push(id); + } + Self { + counters: Arc::new(Counters { + redis, + local: Local::default(), + ttl: settings.ttl, + fail_closed: settings.fail_closed, + }), + settings, + end_users: database.map(|pool| EndUsers { + pool, + cache: Mutex::default(), + }), + groups, + prices: OnceLock::new(), + } + } + + /// The hook for a gateway: its deployments, the coordination Redis and the proxy database. + pub fn for_gateway( + config: &Config, + inference: &litellm_gateway_inference::Gateway, + redis: Option>, + database: Option, + ) -> Self { + let deployments: Vec<(String, Arc)> = config + .model_list + .iter() + .map(|model| model.model_name.clone()) + .zip(inference.models.deployment_ids().map(Arc::from)) + .collect(); + Self::new( + BudgetSettings::from_config(config), + deployments, + redis, + database, + ) + } + + /// Every deployment of the requested group has explicit zero prices: the proxy skips budget + /// checks for free models. + fn free(&self, request: &RequestFacts) -> bool { + let (Some(prices), Some(ids)) = ( + self.prices.get(), + self.groups.get(&*request.requested_model), + ) else { + return false; + }; + let usage = Usage { + input_tokens: 1_000_000, + output_tokens: 1_000_000, + ..Usage::default() + }; + !ids.is_empty() && ids.iter().all(|id| prices.cost(id, &usage) == Some(0.0)) + } + + /// The proxy's estimate of a request's most expensive cost: its input, and as much output + /// as it may ask for, priced on the group's most expensive deployment. `None` when a + /// deployment's prices are unknown. + fn estimate(&self, request: &RequestFacts) -> Option { + let prices = self.prices.get()?; + let ids = self.groups.get(&*request.requested_model)?; + let input_tokens = request.input_tokens(); + let mut most = None::; + for id in ids { + let output_tokens = request.max_tokens.unwrap_or(DEFAULT_OUTPUT_TOKENS).min( + prices + .max_output_tokens(id) + .unwrap_or(DEFAULT_OUTPUT_TOKENS), + ); + let cost = prices.cost( + id, + &Usage { + input_tokens, + output_tokens, + ..Usage::default() + }, + )?; + most = Some(most.map_or(cost, |most| most.max(cost))); + } + most + } + + async fn plan(&self, request: &RequestFacts) -> Plan { + let mut plan = Plan::default(); + let Some(key) = request.key.as_deref() else { + return plan; + }; + let budgets = &key.budgets; + let reserve = |max_budget: Option| { + self.settings.reservations && max_budget.is_some_and(|max| max > 0.0 && max.is_finite()) + }; + let finite = |max_budget: Option| max_budget.filter(|max| max.is_finite()); + let push = |plan: &mut Plan, counter: Counter| { + plan.counters.push(counter); + plan.counters.len() - 1 + }; + // Reservation order: key, team, user, end user, team member. + let mut key_index = None; + if !key.hash.is_empty() { + let max_budget = finite(budgets.key.max_budget); + key_index = Some(push( + &mut plan, + Counter { + owner: Owner::Key { + hash: key.hash.clone(), + descriptor: key_descriptor( + key.alias.as_deref(), + budgets.key_name.as_deref(), + ), + }, + floor: budgets.key.spend, + max_budget, + check: Check::AtLeast, + reserve: reserve(max_budget), + }, + )); + } + let mut team_index = None; + if let Some(team_id) = &key.team_id { + let team = budgets.team.unwrap_or_default(); + let max_budget = finite(team.max_budget); + team_index = Some(push( + &mut plan, + Counter { + owner: Owner::Team(team_id.clone()), + floor: team.spend, + max_budget, + check: Check::Above, + reserve: reserve(max_budget), + }, + )); + } + let mut user_index = None; + if let Some(user_id) = &key.user_id { + let user = budgets.user.unwrap_or_default(); + let applies = key.team_id.is_none() || self.settings.user_budget_on_team_keys; + let max_budget = finite(user.max_budget).filter(|_| applies); + user_index = Some(push( + &mut plan, + Counter { + owner: Owner::User(user_id.clone()), + floor: user.spend, + max_budget, + check: Check::AtLeast, + reserve: reserve(max_budget), + }, + )); + } + let mut end_user_index = None; + if let Some(end_user_id) = request.end_user.as_deref() { + let default_budget_id = budgets + .end_user_budget_id + .as_deref() + .or(self.settings.end_user_budget_id.as_deref()); + let end_user = match &self.end_users { + Some(end_users) => end_users.get(end_user_id, default_budget_id).await, + None => None, + } + .unwrap_or_default(); + let max_budget = finite(end_user.max_budget); + end_user_index = Some(push( + &mut plan, + Counter { + owner: Owner::EndUser(end_user_id.to_owned()), + floor: end_user.spend.unwrap_or(0.0), + max_budget, + // The proxy's check needs an end-user row; its reservation does not. + check: match end_user.spend { + Some(_) => Check::Above, + None => Check::None, + }, + reserve: reserve(max_budget), + }, + )); + } + let mut member_index = None; + if let (Some(user_id), Some(team_id)) = (&key.user_id, &key.team_id) { + let member = budgets.team_member; + let max_budget = member.and_then(|member| { + finite(member.max_budget) + .or_else(|| finite(member.default_max_budget).filter(|max| *max > 0.0)) + }); + let (owner, check) = match max_budget { + Some(max) if max <= 0.0 => ( + Owner::ZeroBudgetMember { + user_id: user_id.clone(), + team_id: team_id.clone(), + }, + Check::AtLeast, + ), + _ => ( + Owner::Member { + user_id: user_id.clone(), + team_id: team_id.clone(), + }, + Check::AtLeast, + ), + }; + member_index = Some(push( + &mut plan, + Counter { + owner, + floor: member.map_or(0.0, |member| member.spend), + max_budget, + check, + reserve: reserve(max_budget), + }, + )); + } + // Check order: the positive member budget and the key while the proxy builds the key, + // then team, user, a zero member budget and the end user in its common checks. + let positive_member = member_index + .filter(|index| matches!(plan.counters[*index].owner, Owner::Member { .. })); + let zero_member = member_index + .filter(|index| matches!(plan.counters[*index].owner, Owner::ZeroBudgetMember { .. })); + plan.checks = [ + positive_member, + key_index, + team_index, + user_index, + zero_member, + end_user_index, + ] + .into_iter() + .flatten() + .filter(|index| plan.counters[*index].max_budget.is_some()) + .collect(); + plan + } + + fn refusal(&self, message: String) -> Refusal { + Refusal::new(self.settings.status, "budget_exceeded", message).without_param() + } +} + +impl Hook for Budgets { + fn name(&self) -> &'static str { + "budgets" + } + + fn reads(&self) -> Facts { + Facts::KEY | Facts::END_USER | Facts::INPUT_TOKENS | Facts::COST + } + + fn writes(&self) -> Targets { + Targets::ADMISSION + } + + fn prepare(&self, prices: &Arc) { + let _ = self.prices.set(prices.clone()); + } + + fn run<'a>( + &'a self, + request: &'a Arc, + _edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + let plan = self.plan(request).await; + if plan.counters.is_empty() { + return Ok(None); + } + let checked = !self.free(request); + let reserving = checked && plan.counters.iter().any(|counter| counter.reserve); + let estimate = match reserving { + true => self.estimate(request).filter(|cost| *cost > 0.0), + false => None, + }; + let charged = if checked && (!plan.checks.is_empty() || estimate.is_some()) { + let (admitted, shared) = self + .counters + .admit(&plan, estimate.unwrap_or(0.0)) + .await + .map_err(|Unavailable| Refusal::new(503, "auth_error", FAIL_CLOSED_READ))?; + match admitted { + Admitted::Charged(charged) => Some((charged, shared)), + Admitted::Over(index, spend) => { + let counter = &plan.counters[index]; + let max_budget = counter.max_budget.unwrap_or_default(); + return Err(self.refusal(counter.owner.check_message(spend, max_budget))); + } + Admitted::Reserved(index, spend) => { + let counter = &plan.counters[index]; + let max_budget = counter.max_budget.unwrap_or_default(); + return Err( + self.refusal(counter.owner.reservation_message(spend, max_budget)) + ); + } + } + } else { + None + }; + let shared = charged + .as_ref() + .map_or(self.counters.redis.is_some(), |(_, shared)| *shared); + let charged = + charged.map_or_else(|| vec![0.0; plan.counters.len()], |(charged, _)| charged); + let reserved: Vec = plan + .counters + .iter() + .zip(&charged) + .filter(|(_, charged)| **charged != 0.0) + .map(|(counter, _)| counter.owner.counter()) + .collect(); + let lease = match shared && !reserved.is_empty() { + true => self.counters.lease(reserved), + false => None, + }; + Ok(Some(Box::new(Reservation { + counters: self.counters.clone(), + plan, + charged, + shared, + _lease: lease, + }) as Box)) + }) + } +} + +/// What a request reserved, settled to its real cost when it ends. +struct Reservation { + counters: Arc, + plan: Plan, + charged: Vec, + shared: bool, + /// Dropping it stops the lease. + _lease: Option>, +} + +impl Settle for Reservation { + fn settle(self: Box, record: &Record) { + let success = (200..300).contains(&record.outcome.status); + let cost = match success { + true => record.outcome.cost.unwrap_or(0.0), + false => 0.0, + }; + let mut keys = Vec::new(); + let mut changes = Vec::new(); + for (counter, charged) in self.plan.counters.iter().zip(&self.charged) { + let delta = cost - charged; + if delta == 0.0 && *charged == 0.0 { + continue; + } + keys.push(counter.owner.counter()); + changes.push(Change { + floor: counter.floor, + delta, + fresh: cost, + }); + } + if keys.is_empty() { + return; + } + let Reservation { + counters, + shared, + _lease, + .. + } = *self; + tokio::spawn(async move { + counters.settle(keys, changes, shared).await; + drop(_lease); + }); + } +} + +/// The proxy names a key by its alias, adding `key_name` only while it still looks masked. +fn key_descriptor(alias: Option<&str>, key_name: Option<&str>) -> String { + let alias = alias.filter(|alias| !alias.is_empty()).unwrap_or("key"); + match key_name.filter(|name| is_masked(name)) { + Some(name) => format!("{alias} ({name})"), + None => alias.to_owned(), + } +} + +/// The proxy's `_MASKED_KEY_NAME_RE`, `sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?`. +fn is_masked(name: &str) -> bool { + name.strip_prefix("sk-...").is_some_and(|rest| { + rest.is_empty() + || (rest.chars().count() == 4 + && rest + .chars() + .all(|character| !character.is_whitespace() && !character.is_control())) + }) +} + +/// Python's `repr` of a float. +fn python_float(value: f64) -> String { + if value.is_nan() { + return "nan".into(); + } + if value.is_infinite() { + return if value > 0.0 { "inf" } else { "-inf" }.into(); + } + let magnitude = value.abs(); + if magnitude != 0.0 && !(1e-4..1e16).contains(&magnitude) { + let scientific = format!("{value:e}"); + let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0")); + let exponent: i32 = exponent.parse().unwrap_or(0); + let sign = if exponent < 0 { '-' } else { '+' }; + return format!("{mantissa}e{sign}{:02}", exponent.abs()); + } + format!("{value:?}") +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::no_alias(None, None, "key")] + #[case::empty_alias(Some(""), None, "key")] + #[case::alias(Some("prod"), None, "prod")] + #[case::masked_name(Some("prod"), Some("sk-...abcd"), "prod (sk-...abcd)")] + #[case::bare_masked_name(None, Some("sk-..."), "key (sk-...)")] + #[case::unmasked_name(Some("prod"), Some("sk-unmasked90"), "prod")] + #[case::name_with_space(Some("prod"), Some("sk-...ab d"), "prod")] + #[case::name_too_long(Some("prod"), Some("sk-...abcde"), "prod")] + fn describes_keys_like_the_proxy( + #[case] alias: Option<&str>, + #[case] name: Option<&str>, + #[case] expected: &str, + ) { + assert_eq!(key_descriptor(alias, name), expected); + } + + #[rstest] + #[case::whole(10.0, "10.0")] + #[case::fraction(0.1, "0.1")] + #[case::zero(0.0, "0.0")] + #[case::small(0.0001, "0.0001")] + #[case::tiny(0.00001, "1e-05")] + #[case::large(1e15, "1000000000000000.0")] + #[case::huge(1e16, "1e+16")] + #[case::huge_fraction(1.5e17, "1.5e+17")] + #[case::negative(-2.5, "-2.5")] + #[case::infinite(f64::INFINITY, "inf")] + fn prints_floats_like_python(#[case] value: f64, #[case] expected: &str) { + assert_eq!(python_float(value), expected); + } + + fn counter(owner: Owner, floor: f64, max_budget: Option, check: Check) -> Counter { + Counter { + owner, + floor, + max_budget, + check, + reserve: max_budget.is_some_and(|max| max > 0.0), + } + } + + fn plan(counters: Vec) -> (Vec, Plan) { + let keys = counters.iter().map(|c| c.owner.counter()).collect(); + let checks = (0..counters.len()).collect(); + (keys, Plan { counters, checks }) + } + + #[rstest] + fn local_counters_check_reserve_shrink_and_refuse() { + let local = Local::default(); + let ttl = Duration::from_secs(60); + let (keys, plan) = plan(vec![ + counter(Owner::Team("t".into()), 4.0, Some(10.0), Check::Above), + counter(Owner::User("u".into()), 0.0, Some(100.0), Check::AtLeast), + ]); + // Fits: both counters take the whole estimate. + assert_eq!( + local.admit(&keys, &plan, 5.0, ttl), + Admitted::Charged(vec![5.0, 5.0]) + ); + // Team has 1.0 left: the reservation shrinks to it on every counter. + assert_eq!( + local.admit(&keys, &plan, 5.0, ttl), + Admitted::Charged(vec![1.0, 1.0]) + ); + // Exactly at budget passes the team's strict check, but nothing is left to reserve. + assert_eq!( + local.admit(&keys, &plan, 5.0, ttl), + Admitted::Reserved(0, 15.0) + ); + // Settling the first request to 0.5 frees 4.5. + local.settle( + &keys, + &[ + Change { + floor: 4.0, + delta: -4.5, + fresh: 0.5, + }, + Change { + floor: 0.0, + delta: -4.5, + fresh: 0.5, + }, + ], + ttl, + ); + assert_eq!( + local.admit(&keys, &plan, 1.0, ttl), + Admitted::Charged(vec![1.0, 1.0]) + ); + // A stored spend above the counter raises it, and the check refuses. + let (keys, plan) = plan_with_floor(30.0); + assert_eq!(local.admit(&keys, &plan, 1.0, ttl), Admitted::Over(0, 30.0)); + } + + fn plan_with_floor(floor: f64) -> (Vec, Plan) { + plan(vec![counter( + Owner::Team("t".into()), + floor, + Some(10.0), + Check::Above, + )]) + } + + #[rstest] + fn refusals_read_like_the_proxy() { + let key = Owner::Key { + hash: "abc".into(), + descriptor: "prod (sk-...abcd)".into(), + }; + assert_eq!( + key.check_message(10.0, 10.0), + "Budget has been exceeded! Key=prod (sk-...abcd) Current cost: 10.0, Max budget: 10.0" + ); + assert_eq!( + key.reservation_message(12.5, 10.0), + "Budget has been exceeded! Key=abc Current cost: 12.5, Max budget: 10.0" + ); + assert_eq!( + Owner::EndUser("e".into()).check_message(2.0, 1.0), + "ExceededBudget: End User=e over budget. Spend=2.0, Budget=1.0" + ); + assert_eq!( + Owner::ZeroBudgetMember { + user_id: "u".into(), + team_id: "t".into() + } + .check_message(0.0, 0.0), + "Budget has been exceeded! User=u in Team=t Current cost: 0.0, Max budget: 0.0" + ); + } +} diff --git a/litellm-rust/crates/gateway/src/coordination.rs b/litellm-rust/crates/gateway/src/coordination.rs index eeb01158e59..48cd6ac14cf 100644 --- a/litellm-rust/crates/gateway/src/coordination.rs +++ b/litellm-rust/crates/gateway/src/coordination.rs @@ -267,10 +267,16 @@ impl Coordination { })) } - fn connection(&self) -> Option { + /// The shared connection, once Redis has answered. + pub(crate) fn connection(&self) -> Option { self.connection.get().cloned() } + /// `key` under the configured namespace, as the proxy prefixes its keys. + pub(crate) fn namespaced(&self, key: &str) -> String { + self.settings.key(key) + } + /// Whether the connection shared state is written and read through is up yet. pub fn connected(&self) -> bool { self.connection.get().is_some() diff --git a/litellm-rust/crates/gateway/src/key_status.rs b/litellm-rust/crates/gateway/src/key_status.rs index f4f54e9ac3d..50d8bf65567 100644 --- a/litellm-rust/crates/gateway/src/key_status.rs +++ b/litellm-rust/crates/gateway/src/key_status.rs @@ -2,7 +2,8 @@ use litellm_db::{Restrictions, VerificationToken}; use std::sync::Arc; use litellm_gateway_auth::keys::{ - AllowedRoutes, KeyModels, KeyProfile, KeyStatus, ModelAccess, Owner, Refusal, Restriction, + AllowedRoutes, KeyBudgets, KeyModels, KeyProfile, KeyStatus, MemberSpend, ModelAccess, Owner, + Refusal, Restriction, Spend, }; use serde_json::Value; @@ -37,6 +38,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K Value::Object(fields) => Arc::new(fields), _ => Arc::default(), }; + let budgets = budgets(&row); let (team_models, team_alias, team_metadata) = match row.team { Some(team) => ( Some(ModelAccess::new(team.models)), @@ -59,10 +61,41 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K team_alias, metadata: object(row.metadata), team_metadata, + budgets, }), } } +/// The spend and budgets the budgets hook checks for this key. +fn budgets(row: &VerificationToken) -> KeyBudgets { + let membership = &row.team_membership; + KeyBudgets { + key: Spend { + spend: row.spend, + max_budget: row.max_budget, + }, + key_name: row.key_name.clone(), + team: row.team.as_ref().map(|team| Spend { + spend: team.spend, + max_budget: team.max_budget, + }), + user: row.user.as_ref().map(|user| Spend { + spend: user.spend, + max_budget: user.max_budget, + }), + team_member: membership.spend.map(|spend| MemberSpend { + spend, + max_budget: membership.max_budget, + default_max_budget: membership.default_max_budget, + }), + end_user_budget_id: row + .metadata + .get("end_user_budget_id") + .and_then(Value::as_str) + .map(str::to_owned), + } +} + fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option { let team = match (&row.team_id, &row.team) { (Some(team_id), None) => { @@ -80,29 +113,6 @@ fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option= max_budget - { - return Some(budget_exceeded( - format!("Key={}", key_label(row)), - row.spend, - max_budget, - )); - } - if let Some((team_id, team)) = team - && let Some(max_budget) = team.max_budget - && max_budget.is_finite() - && team.spend > max_budget - { - return Some(budget_exceeded( - format!("Team={team_id}"), - team.spend, - max_budget, - )); - } unsupported(row, personal_keys_use_teams) .map(|(restriction, owner)| Refusal::Unsupported { restriction, owner }) } @@ -197,57 +207,6 @@ fn is_set(value: &Value) -> bool { } } -fn budget_exceeded(label: String, spend: f64, max_budget: f64) -> Refusal { - Refusal::BudgetExceeded { - owner: label, - spend: python_float(spend), - max_budget: python_float(max_budget), - } -} - -/// The proxy names a key by its alias, adding `key_name` only while it still looks masked. -fn key_label(row: &VerificationToken) -> String { - let alias = row - .key_alias - .as_deref() - .filter(|alias| !alias.is_empty()) - .unwrap_or("key"); - match row.key_name.as_deref().filter(|name| is_masked(name)) { - Some(name) => format!("{alias} ({name})"), - None => alias.to_owned(), - } -} - -/// The proxy's `_MASKED_KEY_NAME_RE`, `sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?`. -fn is_masked(name: &str) -> bool { - name.strip_prefix("sk-...").is_some_and(|rest| { - rest.is_empty() - || (rest.chars().count() == 4 - && rest - .chars() - .all(|character| !character.is_whitespace() && !character.is_control())) - }) -} - -/// Python's `repr` of a float. -fn python_float(value: f64) -> String { - if value.is_nan() { - return "nan".into(); - } - if value.is_infinite() { - return if value > 0.0 { "inf" } else { "-inf" }.into(); - } - let magnitude = value.abs(); - if magnitude != 0.0 && !(1e-4..1e16).contains(&magnitude) { - let scientific = format!("{value:e}"); - let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0")); - let exponent: i32 = exponent.parse().unwrap_or(0); - let sign = if exponent < 0 { '-' } else { '+' }; - return format!("{mantissa}e{sign}{:02}", exponent.abs()); - } - format!("{value:?}") -} - #[cfg(test)] mod tests { use std::time::SystemTime; @@ -307,7 +266,13 @@ mod tests { }, allowed_routes: AllowedRoutes::new(strings(&["llm_api_routes"])), team_id: team.then(|| "team-1".to_owned()), - profile: Default::default(), + profile: Arc::new(KeyProfile { + budgets: KeyBudgets { + team: team.then(Spend::default), + ..KeyBudgets::default() + }, + ..KeyProfile::default() + }), } } @@ -318,14 +283,6 @@ mod tests { } } - fn exceeded(owner: &str, spend: &str, max_budget: &str) -> KeyStatus { - refused(Refusal::BudgetExceeded { - owner: owner.into(), - spend: spend.into(), - max_budget: max_budget.into(), - }) - } - fn unsupported(restriction: Restriction, owner: Owner) -> KeyStatus { refused(Refusal::Unsupported { restriction, owner }) } @@ -358,57 +315,6 @@ mod tests { 0.0, refused(Refusal::TeamBlocked { team_id: "team-1".into() }) )] - #[case::key_at_budget( - false, - None, - None, - Some(10.0), - 10.0, - exceeded("Key=key", "10.0", "10.0") - )] - #[case::key_under_budget(false, None, None, Some(10.0), 9.5, active(false))] - #[case::key_zero_budget(false, None, None, Some(0.0), 0.0, exceeded("Key=key", "0.0", "0.0"))] - #[case::key_infinite_budget(false, None, None, Some(f64::INFINITY), 1e20, active(false))] - #[case::team_over_budget( - false, - Some("team-1"), - team(false, Some(5.0), 5.25), - None, - 0.0, - exceeded("Team=team-1", "5.25", "5.0") - )] - #[case::team_at_budget( - false, - Some("team-1"), - team(false, Some(5.0), 5.0), - None, - 0.0, - active(true) - )] - #[case::team_over_budget_under_key_budget( - false, - Some("team-1"), - team(false, Some(5.0), 9.0), - Some(10.0), - 2.0, - exceeded("Team=team-1", "9.0", "5.0") - )] - #[case::key_and_team_under_budget( - false, - Some("team-1"), - team(false, Some(5.0), 1.0), - Some(10.0), - 2.0, - active(true) - )] - #[case::key_budget_checked_before_team( - false, - Some("team-1"), - team(false, Some(5.0), 9.0), - Some(1.0), - 2.0, - exceeded("Key=key", "2.0", "1.0") - )] fn refuses_like_the_proxy( row: VerificationToken, #[case] blocked: bool, @@ -467,6 +373,8 @@ mod tests { team: team_id.and_then(|_| team(false, None, 0.0)), user_id: Some("user-1".into()), user: Some(User { + spend: 0.0, + max_budget: None, models: strings(models), restrictions, }), @@ -483,6 +391,7 @@ mod tests { team_membership: TeamMembership { rate_limits, allowed_models: strings(allowed_models), + ..TeamMembership::default() }, ..on_team(NONE, json!({}), row) } @@ -691,7 +600,7 @@ mod tests { models, allowed_routes, team_id, - .. + profile, } => KeyStatus::Active { expires_at, models, @@ -699,7 +608,7 @@ mod tests { team_id, profile: Arc::new(KeyProfile { user_id: user_id.map(str::to_owned), - ..KeyProfile::default() + ..(*profile).clone() }), }, other => other, @@ -714,42 +623,4 @@ mod tests { } ); } - - #[rstest] - #[case::no_alias(None, None, "key")] - #[case::empty_alias(Some(""), None, "key")] - #[case::alias(Some("prod"), None, "prod")] - #[case::masked_name(Some("prod"), Some("sk-...abcd"), "prod (sk-...abcd)")] - #[case::bare_masked_name(None, Some("sk-..."), "key (sk-...)")] - #[case::unmasked_name(Some("prod"), Some("sk-unmasked90"), "prod")] - #[case::name_with_space(Some("prod"), Some("sk-...ab d"), "prod")] - #[case::name_too_long(Some("prod"), Some("sk-...abcde"), "prod")] - fn labels_keys_like_the_proxy( - row: VerificationToken, - #[case] alias: Option<&str>, - #[case] name: Option<&str>, - #[case] expected: &str, - ) { - let row = VerificationToken { - key_alias: alias.map(str::to_owned), - key_name: name.map(str::to_owned), - ..row - }; - assert_eq!(key_label(&row), expected); - } - - #[rstest] - #[case::whole(10.0, "10.0")] - #[case::fraction(0.1, "0.1")] - #[case::zero(0.0, "0.0")] - #[case::small(0.0001, "0.0001")] - #[case::tiny(0.00001, "1e-05")] - #[case::large(1e15, "1000000000000000.0")] - #[case::huge(1e16, "1e+16")] - #[case::huge_fraction(1.5e17, "1.5e+17")] - #[case::negative(-2.5, "-2.5")] - #[case::infinite(f64::INFINITY, "inf")] - fn prints_floats_like_python(#[case] value: f64, #[case] expected: &str) { - assert_eq!(python_float(value), expected); - } } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 7e5e50fcab1..dcd3aa994a3 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,4 +1,5 @@ mod auth; +pub mod budgets; pub mod coordination; pub mod drain; mod error; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index fccfa6cfc66..0954fd7634b 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -112,13 +112,22 @@ async fn main() -> Result<(), Box> { coordination.start(inference.clone(), keys.clone(), shutdown.clone()); } // Built-in hooks and observers are added here as the gateway gains them. + let mut hooks: Vec> = Vec::new(); + if database.is_some() { + hooks.push(Arc::new(litellm_gateway::budgets::Budgets::for_gateway( + &config, + &inference, + coordination.clone(), + database.clone(), + ))); + } 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)); } - litellm_gateway::hooks::install(&inference, &config, Vec::new(), observers).await?; + 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()?; // Without TCP_NODELAY, Nagle holds small SSE chunks until the client's delayed ACK (~40 ms per stream). diff --git a/litellm-rust/crates/gateway/tests/budgets.rs b/litellm-rust/crates/gateway/tests/budgets.rs new file mode 100644 index 00000000000..83a7778c732 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/budgets.rs @@ -0,0 +1,495 @@ +//! The budgets hook on real requests: refusals as the proxy words them, reservations that keep +//! concurrent spend under a budget, and counters shared through Redis. Set +//! `LITELLM_TEST_REDIS_URL` to also run the shared cases against a live Redis. + +use std::{ + collections::HashMap, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway::{ + budgets::{BudgetSettings, Budgets}, + coordination::{Coordination, RedisSettings}, +}; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyBudgets, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, + MemberSpend, Spend, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +/// Costs 3 * 0.001 + 4 * 0.002 = 0.011 at the test prices. +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; +const COST: f64 = 0.011; + +async fn upstream() -> (String, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counted = calls.clone(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(move |_: Request| { + let counted = counted.clone(); + async move { + counted.fetch_add(1, Ordering::SeqCst); + ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response() + } + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + (format!("http://{address}"), calls) +} + +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 hash(token: &str) -> KeyHash { + KeyHash::from_token(&SecretValue::new(token)) +} + +fn key(team_id: Option<&str>, budgets: KeyBudgets) -> KeyStatus { + KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: team_id.map(str::to_owned), + profile: Arc::new(KeyProfile { + alias: Some("prod".into()), + user_id: Some("user-1".into()), + budgets, + ..KeyProfile::default() + }), + } +} + +fn budget(spend: f64, max_budget: f64) -> Spend { + Spend { + spend, + max_budget: Some(max_budget), + } +} + +struct Gateway { + app: Router, + calls: Arc, +} + +async fn gateway( + keys: Vec<(&str, KeyStatus)>, + redis: Option>, + settings: &str, +) -> Gateway { + gateway_with_database(keys, redis, None, settings).await +} + +async fn gateway_with_database( + keys: Vec<(&str, KeyStatus)>, + redis: Option>, + database: Option, + settings: &str, +) -> Gateway { + let (upstream, calls) = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream + input_cost_per_token: 0.001 + output_cost_per_token: 0.002 + model_info: + id: claude-1 + - model_name: free + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream + input_cost_per_token: 0 + output_cost_per_token: 0 + model_info: + id: free-1 +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +{settings}" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let budgets = Budgets::for_gateway(&config, &inference, redis, database); + litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(budgets)], Vec::new()) + .await + .unwrap(); + let keys: Arc = Arc::new(Keys( + keys.into_iter() + .map(|(token, status)| (hash(token), status)) + .collect(), + )); + let app = litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys)); + Gateway { app, calls } +} + +async fn send(app: &Router, token: &str, model: &str) -> (u16, Value) { + send_as(app, token, model, None).await +} + +async fn send_as(app: &Router, token: &str, model: &str, end_user: Option<&str>) -> (u16, Value) { + let mut body = json!({ + "model": model, + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + if let Some(end_user) = end_user { + body["metadata"] = json!({"user_id": end_user}); + } + let response = app + .clone() + .oneshot( + Request::post("/v1/messages") + .header("content-type", "application/json") + .header("authorization", format!("Bearer {token}")) + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status().as_u16(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + (status, serde_json::from_slice(&body).unwrap_or(Value::Null)) +} + +/// Settling runs after the response, off the request path. +async fn settled() { + tokio::time::sleep(Duration::from_millis(100)).await; +} + +#[rstest] +#[case::key_at_budget( + None, + KeyBudgets { key: budget(10.0, 10.0), ..KeyBudgets::default() }, + "Budget has been exceeded! Key=prod Current cost: 10.0, Max budget: 10.0" +)] +#[case::zero_key_budget( + None, + KeyBudgets { key: budget(0.0, 0.0), ..KeyBudgets::default() }, + "Budget has been exceeded! Key=prod Current cost: 0.0, Max budget: 0.0" +)] +#[case::team_over_budget( + Some("team-1"), + KeyBudgets { team: Some(budget(5.25, 5.0)), ..KeyBudgets::default() }, + "Budget has been exceeded! Team=team-1 Current cost: 5.25, Max budget: 5.0" +)] +#[case::personal_user_over_budget( + None, + KeyBudgets { user: Some(budget(3.0, 3.0)), ..KeyBudgets::default() }, + "ExceededBudget: User=user-1 over budget. Spend=3.0, Budget=3.0" +)] +#[case::member_before_key( + Some("team-1"), + KeyBudgets { + key: budget(10.0, 10.0), + team_member: Some(MemberSpend { spend: 2.0, max_budget: Some(1.0), default_max_budget: None }), + ..KeyBudgets::default() + }, + "Budget has been exceeded! TeamMember=user-1:team-1 Current cost: 2.0, Max budget: 1.0" +)] +#[case::zero_member_budget( + Some("team-1"), + KeyBudgets { + team_member: Some(MemberSpend { spend: 0.0, max_budget: Some(0.0), default_max_budget: None }), + ..KeyBudgets::default() + }, + "Budget has been exceeded! User=user-1 in Team=team-1 Current cost: 0.0, Max budget: 0.0" +)] +#[case::team_default_member_budget( + Some("team-1"), + KeyBudgets { + team_member: Some(MemberSpend { spend: 4.0, max_budget: None, default_max_budget: Some(4.0) }), + ..KeyBudgets::default() + }, + "Budget has been exceeded! TeamMember=user-1:team-1 Current cost: 4.0, Max budget: 4.0" +)] +#[tokio::test] +async fn over_budget_requests_are_refused_like_the_proxy( + #[case] team_id: Option<&str>, + #[case] budgets: KeyBudgets, + #[case] message: &str, +) { + let gateway = gateway(vec![("sk-key", key(team_id, budgets))], None, "").await; + let (status, body) = send(&gateway.app, "sk-key", "claude").await; + assert_eq!(status, 422); + assert_eq!( + body, + json!({"error": {"message": message, "type": "budget_exceeded", "param": null, "code": "422"}}) + ); + assert_eq!(gateway.calls.load(Ordering::SeqCst), 0); +} + +#[rstest] +#[case::team_at_budget( + Some("team-1"), + KeyBudgets { team: Some(budget(5.0, 5.0)), ..KeyBudgets::default() } +)] +#[case::user_budget_on_team_key( + Some("team-1"), + KeyBudgets { user: Some(budget(3.0, 3.0)), ..KeyBudgets::default() } +)] +#[tokio::test] +async fn budgets_the_proxy_does_not_apply_pass( + #[case] team_id: Option<&str>, + #[case] budgets: KeyBudgets, +) { + // A team exactly at its budget passes the proxy's check; only a reservation would refuse it. + let gateway = gateway( + vec![("sk-key", key(team_id, budgets))], + None, + " disable_budget_reservation: true\n", + ) + .await; + assert_eq!(send(&gateway.app, "sk-key", "claude").await.0, 200); +} + +#[rstest] +#[tokio::test] +async fn free_models_skip_budgets_and_the_status_code_is_configurable() { + let budgets = KeyBudgets { + key: budget(10.0, 10.0), + ..KeyBudgets::default() + }; + let gateway = gateway( + vec![("sk-key", key(None, budgets))], + None, + "litellm_settings:\n budget_exceeded_status_code: 429\n", + ) + .await; + assert_eq!(send(&gateway.app, "sk-key", "free").await.0, 200); + assert_eq!(send(&gateway.app, "sk-key", "claude").await.0, 429); +} + +/// Spends a key with a 0.02 budget down: the first request fits, the second's reservation +/// shrinks to what is left, and the third finds the budget spent. +async fn spend_down(app: &Router) { + assert_eq!(send(app, "sk-key", "claude").await.0, 200); + settled().await; + assert_eq!(send(app, "sk-key", "claude").await.0, 200); + settled().await; + let (status, body) = send(app, "sk-key", "claude").await; + assert_eq!(status, 422, "{body}"); + let message = body["error"]["message"].as_str().unwrap(); + assert!( + message.starts_with("Budget has been exceeded! Key=prod Current cost: 0.02"), + "{message}" + ); +} + +fn budgeted_key() -> KeyStatus { + key( + None, + KeyBudgets { + key: budget(0.0, 0.02), + ..KeyBudgets::default() + }, + ) +} + +#[rstest] +#[tokio::test] +async fn spend_is_counted_on_the_replica_without_redis() { + let gateway = gateway(vec![("sk-key", budgeted_key())], None, "").await; + spend_down(&gateway.app).await; + assert_eq!(gateway.calls.load(Ordering::SeqCst), 2); +} + +#[rstest] +#[tokio::test] +async fn reservations_hold_back_concurrent_requests() { + // The first request's reservation takes the whole budget, so the second, sent before the + // first settles, finds it spent. + let key = key( + None, + KeyBudgets { + key: budget(0.0, 0.0001), + ..KeyBudgets::default() + }, + ); + let gateway = gateway(vec![("sk-key", key)], None, "").await; + let (first, second) = tokio::join!( + send(&gateway.app, "sk-key", "claude"), + send(&gateway.app, "sk-key", "claude") + ); + let mut statuses = [first.0, second.0]; + statuses.sort_unstable(); + assert_eq!(statuses, [200, 422]); +} + +fn redis_url() -> Option { + std::env::var("LITELLM_TEST_REDIS_URL").ok() +} + +async fn redis(url: &str, namespace: &str) -> Arc { + let coordination = + Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap(); + let router: Arc = Arc::new(NoRouter( + litellm_gateway_inference::ModelRouter::from_config( + &Config::from_yaml("model_list: []").unwrap(), + ) + .unwrap(), + )); + coordination.start(router, None, tokio_util::sync::CancellationToken::new()); + for _ in 0..100 { + if coordination.connected() { + return coordination; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("Redis did not connect"); +} + +struct NoRouter(litellm_gateway_inference::ModelRouter); + +impl litellm_gateway::coordination::RouterHandle for NoRouter { + fn router(&self) -> &litellm_gateway_inference::ModelRouter { + &self.0 + } +} + +#[rstest] +#[tokio::test] +async fn replicas_share_spend_through_the_proxy_counters() { + let Some(url) = redis_url() else { return }; + let namespace = format!("budgets-{}", uuid::Uuid::new_v4()); + let first = gateway( + vec![("sk-key", budgeted_key())], + Some(redis(&url, &namespace).await), + "", + ) + .await; + let second = gateway( + vec![("sk-key", budgeted_key())], + Some(redis(&url, &namespace).await), + "", + ) + .await; + assert_eq!(send(&first.app, "sk-key", "claude").await.0, 200); + settled().await; + assert_eq!(send(&second.app, "sk-key", "claude").await.0, 200); + settled().await; + assert_eq!(send(&first.app, "sk-key", "claude").await.0, 422); + + let mut connection = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let counter = format!("{namespace}:spend:key:{}", hash("sk-key").as_str()); + let spend: String = redis::cmd("GET") + .arg(&counter) + .query_async(&mut connection) + .await + .unwrap(); + assert!( + (spend.parse::().unwrap() - 2.0 * COST).abs() < 1e-9, + "{spend}" + ); + let ttl: i64 = redis::cmd("TTL") + .arg(&counter) + .query_async(&mut connection) + .await + .unwrap(); + assert!((1..=60).contains(&ttl), "{ttl}"); + let user: String = redis::cmd("GET") + .arg(format!("{namespace}:spend:user:user-1")) + .query_async(&mut connection) + .await + .unwrap(); + assert!( + (user.parse::().unwrap() - 2.0 * COST).abs() < 1e-9, + "{user}" + ); +} + +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows; the lookups under test use the checked macro" +)] +async fn end_users(pool: &litellm_db::PgPool) { + sqlx::raw_sql( + r#" +INSERT INTO "LiteLLM_BudgetTable" (budget_id, max_budget, created_by, updated_by) + VALUES ('small', 1.0, 'x', 'x'), ('default', 0.0001, 'x', 'x'); +INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id) + VALUES ('over', 1.5, 'small'), ('at', 1.0, 'small'), ('unbudgeted', 9.0, NULL); +"#, + ) + .execute(pool) + .await + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn end_users_are_held_to_their_budgets() { + let postgres = litellm_db_testing::MigratedPostgres::start().await.unwrap(); + end_users(postgres.pool()).await; + let key = key(None, KeyBudgets::default()); + let gateway = gateway_with_database( + vec![("sk-key", key)], + None, + Some(postgres.pool().clone()), + " disable_budget_reservation: true\nlitellm_settings:\n max_end_user_budget_id: default\n", + ) + .await; + let (status, body) = send_as(&gateway.app, "sk-key", "claude", Some("over")).await; + assert_eq!(status, 422); + assert_eq!( + body["error"]["message"], + "ExceededBudget: End User=over over budget. Spend=1.5, Budget=1.0" + ); + // Strictly over: an end user exactly at budget passes the check. + assert_eq!( + send_as(&gateway.app, "sk-key", "claude", Some("at")) + .await + .0, + 200 + ); + // Without a row, only a reservation applies the default budget, and reservations are off. + assert_eq!( + send_as(&gateway.app, "sk-key", "claude", Some("new")) + .await + .0, + 200 + ); + // A row without a budget takes the default one. + let (status, body) = send_as(&gateway.app, "sk-key", "claude", Some("unbudgeted")).await; + assert_eq!(status, 422); + assert_eq!( + body["error"]["message"], + "ExceededBudget: End User=unbudgeted over budget. Spend=9.0, Budget=0.0001" + ); +} + +#[rstest] +fn settings_follow_the_proxy_defaults() { + let settings = BudgetSettings::from_config(&Config::from_yaml("model_list: []").unwrap()); + assert_eq!(settings.status, 422); + assert!(settings.reservations); + assert_eq!(settings.ttl, Duration::from_secs(60)); +} diff --git a/litellm-rust/docs/customer-proxy/budgets.svg b/litellm-rust/docs/customer-proxy/budgets.svg new file mode 100644 index 00000000000..3ff96dd5c81 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/budgets.svg @@ -0,0 +1,51 @@ + + + + + + + Budgets hook: proxy checks on the proxy's shared spend counters + + + Key store row (cached with the key) + key spend + max_budget (or budget row's) + team, user spend + budgets + team member spend, member budget + (+ active temp increase, or team default) + end users: own lookup, cached 60 s + + + One Redis script per request (atomic) + 1. read spend:key / team / user / end_user / team_member counters; missing or below + the stored spend: seed or raise to it (the proxy's reseed + set-max repair) + 2. checks in the proxy's order: member, key, team (>), user, zero member, end user (>) + 3. reserve the max-cost estimate on every counter with a positive budget; + shrink to what is left, refuse when nothing is (1e-12), as the proxy + + + + Refusal: 422 (budget_exceeded_status_code) + {"error": {"message", "type": "budget_exceeded", + "param": null, "code": "422"}} on both endpoints + proxy messages, Python float repr; no provider call + + + + Admitted: request runs + estimate = input tokens x input price + + min(max_tokens or 16384, model max) x output price + lease renews reserved counters every TTL/2 + + + + Settle (off the request path) + success: INCRBYFLOAT (actual - reserved) on reserved counters, actual on the others; failure: release the reservation + a counter that expired mid-request is reseeded and charged the whole cost; EXPIRE ttl (default_redis_ttl, 60 s) refreshed + script failure: the counters are deleted so the next read reseeds them, as the proxy does + + + + Same keys, values and TTLs as the Python proxy, under its namespace + gateway replicas and Python proxies on one Redis enforce one budget; the proxy's reset job deletes counters and they reseed. + No Redis, or Redis unreachable: the same logic on this replica's counters (fail_closed_budget_enforcement: 503 instead) + From 0919823385f46bf55baa74910c120abd67e2bf3c Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:44 -0700 Subject: [PATCH 04/10] feat(gateway): enforce key, user, team, member and end-user rate limits on shared counters (#47050) * feat(gateway): enforce key, user, team, member and end-user rate limits on shared counters Adds a built-in rate limits hook with the proxy's v3 limiter semantics: requests per window, tokens per window and in-flight requests per key, on the proxy's own {scope:id}:window|requests|tokens|max_parallel_requests Redis keys and with its scripts, so replicas and Python proxies sharing the Redis enforce one limit. Tokens are reserved from the proxy's estimate and settled to usage when the response ends. Refusals carry the proxy's message, body and headers on each endpoint, and complete responses carry its x-ratelimit-* headers. Without Redis the counters are kept per replica. The key store now reads key, user, team and member rpm/tpm limits and the key's max_parallel_requests, and end users' budget limits are read with their budgets; keys with plain rate limits are no longer refused. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * test(gateway-auth): expect the narrowed unsupported rate limit message Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 --------- Co-authored-by: Claude Opus 5.5 (1M context) --- ...ceebb7cf0688affae913dd9d25bddcd2a4d6.json} | 30 +- ...598589a6afcbd4fb2cf1c38c6f1fa46d7773.json} | 172 ++- litellm-rust/crates/db/src/end_user.rs | 23 +- litellm-rust/crates/db/src/lib.rs | 3 +- .../crates/db/src/verification_token.rs | 64 +- .../crates/db/tests/verification_token.rs | 91 +- litellm-rust/crates/gateway-auth/src/error.rs | 2 +- litellm-rust/crates/gateway-auth/src/keys.rs | 30 +- .../crates/gateway-hooks/src/facts.rs | 3 +- litellm-rust/crates/gateway-hooks/src/hook.rs | 10 + .../crates/gateway-inference/src/error.rs | 18 +- .../crates/gateway-inference/src/hooks.rs | 1 + litellm-rust/crates/gateway/src/budgets.rs | 54 +- litellm-rust/crates/gateway/src/end_users.rs | 70 + litellm-rust/crates/gateway/src/key_status.rs | 100 +- litellm-rust/crates/gateway/src/lib.rs | 2 + litellm-rust/crates/gateway/src/main.rs | 7 + .../crates/gateway/src/rate_limits.rs | 1193 +++++++++++++++++ litellm-rust/crates/gateway/src/spend.rs | 2 +- .../crates/gateway/tests/rate_limits.rs | 515 +++++++ .../docs/customer-proxy/rate-limits.svg | 55 + 21 files changed, 2293 insertions(+), 152 deletions(-) rename litellm-rust/crates/db/.sqlx/{query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json => query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json} (55%) rename litellm-rust/crates/db/.sqlx/{query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json => query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json} (64%) create mode 100644 litellm-rust/crates/gateway/src/end_users.rs create mode 100644 litellm-rust/crates/gateway/src/rate_limits.rs create mode 100644 litellm-rust/crates/gateway/tests/rate_limits.rs create mode 100644 litellm-rust/docs/customer-proxy/rate-limits.svg diff --git a/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json b/litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json similarity index 55% rename from litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json rename to litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json index c9805198db5..de185711b56 100644 --- a/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json +++ b/litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n (SELECT d.max_budget FROM \"LiteLLM_BudgetTable\" AS d\n WHERE d.budget_id = $2) AS \"default_max_budget?\"\n FROM (SELECT 1) AS one\n LEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\n LEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = e.budget_id", + "query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n b.rpm_limit AS \"rpm_limit?\",\n b.tpm_limit AS \"tpm_limit?\"\n FROM (SELECT 1) AS one\n LEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\n LEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = COALESCE($3, e.budget_id, $2)", "describe": { "columns": [ { @@ -38,13 +38,30 @@ }, { "ordinal": 3, - "name": "default_max_budget?", - "type_info": "Float8", - "origin": "Expression" + "name": "rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 4, + "name": "tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "tpm_limit" + } + } } ], "parameters": { "Left": [ + "Text", "Text", "Text" ] @@ -53,8 +70,9 @@ true, true, true, - null + true, + true ] }, - "hash": "a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe" + "hash": "623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6" } diff --git a/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json b/litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json similarity index 64% rename from litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json rename to litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json index 1a9ad366cfa..d8f743ae315 100644 --- a/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json +++ b/litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n COALESCE(v.max_budget, kb.max_budget) AS max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL\n OR v.max_parallel_requests IS NOT NULL\n OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL\n OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL\n OR t.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL\n OR mb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n tm.spend AS \"member_spend?\",\n mb.max_budget AS \"member_max_budget?\",\n CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC')\n THEN mb.temp_budget_increase END AS \"member_budget_increase?\",\n tdb.max_budget AS \"member_default_budget?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n u.spend AS \"user_spend?\",\n u.max_budget AS \"user_max_budget?\",\n (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL\n OR u.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS tdb\n ON tdb.budget_id = t.metadata ->> 'team_member_budget_id'\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", + "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n COALESCE(v.max_budget, kb.max_budget) AS max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n COALESCE(v.rpm_limit, kb.rpm_limit) AS rpm_limit,\n COALESCE(v.tpm_limit, kb.tpm_limit) AS tpm_limit,\n v.max_parallel_requests,\n (COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n t.rpm_limit AS \"team_rpm_limit?\",\n t.tpm_limit AS \"team_tpm_limit?\",\n (COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n mb.rpm_limit AS \"member_rpm_limit?\",\n mb.tpm_limit AS \"member_tpm_limit?\",\n (COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n tm.spend AS \"member_spend?\",\n mb.max_budget AS \"member_max_budget?\",\n CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC')\n THEN mb.temp_budget_increase END AS \"member_budget_increase?\",\n tdb.max_budget AS \"member_default_budget?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n u.spend AS \"user_spend?\",\n u.max_budget AS \"user_max_budget?\",\n u.rpm_limit AS \"user_rpm_limit?\",\n u.tpm_limit AS \"user_tpm_limit?\",\n (COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS tdb\n ON tdb.budget_id = t.metadata ->> 'team_member_budget_id'\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", "describe": { "columns": [ { @@ -105,54 +105,77 @@ }, { "ordinal": 10, + "name": "rpm_limit", + "type_info": "Int8", + "origin": "Expression" + }, + { + "ordinal": 11, + "name": "tpm_limit", + "type_info": "Int8", + "origin": "Expression" + }, + { + "ordinal": 12, + "name": "max_parallel_requests", + "type_info": "Int4", + "origin": { + "Table": { + "table": "\"LiteLLM_VerificationToken\"", + "name": "max_parallel_requests" + } + } + }, + { + "ordinal": 13, "name": "key_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 11, + "ordinal": 14, "name": "key_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 12, + "ordinal": 15, "name": "key_model_aliases!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 13, + "ordinal": 16, "name": "key_object_permissions!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 14, + "ordinal": 17, "name": "key_cache_controls!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 15, + "ordinal": 18, "name": "key_project!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 16, + "ordinal": 19, "name": "key_organization!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 17, + "ordinal": 20, "name": "key_agent!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 18, + "ordinal": 21, "name": "team_id", "type_info": "Text", "origin": { @@ -163,7 +186,7 @@ } }, { - "ordinal": 19, + "ordinal": 22, "name": "team_row_id?", "type_info": "Text", "origin": { @@ -174,7 +197,7 @@ } }, { - "ordinal": 20, + "ordinal": 23, "name": "team_alias?", "type_info": "Text", "origin": { @@ -185,7 +208,7 @@ } }, { - "ordinal": 21, + "ordinal": 24, "name": "team_blocked?", "type_info": "Bool", "origin": { @@ -196,7 +219,7 @@ } }, { - "ordinal": 22, + "ordinal": 25, "name": "team_models?", "type_info": "TextArray", "origin": { @@ -207,7 +230,7 @@ } }, { - "ordinal": 23, + "ordinal": 26, "name": "team_spend?", "type_info": "Float8", "origin": { @@ -218,7 +241,7 @@ } }, { - "ordinal": 24, + "ordinal": 27, "name": "team_max_budget?", "type_info": "Float8", "origin": { @@ -229,7 +252,7 @@ } }, { - "ordinal": 25, + "ordinal": 28, "name": "team_metadata?", "type_info": "Jsonb", "origin": { @@ -240,43 +263,87 @@ } }, { - "ordinal": 26, + "ordinal": 29, + "name": "team_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 30, + "name": "team_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 31, "name": "team_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 27, + "ordinal": 32, "name": "team_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 28, + "ordinal": 33, "name": "team_model_aliases!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 29, + "ordinal": 34, "name": "team_object_permissions!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 30, + "ordinal": 35, "name": "team_organization!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 31, + "ordinal": 36, + "name": "member_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 37, + "name": "member_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 38, "name": "member_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 32, + "ordinal": 39, "name": "member_allowed_models?", "type_info": "TextArray", "origin": { @@ -287,7 +354,7 @@ } }, { - "ordinal": 33, + "ordinal": 40, "name": "member_spend?", "type_info": "Float8", "origin": { @@ -298,7 +365,7 @@ } }, { - "ordinal": 34, + "ordinal": 41, "name": "member_max_budget?", "type_info": "Float8", "origin": { @@ -309,13 +376,13 @@ } }, { - "ordinal": 35, + "ordinal": 42, "name": "member_budget_increase?", "type_info": "Float8", "origin": "Expression" }, { - "ordinal": 36, + "ordinal": 43, "name": "member_default_budget?", "type_info": "Float8", "origin": { @@ -326,7 +393,7 @@ } }, { - "ordinal": 37, + "ordinal": 44, "name": "user_id", "type_info": "Text", "origin": { @@ -337,7 +404,7 @@ } }, { - "ordinal": 38, + "ordinal": 45, "name": "user_row_id?", "type_info": "Text", "origin": { @@ -348,7 +415,7 @@ } }, { - "ordinal": 39, + "ordinal": 46, "name": "user_models?", "type_info": "TextArray", "origin": { @@ -359,7 +426,7 @@ } }, { - "ordinal": 40, + "ordinal": 47, "name": "user_spend?", "type_info": "Float8", "origin": { @@ -370,7 +437,7 @@ } }, { - "ordinal": 41, + "ordinal": 48, "name": "user_max_budget?", "type_info": "Float8", "origin": { @@ -381,19 +448,41 @@ } }, { - "ordinal": 42, + "ordinal": 49, + "name": "user_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 50, + "name": "user_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 51, "name": "user_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 43, + "ordinal": 52, "name": "user_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 44, + "ordinal": 53, "name": "user_object_permissions!", "type_info": "Bool", "origin": "Expression" @@ -417,6 +506,9 @@ false, null, null, + true, + null, + null, null, null, null, @@ -431,7 +523,8 @@ true, true, true, - null, + true, + true, null, null, null, @@ -439,6 +532,9 @@ null, true, true, + null, + true, + true, true, null, true, @@ -447,10 +543,12 @@ true, true, true, + true, + true, null, null, null ] }, - "hash": "08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9" + "hash": "77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773" } diff --git a/litellm-rust/crates/db/src/end_user.rs b/litellm-rust/crates/db/src/end_user.rs index a4d98eef2aa..4f7d05a32ec 100644 --- a/litellm-rust/crates/db/src/end_user.rs +++ b/litellm-rust/crates/db/src/end_user.rs @@ -1,6 +1,6 @@ use sqlx::PgPool; -use crate::Error; +use crate::{Error, RateLimits}; /// An end user's spend and budget, as the proxy's end-user budget checks read them. #[derive(Clone, Debug, Default, PartialEq)] @@ -8,33 +8,42 @@ pub struct EndUser { /// `LiteLLM_EndUserTable.spend`, when the end user has a row. pub spend: Option, pub blocked: bool, - /// The row's budget, else the default budget's. + /// The budget's `max_budget`: the override's, else the row's, else the default's. pub max_budget: Option, + /// That budget's `rpm_limit` and `tpm_limit`. + pub limits: RateLimits, } -/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own. +/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own +/// and `override_budget_id`, when given, as its budget whatever the row says. pub async fn find_end_user( pool: &PgPool, end_user_id: &str, default_budget_id: Option<&str>, + override_budget_id: Option<&str>, ) -> Result { let row = sqlx::query!( r#"SELECT e.spend AS "spend?", e.blocked AS "blocked?", b.max_budget AS "max_budget?", - (SELECT d.max_budget FROM "LiteLLM_BudgetTable" AS d - WHERE d.budget_id = $2) AS "default_max_budget?" + b.rpm_limit AS "rpm_limit?", + b.tpm_limit AS "tpm_limit?" FROM (SELECT 1) AS one LEFT JOIN "LiteLLM_EndUserTable" AS e ON e.user_id = $1 - LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = e.budget_id"#, + LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = COALESCE($3, e.budget_id, $2)"#, end_user_id, default_budget_id, + override_budget_id, ) .fetch_one(pool) .await?; Ok(EndUser { spend: row.spend, blocked: row.blocked.unwrap_or(false), - max_budget: row.max_budget.or(row.default_max_budget), + max_budget: row.max_budget, + limits: RateLimits { + rpm: row.rpm_limit, + tpm: row.tpm_limit, + }, }) } diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index b41ee5ace78..20d7d7616bb 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -16,5 +16,6 @@ pub use spend::{ }; pub use sqlx::PgPool; pub use verification_token::{ - Restrictions, Team, TeamMembership, User, VerificationToken, find_verification_token, + RateLimits, Restrictions, Team, TeamMembership, User, VerificationToken, + find_verification_token, }; diff --git a/litellm-rust/crates/db/src/verification_token.rs b/litellm-rust/crates/db/src/verification_token.rs index 3ff008b0cff..5db37ae5c9b 100644 --- a/litellm-rust/crates/db/src/verification_token.rs +++ b/litellm-rust/crates/db/src/verification_token.rs @@ -19,6 +19,9 @@ pub struct VerificationToken { pub key_alias: Option, pub key_name: Option, pub metadata: Value, + /// The key's `rpm_limit` and `tpm_limit`, else its budget table's. + pub limits: RateLimits, + pub max_parallel_requests: Option, pub restrictions: Restrictions, pub team_id: Option, /// `None` when the key has no team or the team row is missing. @@ -38,9 +41,17 @@ pub struct Team { pub spend: f64, pub max_budget: Option, pub metadata: Value, + pub limits: RateLimits, pub restrictions: Restrictions, } +/// Requests and tokens per rate limit window. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct RateLimits { + pub rpm: Option, + pub tpm: Option, +} + /// The key's team membership: its spend and budget, and what its budget table sets beyond them. #[derive(Clone, Debug, Default, PartialEq)] pub struct TeamMembership { @@ -50,7 +61,9 @@ pub struct TeamMembership { pub max_budget: Option, /// The team's default member budget, `metadata.team_member_budget_id`. pub default_max_budget: Option, - /// `rpm_limit`, `tpm_limit`, `tpd_limit` or `max_parallel_requests`. + /// The membership budget table's `rpm_limit` and `tpm_limit`. + pub limits: RateLimits, + /// An `rpm_limit` or `tpm_limit` inside the budget table's `model_max_budget`. pub rate_limits: bool, /// The member's own model scope; empty inherits the team's. pub allowed_models: Vec, @@ -61,6 +74,7 @@ pub struct User { pub spend: f64, pub max_budget: Option, pub models: Vec, + pub limits: RateLimits, pub restrictions: Restrictions, } @@ -68,8 +82,8 @@ pub struct User { /// `allowed_routes`. #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct Restrictions { - /// `rpm_limit`, `tpm_limit`, `tpd_limit` or `max_parallel_requests`, on the row or on a key's - /// budget table, or an `rpm_limit` or `tpm_limit` inside `model_max_budget`. + /// Per-model limits: an `rpm_limit` or `tpm_limit` inside `model_max_budget`, on the row or + /// on a key's budget table. pub rate_limits: bool, pub policies: bool, /// `aliases`, `router_settings`, or a team's `model_id`. @@ -98,11 +112,10 @@ pub async fn find_verification_token( v.key_alias, v.key_name, v.metadata, - (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL - OR v.max_parallel_requests IS NOT NULL - OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL - OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(v.model_max_budget, + COALESCE(v.rpm_limit, kb.rpm_limit) AS rpm_limit, + COALESCE(v.tpm_limit, kb.tpm_limit) AS tpm_limit, + v.max_parallel_requests, + (COALESCE(jsonb_path_exists(v.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false) OR COALESCE(jsonb_path_exists(kb.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) @@ -123,9 +136,9 @@ pub async fn find_verification_token( t.spend AS "team_spend?", t.max_budget AS "team_max_budget?", t.metadata AS "team_metadata?", - (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL - OR t.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(t.model_max_budget, + t.rpm_limit AS "team_rpm_limit?", + t.tpm_limit AS "team_tpm_limit?", + (COALESCE(jsonb_path_exists(t.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "team_rate_limits!", COALESCE(cardinality(t.policies) > 0, false) AS "team_policies!", @@ -133,9 +146,9 @@ pub async fn find_verification_token( OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS "team_model_aliases!", (t.object_permission_id IS NOT NULL) AS "team_object_permissions!", (t.organization_id IS NOT NULL) AS "team_organization!", - (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL - OR mb.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(mb.model_max_budget, + mb.rpm_limit AS "member_rpm_limit?", + mb.tpm_limit AS "member_tpm_limit?", + (COALESCE(jsonb_path_exists(mb.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "member_rate_limits!", mb.allowed_models AS "member_allowed_models?", @@ -149,9 +162,9 @@ pub async fn find_verification_token( u.models AS "user_models?", u.spend AS "user_spend?", u.max_budget AS "user_max_budget?", - (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL - OR u.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(u.model_max_budget, + u.rpm_limit AS "user_rpm_limit?", + u.tpm_limit AS "user_tpm_limit?", + (COALESCE(jsonb_path_exists(u.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "user_rate_limits!", COALESCE(cardinality(u.policies) > 0, false) AS "user_policies!", @@ -181,6 +194,11 @@ pub async fn find_verification_token( key_alias: row.key_alias, key_name: row.key_name, metadata: row.metadata, + limits: RateLimits { + rpm: row.rpm_limit, + tpm: row.tpm_limit, + }, + max_parallel_requests: row.max_parallel_requests.map(i64::from), restrictions: Restrictions { rate_limits: row.key_rate_limits, policies: row.key_policies, @@ -199,6 +217,10 @@ pub async fn find_verification_token( spend: row.team_spend.unwrap_or(0.0), max_budget: row.team_max_budget, metadata: row.team_metadata.unwrap_or(Value::Null), + limits: RateLimits { + rpm: row.team_rpm_limit, + tpm: row.team_tpm_limit, + }, restrictions: Restrictions { rate_limits: row.team_rate_limits, policies: row.team_policies, @@ -214,6 +236,10 @@ pub async fn find_verification_token( .member_max_budget .map(|budget| budget + row.member_budget_increase.unwrap_or(0.0)), default_max_budget: row.member_default_budget, + limits: RateLimits { + rpm: row.member_rpm_limit, + tpm: row.member_tpm_limit, + }, rate_limits: row.member_rate_limits, allowed_models: row.member_allowed_models.unwrap_or_default(), }, @@ -222,6 +248,10 @@ pub async fn find_verification_token( spend: row.user_spend.unwrap_or(0.0), max_budget: row.user_max_budget, models: row.user_models.unwrap_or_default(), + limits: RateLimits { + rpm: row.user_rpm_limit, + tpm: row.user_tpm_limit, + }, restrictions: Restrictions { rate_limits: row.user_rate_limits, policies: row.user_policies, diff --git a/litellm-rust/crates/db/tests/verification_token.rs b/litellm-rust/crates/db/tests/verification_token.rs index ffe0b4d66f6..fc4cb83e04f 100644 --- a/litellm-rust/crates/db/tests/verification_token.rs +++ b/litellm-rust/crates/db/tests/verification_token.rs @@ -1,8 +1,8 @@ use std::time::{Duration, SystemTime}; use litellm_db::{ - EndUser, Restrictions, Team, TeamMembership, User, VerificationToken, find_end_user, - find_verification_token, + EndUser, RateLimits, Restrictions, Team, TeamMembership, User, VerificationToken, + find_end_user, find_verification_token, }; use litellm_db_testing::MigratedPostgres; use rstest::rstest; @@ -93,6 +93,8 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { key_alias: None, key_name: None, metadata: json!({}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: None, team: None, @@ -115,6 +117,8 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { key_alias: Some("prod".into()), key_name: Some("sk-...abcd".into()), metadata: json!({"tags": ["a"]}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: Some("team-1".into()), team: Some(Team { @@ -124,16 +128,21 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: 2.5, max_budget: Some(10.0), metadata: json!({"require_trace_id": true}), - restrictions: Restrictions { - rate_limits: true, - ..Restrictions::default() + limits: RateLimits { + rpm: None, + tpm: Some(100), }, + restrictions: Restrictions::default(), }), team_membership: TeamMembership { spend: Some(0.0), max_budget: None, default_max_budget: None, - rate_limits: true, + limits: RateLimits { + rpm: Some(3), + tpm: None, + }, + rate_limits: false, allowed_models: vec![], }, user_id: Some("user-1".into()), @@ -141,8 +150,11 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: 0.0, max_budget: Some(1.0), models: vec!["gpt-b".into()], + limits: RateLimits { + rpm: Some(5), + tpm: None, + }, restrictions: Restrictions { - rate_limits: true, policies: true, ..Restrictions::default() }, @@ -152,7 +164,7 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { assert_eq!( find("restricted").await.restrictions, Restrictions { - rate_limits: true, + rate_limits: false, policies: true, model_aliases: true, object_permissions: true, @@ -162,16 +174,33 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { agent: true, } ); + let limited = find("limited-only").await; assert_eq!( - find("limited-only").await.restrictions, - Restrictions { - rate_limits: true, - ..Restrictions::default() + (limited.max_parallel_requests, limited.restrictions), + (Some(2), Restrictions::default()) + ); + // A key's own limits win; its budget row's fill the gaps. + assert_eq!( + find("restricted").await.limits, + RateLimits { + rpm: Some(10), + tpm: None } ); - // Budgets are not restrictions; limits held in a budget row or per model are. + let budget_row = find("budget-row-limits").await; + assert_eq!( + (budget_row.limits, budget_row.restrictions), + ( + RateLimits { + rpm: Some(3), + tpm: None + }, + Restrictions::default() + ) + ); + // Budgets and plain limits are not restrictions; limits per model are. assert_eq!(find("budgeted").await.restrictions, Restrictions::default()); - for token in ["budget-row-limits", "budget-row-model-limits", "model-rpm"] { + for token in ["budget-row-model-limits", "model-rpm"] { assert_eq!( find(token).await.restrictions, Restrictions { @@ -188,6 +217,7 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: Some(0.75), max_budget: Some(5.0), default_max_budget: None, + limits: RateLimits::default(), rate_limits: false, allowed_models: vec!["gpt-a".into()], } @@ -216,15 +246,17 @@ async fn reads_end_user_budgets_with_their_default() { let postgres = MigratedPostgres::start().await.unwrap(); let pool = postgres.pool(); setup(pool).await; - let find = - async |id: &str, default: Option<&str>| find_end_user(pool, id, default).await.unwrap(); + let find = async |id: &str, default: Option<&str>| { + find_end_user(pool, id, default, None).await.unwrap() + }; assert_eq!( find("end-1", Some("b-2")).await, EndUser { spend: Some(1.5), blocked: false, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); assert_eq!( @@ -232,16 +264,37 @@ async fn reads_end_user_budgets_with_their_default() { EndUser { spend: Some(0.25), blocked: true, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); assert_eq!(find("end-2", None).await.max_budget, None); + // An override replaces the row's own budget. + assert_eq!( + find_end_user(pool, "end-1", None, Some("b-2")) + .await + .unwrap() + .limits, + RateLimits { + rpm: Some(3), + tpm: None + } + ); + // A row without its own budget takes the default's limits too. + assert_eq!( + find("end-2", Some("b-2")).await.limits, + RateLimits { + rpm: Some(3), + tpm: None + } + ); assert_eq!( find("unknown", Some("b-1")).await, EndUser { spend: None, blocked: false, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); } diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs index 6310e2eda90..30494a7e1a3 100644 --- a/litellm-rust/crates/gateway-auth/src/error.rs +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -269,7 +269,7 @@ mod tests { restriction: Restriction::RateLimits, owner: Owner::TeamMembership, }), - "Authentication Error, keys with rate limits set on its team membership are not supported by this gateway yet." + "Authentication Error, keys with per-model, per-tag or dynamic rate limits set on its team membership are not supported by this gateway yet." )] #[case::banned_param( Error::BannedParam("aws_bedrock_runtime_endpoint"), diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index ff99572a1ee..654b8503d3e 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -241,7 +241,7 @@ pub enum Refusal { #[derive(Clone, Copy, Debug, Eq, PartialEq, strum::Display, strum::VariantArray)] pub enum Restriction { - #[strum(serialize = "rate limits")] + #[strum(serialize = "per-model, per-tag or dynamic rate limits")] RateLimits, #[strum(serialize = "guardrails or policies")] Guardrails, @@ -298,6 +298,34 @@ pub struct KeyProfile { pub metadata: Arc>, pub team_metadata: Arc>, pub budgets: KeyBudgets, + pub limits: KeyLimits, +} + +/// Requests and tokens per rate limit window. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Limits { + pub rpm: Option, + pub tpm: Option, +} + +impl Limits { + pub fn is_set(&self) -> bool { + self.rpm.is_some() || self.tpm.is_some() + } +} + +/// The rate limits of a key and of who it belongs to, read with the key and refreshed with it. +/// End users' limits are read on their own. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct KeyLimits { + /// The key's own limits, else its budget table's. + pub key: Limits, + /// Requests the key may have in flight at once. + pub max_parallel_requests: Option, + pub user: Limits, + pub team: Limits, + /// The key's user within its team, from the membership's budget table. + pub team_member: Limits, } /// Stored spend against an optional budget, as the key store last read it. diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs index 437768707d0..d696427fd8d 100644 --- a/litellm-rust/crates/gateway-hooks/src/facts.rs +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -8,7 +8,7 @@ use bytes::Bytes; use http::HeaderMap; use serde_json::{Map, Value}; -use litellm_gateway_auth::keys::KeyBudgets; +use litellm_gateway_auth::keys::{KeyBudgets, KeyLimits}; use crate::usage::Usage; @@ -104,6 +104,7 @@ pub struct KeyFacts { pub metadata: Arc>, pub team_metadata: Arc>, pub budgets: KeyBudgets, + pub limits: KeyLimits, } /// What the gateway knows about a request before it is sent, shared by every hook and observer. diff --git a/litellm-rust/crates/gateway-hooks/src/hook.rs b/litellm-rust/crates/gateway-hooks/src/hook.rs index 3e57e6b12c8..9246f77d1a2 100644 --- a/litellm-rust/crates/gateway-hooks/src/hook.rs +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -100,6 +100,8 @@ pub struct Refusal { pub param: Option<&'static str>, pub message: String, pub headers: Vec<(HeaderName, HeaderValue)>, + /// Send Anthropic's error envelope, with `kind` as its error type, instead of the proxy's. + pub anthropic: bool, } impl Refusal { @@ -114,9 +116,17 @@ impl Refusal { param: Some("None"), message: message.into(), headers: Vec::new(), + anthropic: false, } } + /// Sends `{"type": "error", "error": {"type": kind, "message": message}}`, as the proxy's + /// `/v1/messages` route does for some errors. + pub fn anthropic(mut self) -> Self { + self.anthropic = true; + self + } + /// Sends a `null` `param`, as the proxy's budget errors do. pub fn without_param(mut self) -> Self { self.param = None; diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 05009604fba..3aa4b44f9aa 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -177,12 +177,18 @@ impl Error { /// A hook's refusal, in the proxy's error shape and with the hook's headers. fn refused(refusal: Refusal) -> Response { let status = StatusCode::from_u16(refusal.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); - let body = json!({"error": { - "message": refusal.message, - "type": refusal.kind, - "param": refusal.param, - "code": status.as_u16().to_string(), - }}); + let body = match refusal.anthropic { + true => json!({"type": "error", "error": { + "type": refusal.kind, + "message": refusal.message, + }}), + false => json!({"error": { + "message": refusal.message, + "type": refusal.kind, + "param": refusal.param, + "code": status.as_u16().to_string(), + }}), + }; let mut response = (status, Json(body)).into_response(); response.headers_mut().extend(refusal.headers); response diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index e60bcd140bc..07833d8cc93 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -121,6 +121,7 @@ fn key_facts(hooks: &Hooks, identity: &AuthenticatedRequest) -> Option Some(Arc::new(KeyFacts { diff --git a/litellm-rust/crates/gateway/src/budgets.rs b/litellm-rust/crates/gateway/src/budgets.rs index be04ae028be..7196bf132da 100644 --- a/litellm-rust/crates/gateway/src/budgets.rs +++ b/litellm-rust/crates/gateway/src/budgets.rs @@ -17,22 +17,18 @@ use std::{ use futures_util::future::BoxFuture; use litellm_config::Config; -use litellm_db::{EndUser, PgPool}; +use litellm_db::PgPool; use litellm_gateway_hooks::{ Facts, Hook, Prices, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets, Usage, }; use redis::aio::ConnectionManager; -use crate::coordination::Coordination; +use crate::{coordination::Coordination, end_users::EndUsers}; /// The output tokens the proxy assumes when neither the request nor the model caps them. const DEFAULT_OUTPUT_TOKENS: u64 = 16_384; /// The proxy's slack when a reservation exactly meets a budget. const EPSILON: f64 = 1e-12; -/// How long an end user's budget is reused before it is read again. -const END_USER_TTL: Duration = Duration::from_secs(60); -/// End users held in memory; the cache starts over past this. -const MAX_END_USERS: usize = 10_000; const FAIL_CLOSED_READ: &str = "{'error': 'Budget enforcement unavailable: current spend could not be verified against Redis or the database, and fail_closed_budget_enforcement is enabled, so the request was rejected to avoid exceeding the configured budget. Retry shortly.'}"; /// Seeds or raises each counter to its floor, runs the budget checks in the order given and @@ -600,45 +596,6 @@ async fn admit_redis( } } -/// An end user and the default budget it was read with. -type EndUserKey = (String, Option); - -/// End users' budgets, read from the proxy's database and reused for a minute. -struct EndUsers { - pool: PgPool, - cache: Mutex>, -} - -impl EndUsers { - async fn get(&self, id: &str, default_budget_id: Option<&str>) -> Option { - let key = (id.to_owned(), default_budget_id.map(str::to_owned)); - let now = Instant::now(); - if let Some((end_user, read_at)) = self - .cache - .lock() - .unwrap_or_else(|e| e.into_inner()) - .get(&key) - && now.duration_since(*read_at) < END_USER_TTL - { - return Some(end_user.clone()); - } - match litellm_db::find_end_user(&self.pool, id, default_budget_id).await { - Ok(end_user) => { - let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); - if cache.len() >= MAX_END_USERS { - cache.clear(); - } - cache.insert(key, (end_user.clone(), now)); - Some(end_user) - } - Err(error) => { - tracing::warn!(%error, "reading an end user's budget failed; not checking it"); - None - } - } - } -} - /// The `budgets` hook. pub struct Budgets { settings: BudgetSettings, @@ -670,10 +627,7 @@ impl Budgets { fail_closed: settings.fail_closed, }), settings, - end_users: database.map(|pool| EndUsers { - pool, - cache: Mutex::default(), - }), + end_users: database.map(EndUsers::new), groups, prices: OnceLock::new(), } @@ -817,7 +771,7 @@ impl Budgets { .as_deref() .or(self.settings.end_user_budget_id.as_deref()); let end_user = match &self.end_users { - Some(end_users) => end_users.get(end_user_id, default_budget_id).await, + Some(end_users) => end_users.get(end_user_id, default_budget_id, None).await, None => None, } .unwrap_or_default(); diff --git a/litellm-rust/crates/gateway/src/end_users.rs b/litellm-rust/crates/gateway/src/end_users.rs new file mode 100644 index 00000000000..5c23e4f6328 --- /dev/null +++ b/litellm-rust/crates/gateway/src/end_users.rs @@ -0,0 +1,70 @@ +//! End users' budgets and limits, read from the proxy's database and reused for a minute. + +use std::{ + collections::HashMap, + sync::Mutex, + time::{Duration, Instant}, +}; + +use litellm_db::{EndUser, PgPool}; + +/// How long an end user's budget is reused before it is read again. +const END_USER_TTL: Duration = Duration::from_secs(60); +/// End users held in memory; the cache starts over past this. +const MAX_END_USERS: usize = 10_000; + +/// An end user and the default and override budgets it was read with. +type EndUserKey = (String, Option, Option); + +pub(crate) struct EndUsers { + pool: PgPool, + cache: Mutex>, +} + +impl EndUsers { + pub(crate) fn new(pool: PgPool) -> Self { + Self { + pool, + cache: Mutex::default(), + } + } + + /// `None` when the database could not be read. + pub(crate) async fn get( + &self, + id: &str, + default_budget_id: Option<&str>, + override_budget_id: Option<&str>, + ) -> Option { + let key = ( + id.to_owned(), + default_budget_id.map(str::to_owned), + override_budget_id.map(str::to_owned), + ); + let now = Instant::now(); + if let Some((end_user, read_at)) = self + .cache + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&key) + && now.duration_since(*read_at) < END_USER_TTL + { + return Some(end_user.clone()); + } + match litellm_db::find_end_user(&self.pool, id, default_budget_id, override_budget_id).await + { + Ok(end_user) => { + let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); + if cache.len() >= MAX_END_USERS { + cache.clear(); + } + cache.insert(key, (end_user.clone(), now)); + Some(end_user) + } + Err(error) => { + tracing::warn!(%error, "reading an end user's budget failed; not checking it"); + None + } + } + } +} diff --git a/litellm-rust/crates/gateway/src/key_status.rs b/litellm-rust/crates/gateway/src/key_status.rs index 50d8bf65567..20e2e795336 100644 --- a/litellm-rust/crates/gateway/src/key_status.rs +++ b/litellm-rust/crates/gateway/src/key_status.rs @@ -1,20 +1,21 @@ -use litellm_db::{Restrictions, VerificationToken}; +use litellm_db::{RateLimits, Restrictions, VerificationToken}; use std::sync::Arc; use litellm_gateway_auth::keys::{ - AllowedRoutes, KeyBudgets, KeyModels, KeyProfile, KeyStatus, MemberSpend, ModelAccess, Owner, - Refusal, Restriction, Spend, + AllowedRoutes, KeyBudgets, KeyLimits, KeyModels, KeyProfile, KeyStatus, Limits, MemberSpend, + ModelAccess, Owner, Refusal, Restriction, Spend, }; use serde_json::Value; /// Key and team `metadata` fields the proxy enforces and this gateway does not. -const METADATA_RESTRICTIONS: [(&str, Restriction); 11] = [ +const METADATA_RESTRICTIONS: [(&str, Restriction); 12] = [ ("guardrails", Restriction::Guardrails), ("policies", Restriction::Guardrails), ("tags", Restriction::Tags), ("model_rpm_limit", Restriction::RateLimits), ("model_tpm_limit", Restriction::RateLimits), ("tag_rpm_limit", Restriction::RateLimits), + ("throttle_on_budget_exceeded", Restriction::RateLimits), ("allowed_tools", Restriction::ToolAllowlists), ("enforced_params", Restriction::EnforcedParams), ("priority", Restriction::Priority), @@ -39,6 +40,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K _ => Arc::default(), }; let budgets = budgets(&row); + let limits = limits(&row); let (team_models, team_alias, team_metadata) = match row.team { Some(team) => ( Some(ModelAccess::new(team.models)), @@ -62,6 +64,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K metadata: object(row.metadata), team_metadata, budgets, + limits, }), } } @@ -96,6 +99,29 @@ fn budgets(row: &VerificationToken) -> KeyBudgets { } } +/// The rate limits the rate limits hook enforces for this key. +fn limits(row: &VerificationToken) -> KeyLimits { + let limits = |limits: RateLimits| Limits { + rpm: limits.rpm, + tpm: limits.tpm, + }; + KeyLimits { + key: limits(row.limits), + max_parallel_requests: row.max_parallel_requests, + user: row + .user + .as_ref() + .map(|user| limits(user.limits)) + .unwrap_or_default(), + team: row + .team + .as_ref() + .map(|team| limits(team.limits)) + .unwrap_or_default(), + team_member: limits(row.team_membership.limits), + } +} + fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option { let team = match (&row.team_id, &row.team) { (Some(team_id), None) => { @@ -145,7 +171,7 @@ fn unsupported( }) } -/// What the membership's budget table sets: rate limits, or an `allowed_models` list the proxy's +/// What the membership's budget table sets: per-model rate limits, or an `allowed_models` list the proxy's /// `check_team_member_model_access` can deny with (it skips an empty list, and `*` and /// `all-proxy-models` grant every model). fn membership(row: &VerificationToken) -> Option { @@ -195,6 +221,14 @@ fn metadata_restriction(metadata: &Value) -> Option { .iter() .find(|(field, _)| metadata.get(field).is_some_and(is_set)) .map(|(_, restriction)| *restriction) + .or_else(|| { + // Dynamic limits apply only while a deployment is failing, which the gateway does + // not track. + ["rpm_limit_type", "tpm_limit_type"] + .iter() + .any(|field| metadata.get(field).and_then(Value::as_str) == Some("dynamic")) + .then_some(Restriction::RateLimits) + }) } fn is_set(value: &Value) -> bool { @@ -236,6 +270,8 @@ mod tests { key_alias: None, key_name: None, metadata: json!({}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: None, team: None, @@ -253,6 +289,7 @@ mod tests { spend, max_budget, metadata: json!({}), + limits: RateLimits::default(), restrictions: Restrictions::default(), }) } @@ -297,6 +334,46 @@ mod tests { assert_eq!(status(row, false), active(true)); } + #[rstest] + fn carries_the_key_and_owner_rate_limits(row: VerificationToken) { + let limits = |rpm, tpm| RateLimits { rpm, tpm }; + let row = VerificationToken { + limits: limits(Some(10), None), + max_parallel_requests: Some(2), + team_id: Some("team-1".into()), + team: team(false, None, 0.0).map(|team| Team { + limits: limits(None, Some(1_000)), + ..team + }), + team_membership: TeamMembership { + limits: limits(Some(3), Some(30)), + ..TeamMembership::default() + }, + ..on_user(Some("team-1"), &[], NONE, row) + }; + let row = VerificationToken { + user: row.user.clone().map(|user| User { + limits: limits(Some(5), None), + ..user + }), + ..row + }; + let KeyStatus::Active { profile, .. } = status(row, false) else { + panic!("expected an active key"); + }; + let limits = |rpm, tpm| Limits { rpm, tpm }; + assert_eq!( + profile.limits, + KeyLimits { + key: limits(Some(10), None), + max_parallel_requests: Some(2), + user: limits(Some(5), None), + team: limits(None, Some(1_000)), + team_member: limits(Some(3), Some(30)), + } + ); + } + #[rstest] #[case::blocked_key(true, None, None, None, 0.0, KeyStatus::Blocked)] #[case::team_row_missing( @@ -376,6 +453,7 @@ mod tests { spend: 0.0, max_budget: None, models: strings(models), + limits: RateLimits::default(), restrictions, }), ..row @@ -473,6 +551,18 @@ mod tests { |row| key_metadata(json!({"tag_rpm_limit": {"a": 1}}), row), Some((Restriction::RateLimits, Owner::Key)) )] + #[case::metadata_dynamic_rpm( + |row| key_metadata(json!({"rpm_limit_type": "dynamic"}), row), + Some((Restriction::RateLimits, Owner::Key)) + )] + #[case::metadata_guaranteed_tpm( + |row| key_metadata(json!({"tpm_limit_type": "guaranteed_throughput"}), row), + None + )] + #[case::metadata_budget_throttle( + |row| key_metadata(json!({"throttle_on_budget_exceeded": true}), row), + Some((Restriction::RateLimits, Owner::Key)) + )] #[case::metadata_allowed_tools( |row| key_metadata(json!({"allowed_tools": ["search"]}), row), Some((Restriction::ToolAllowlists, Owner::Key)) diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index dcd3aa994a3..57f05ec9322 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -2,12 +2,14 @@ mod auth; pub mod budgets; pub mod coordination; pub mod drain; +mod end_users; mod error; pub mod health; pub mod hooks; mod key_status; mod keys; pub mod metrics; +pub mod rate_limits; 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 0954fd7634b..3e648f5a758 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -120,6 +120,13 @@ async fn main() -> Result<(), Box> { coordination.clone(), database.clone(), ))); + hooks.push(Arc::new( + litellm_gateway::rate_limits::RateLimiter::for_gateway( + &config, + coordination.clone(), + database.clone(), + ), + )); } let mut observers: Vec> = Vec::new(); let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config)); diff --git a/litellm-rust/crates/gateway/src/rate_limits.rs b/litellm-rust/crates/gateway/src/rate_limits.rs new file mode 100644 index 00000000000..2b9ca57cbbc --- /dev/null +++ b/litellm-rust/crates/gateway/src/rate_limits.rs @@ -0,0 +1,1193 @@ +//! The rate limits hook: refuses a request when its key, user, team, team member or end user is +//! over its requests or tokens per window, or its key has too many requests in flight, with the +//! proxy's counters, order, messages and headers. +//! +//! The counters are the proxy's own Redis keys, `{:}:window|requests|tokens` and +//! `{api_key:}:max_parallel_requests`, changed by scripts with the proxy's semantics, so +//! gateway replicas and Python proxies sharing that Redis enforce one limit. Windows are fixed and +//! start on their first request. A request counts against requests whether or not it is let +//! through; tokens are reserved up front from an estimate and settled to the real count when the +//! response ends. Without Redis, or while it is unreachable, the counters live in this replica's +//! memory. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use futures_util::future::BoxFuture; +use http::{HeaderName, HeaderValue}; +use litellm_config::Config; +use litellm_db::PgPool; +use litellm_gateway_auth::keys::Limits; +use litellm_gateway_hooks::{ + Facts, Hook, KeyFacts, Record, Refusal, RequestEdit, RequestFacts, Route, Settle, Targets, +}; +use redis::aio::ConnectionManager; +use serde_json::{Map, Value}; + +use crate::{coordination::Coordination, end_users::EndUsers, spend::civil}; + +const DEFAULT_WINDOW: u64 = 60; +/// How long an in-flight slot counts before it is presumed leaked. +const SLOT_TTL: u64 = 3600; +/// The output tokens assumed at most when a request does not cap them. +const MAX_OUTPUT_FLOOR: i64 = 1024; +/// Windows and slots held in memory before expired ones are dropped. +const MAX_LOCAL: usize = 100_000; +const FAIL_CLOSED: &str = "Rate limit enforcement unavailable: request counters could not be verified against Redis, and fail_closed_rate_limit_enforcement is enabled, so the request was rejected to avoid exceeding the configured rate limit. Retry shortly."; + +/// Counts a request on every `(window, requests)` pair, starting a window that is missing or +/// over, and returns each window's start and count. The caller compares; the count always rises. +const REQUESTS: &str = r#" +local results = {} +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local window_start = redis.call('GET', window_key) + if not window_start or (now - tonumber(window_start)) >= window_size then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, 1) + redis.call('EXPIRE', window_key, window_size) + redis.call('EXPIRE', counter_key, window_size) + table.insert(results, tostring(now)) + table.insert(results, 1) + else + local counter = redis.call('INCR', counter_key) + if redis.call('TTL', counter_key) == -1 then + redis.call('EXPIRE', counter_key, window_size) + end + table.insert(results, window_start) + table.insert(results, counter) + end +end +return results +"#; + +/// Adds the increment to every `(window, tokens)` pair, or to none when one would go over its +/// limit: `{0, count, start, ...}` or `{1, index, current, limit}`. +const RESERVE: &str = r#" +local now = tonumber(redis.call('TIME')[1]) +local count = #KEYS / 2 +local state = {} +for i = 1, count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[base]) + local increment = tonumber(ARGV[base + 1]) + local window_size = tonumber(ARGV[base + 3]) + local window_start = redis.call('GET', window_key) + local expired = (not window_start) or ((now - tonumber(window_start)) >= window_size) + local current = 0 + if not expired then + current = tonumber(redis.call('GET', counter_key) or 0) + end + local blocked + if increment > 0 then + blocked = current + increment > limit + else + blocked = current >= limit + end + if blocked then + return { 1, i, current, limit } + end + state[i] = { expired, window_start } +end +local results = { 0 } +local reset = {} +for i = 1, count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[base + 1]) + local ttl = tonumber(ARGV[base + 2]) + local window_size = tonumber(ARGV[base + 3]) + if state[i][1] then + if not reset[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset[window_key] = true + end + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + table.insert(results, now) + else + local counter = redis.call('INCRBY', counter_key, increment) + if redis.call('TTL', counter_key) == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, counter) + table.insert(results, tonumber(state[i][2])) + end +end +return results +"#; + +/// Takes an in-flight slot unless the gauge is full: `{0, in_flight}` or `{1, 1, in_flight, limit}`. +const ACQUIRE: &str = r#" +local now = tonumber(redis.call('TIME')[1]) +local limit = tonumber(ARGV[1]) +local slot_ttl = tonumber(ARGV[2]) +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now - slot_ttl) +local in_flight = redis.call('ZCARD', KEYS[1]) +if in_flight + 1 > limit then + return { 1, 1, in_flight, limit } +end +redis.call('ZADD', KEYS[1], now, ARGV[3]) +redis.call('EXPIRE', KEYS[1], slot_ttl) +return { 0, redis.call('ZCARD', KEYS[1]) } +"#; + +/// Adds each token delta, giving a counter without an expiry the window's. +const CHARGE: &str = r#" +local results = {} +for i = 1, #KEYS do + local value = redis.call('INCRBYFLOAT', KEYS[i], tonumber(ARGV[i * 2 - 1])) + local ttl = tonumber(ARGV[i * 2]) + if ttl and ttl > 0 and redis.call('TTL', KEYS[i]) == -1 then + redis.call('EXPIRE', KEYS[i], ttl) + end + table.insert(results, value) +end +return results +"#; + +#[derive(Clone, Debug, PartialEq)] +pub struct RateLimitSettings { + /// `LITELLM_RATE_LIMIT_WINDOW_SIZE`, in seconds; 60 by default. + pub window: u64, + /// `general_settings.fail_closed_rate_limit_enforcement`: refuse with a 503 rather than fall + /// back to this replica's counters while Redis is unreachable. + pub fail_closed: bool, + /// `litellm_settings.max_end_user_budget_id`: the limits of end users without their own. + pub end_user_budget_id: Option, +} + +impl Default for RateLimitSettings { + fn default() -> Self { + Self { + window: DEFAULT_WINDOW, + fail_closed: false, + end_user_budget_id: None, + } + } +} + +impl RateLimitSettings { + pub fn from_config(config: &Config) -> Self { + let general = &config.general_settings.additional_fields; + Self { + window: std::env::var("LITELLM_RATE_LIMIT_WINDOW_SIZE") + .ok() + .and_then(|window| window.trim().parse().ok()) + .filter(|window| *window > 0) + .unwrap_or(DEFAULT_WINDOW), + fail_closed: match general.get("fail_closed_rate_limit_enforcement") { + Some(litellm_config::Value::Bool(flag)) => *flag, + Some(litellm_config::Value::String(text)) => text.eq_ignore_ascii_case("true"), + _ => false, + }, + end_user_budget_id: config + .litellm_settings + .additional_fields + .get("max_end_user_budget_id") + .and_then(litellm_config::Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned), + } + } +} + +/// One scope a request counts against: `{key}:{value}` with its limits. +#[derive(Clone, Debug, PartialEq)] +struct Descriptor { + key: &'static str, + value: String, + limits: Limits, + max_parallel_requests: Option, +} + +impl Descriptor { + fn prefix(&self) -> String { + format!("{{{}:{}}}", self.key, self.value) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Kind { + Requests, + Tokens, + Parallel, +} + +impl Kind { + fn name(self) -> &'static str { + match self { + Self::Requests => "requests", + Self::Tokens => "tokens", + Self::Parallel => "max_parallel_requests", + } + } +} + +/// Where a descriptor stands after this request. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct Status { + descriptor: usize, + kind: Kind, + limit: i64, + remaining: i64, +} + +/// The token reservation went over on the descriptor at `index`, which held `current`. +#[derive(Debug, PartialEq)] +enum Reserved { + Counts(Vec), + Over { index: usize, current: i64 }, +} + +/// Redis could not be used and the settings say to refuse rather than fall back. +struct Unavailable; + +#[derive(Debug)] +struct Window { + start: u64, + requests: i64, + tokens: i64, +} + +/// This replica's counters, used without Redis or while it is unreachable. +#[derive(Default)] +struct Local { + windows: HashMap, + slots: HashMap>, +} + +impl Local { + fn window(&mut self, prefix: &str, now: u64, size: u64) -> (&mut Window, bool) { + if self.windows.len() >= MAX_LOCAL { + self.windows + .retain(|_, window| now.saturating_sub(window.start) < size); + } + let window = self.windows.entry(prefix.to_owned()).or_insert(Window { + start: now, + requests: 0, + tokens: 0, + }); + let expired = now.saturating_sub(window.start) >= size; + (window, expired) + } + + fn requests(&mut self, prefixes: &[String], now: u64, size: u64) -> Vec { + prefixes + .iter() + .map(|prefix| { + let (window, expired) = self.window(prefix, now, size); + if expired { + *window = Window { + start: now, + requests: 0, + tokens: 0, + }; + } + window.requests += 1; + window.requests + }) + .collect() + } + + fn reserve( + &mut self, + scopes: &[(String, i64)], + increment: i64, + now: u64, + size: u64, + ) -> Reserved { + for (index, (prefix, limit)) in scopes.iter().enumerate() { + let current = match self.windows.get(prefix) { + Some(window) if now.saturating_sub(window.start) < size => window.tokens, + _ => 0, + }; + if blocked(current, increment, *limit) { + return Reserved::Over { index, current }; + } + } + Reserved::Counts( + scopes + .iter() + .map(|(prefix, _)| { + let (window, expired) = self.window(prefix, now, size); + if expired { + *window = Window { + start: now, + requests: 0, + tokens: 0, + }; + } + window.tokens += increment; + window.tokens + }) + .collect(), + ) + } + + fn charge(&mut self, deltas: &[(String, i64)]) { + for (prefix, delta) in deltas { + if let Some(window) = self.windows.get_mut(prefix) { + window.tokens += delta; + } + } + } + + fn acquire(&mut self, prefix: &str, limit: i64, slot: &str) -> Result { + let ttl = Duration::from_secs(SLOT_TTL); + if self.slots.len() >= MAX_LOCAL { + self.slots.retain(|_, slots| !slots.is_empty()); + } + let slots = self.slots.entry(prefix.to_owned()).or_default(); + slots.retain(|_, taken| taken.elapsed() < ttl); + let in_flight = slots.len() as i64; + if in_flight + 1 > limit { + return Err(in_flight); + } + slots.insert(slot.to_owned(), Instant::now()); + Ok(in_flight + 1) + } + + fn release(&mut self, prefix: &str, slot: &str) { + if let Some(slots) = self.slots.get_mut(prefix) { + slots.remove(slot); + } + } +} + +fn blocked(current: i64, increment: i64, limit: i64) -> bool { + match increment > 0 { + true => current + increment > limit, + false => current >= limit, + } +} + +/// Where the counters live. +struct Store { + redis: Option>, + local: Mutex, + window: u64, + fail_closed: bool, +} + +impl Store { + fn connection(&self) -> Option<(ConnectionManager, &Coordination)> { + let redis = self.redis.as_deref()?; + Some((redis.connection()?, redis)) + } + + fn local(&self) -> std::sync::MutexGuard<'_, Local> { + self.local.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Logs why Redis was not used; refuses when the settings say to. + fn fall_back(&self, error: Option) -> Result<(), Unavailable> { + match error { + Some(error) => { + tracing::warn!(%error, "rate limit counters failed; using this replica's") + } + None => { + tracing::debug!("Redis not connected; using this replica's rate limit counters") + } + } + match self.fail_closed { + true => Err(Unavailable), + false => Ok(()), + } + } + + /// Counts the request on each prefix's window and returns the counts. + async fn requests(&self, prefixes: &[String]) -> Result, Unavailable> { + let now = epoch_seconds(); + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let mut command = redis::cmd("EVAL"); + command.arg(REQUESTS).arg(prefixes.len() * 2); + for prefix in prefixes { + command + .arg(redis.namespaced(&format!("{prefix}:window"))) + .arg(redis.namespaced(&format!("{prefix}:requests"))); + } + command.arg(now).arg(self.window); + match integers(command.query_async(&mut connection).await) { + Ok(reply) if reply.len() == prefixes.len() * 2 => { + return Ok(reply.iter().skip(1).step_by(2).copied().collect()); + } + Ok(_) => Some(malformed()), + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok(self.local().requests(prefixes, now, self.window)) + } + + /// Takes an in-flight slot: the count with it, or the full count without. `true` when the + /// slot is in Redis. + async fn acquire( + &self, + prefix: &str, + limit: i64, + slot: &str, + ) -> Result<(Result, bool), Unavailable> { + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let reply = redis::cmd("EVAL") + .arg(ACQUIRE) + .arg(1) + .arg(redis.namespaced(&format!("{prefix}:max_parallel_requests"))) + .arg(limit) + .arg(SLOT_TTL) + .arg(slot) + .query_async(&mut connection) + .await; + match integers(reply) { + Ok(reply) => match reply.as_slice() { + [0, in_flight] => return Ok((Ok(*in_flight), true)), + [1, _, in_flight, _] => return Ok((Err(*in_flight), true)), + _ => Some(malformed()), + }, + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok((self.local().acquire(prefix, limit, slot), false)) + } + + async fn release(&self, prefix: &str, slot: &str, shared: bool) { + if shared && let Some((mut connection, redis)) = self.connection() { + let key = redis.namespaced(&format!("{prefix}:max_parallel_requests")); + if let Err(error) = redis::cmd("ZREM") + .arg(key) + .arg(slot) + .query_async::(&mut connection) + .await + { + tracing::warn!(%error, "releasing an in-flight slot failed; it expires on its own"); + } + return; + } + self.local().release(prefix, slot); + } + + /// Reserves `increment` tokens on every `(prefix, limit)`, or on none. `true` when the + /// reservation is in Redis. + async fn reserve( + &self, + scopes: &[(String, i64)], + increment: i64, + ) -> Result<(Reserved, bool), Unavailable> { + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let mut command = redis::cmd("EVAL"); + command.arg(RESERVE).arg(scopes.len() * 2); + for (prefix, _) in scopes { + command + .arg(redis.namespaced(&format!("{prefix}:window"))) + .arg(redis.namespaced(&format!("{prefix}:tokens"))); + } + for (_, limit) in scopes { + command + .arg(*limit) + .arg(increment) + .arg(self.window) + .arg(self.window); + } + match integers(command.query_async(&mut connection).await) { + Ok(reply) if reply.first() == Some(&0) => { + return Ok(( + Reserved::Counts( + reply.iter().skip(1).step_by(2).copied().collect(), + ), + true, + )); + } + Ok(reply) => match reply.as_slice() { + [1, index, current, _] if (1..=scopes.len() as i64).contains(index) => { + return Ok(( + Reserved::Over { + index: (*index - 1) as usize, + current: *current, + }, + true, + )); + } + _ => Some(malformed()), + }, + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok(( + self.local() + .reserve(scopes, increment, epoch_seconds(), self.window), + false, + )) + } + + /// Adds each delta to its prefix's tokens. + async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) { + if shared && let Some((mut connection, redis)) = self.connection() { + let mut command = redis::cmd("EVAL"); + command.arg(CHARGE).arg(deltas.len()); + for (prefix, _) in &deltas { + command.arg(redis.namespaced(&format!("{prefix}:tokens"))); + } + for (_, delta) in &deltas { + command.arg(*delta).arg(self.window); + } + if let Err(error) = command.query_async::(&mut connection).await { + tracing::warn!(%error, "settling reserved tokens failed"); + } + return; + } + self.local().charge(&deltas); + } +} + +fn malformed() -> redis::RedisError { + redis::RedisError::from(( + redis::ErrorKind::UnexpectedReturnType, + "unexpected rate limit script reply", + )) +} + +/// A script's array reply as integers; counters may come back as strings. +fn integers(reply: redis::RedisResult>) -> redis::RedisResult> { + reply? + .iter() + .map(|value| { + let text = match value { + redis::Value::Int(number) => return Ok(*number), + redis::Value::BulkString(bytes) => std::str::from_utf8(bytes).ok(), + redis::Value::SimpleString(text) => Some(text.as_str()), + _ => None, + }; + text.and_then(|text| text.parse::().ok()) + .map(|number| number as i64) + .ok_or_else(malformed) + }) + .collect() +} + +fn epoch_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +/// `%Y-%m-%d %H:%M:%S UTC`, as the proxy prints when a limit resets. +fn utc(seconds: u64) -> String { + let (year, month, day) = civil((seconds / 86_400) as i64); + let time = seconds % 86_400; + format!( + "{year:04}-{month:02}-{day:02} {:02}:{:02}:{:02} UTC", + time / 3600, + time % 3600 / 60, + time % 60 + ) +} + +/// The proxy's token estimate for a request: a token per four characters of message text, +/// plus the output it may produce for every candidate. +#[derive(Debug, PartialEq)] +struct Estimate { + tokens: i64, + /// The output cap to send upstream when the request has none and the limits are small. + implicit_max_tokens: Option, +} + +fn estimate( + body: &Map, + key: Option<&KeyFacts>, + model: &str, + min_tpm: i64, +) -> Estimate { + let chars: usize = body + .get("messages") + .and_then(Value::as_array) + .into_iter() + .flatten() + .map(|message| match message.get("content") { + Some(Value::String(text)) => text.chars().count(), + Some(Value::Array(blocks)) => blocks + .iter() + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .map(|text| text.chars().count()) + .sum(), + _ => 0, + }) + .sum(); + let chars = i64::try_from(chars).unwrap_or(i64::MAX); + let input = match chars { + 0 => 0, + chars => (chars / 4).max(1), + }; + let integer = |value: &Value| match value { + Value::Number(number) => number.as_f64().map(|number| number as i64), + Value::String(text) => text.trim().parse::().ok().map(|number| number as i64), + _ => None, + }; + let cap = ["max_tokens", "max_completion_tokens"] + .iter() + .filter_map(|field| body.get(*field).and_then(integer)) + .max(); + let configured = key.and_then(|key| { + [&key.metadata, &key.team_metadata] + .into_iter() + .find_map(|metadata| { + metadata + .get("default_estimated_output_tokens_per_model") + .and_then(|models| models.get(model)) + .and_then(integer) + .or_else(|| { + metadata + .get("default_estimated_output_tokens") + .and_then(integer) + }) + }) + }); + let floor = (min_tpm / 4).clamp(1, MAX_OUTPUT_FLOOR); + let output = match (cap, configured) { + (Some(cap), _) => cap, + (None, Some(configured)) => configured, + (None, None) if chars == 0 => 0, + (None, None) => input.max(floor), + }; + let candidates = ["n", "best_of"] + .iter() + .filter_map(|field| body.get(*field).and_then(integer)) + .fold(1, i64::max); + Estimate { + tokens: input + .saturating_add(output.saturating_mul(candidates)) + .max(1), + implicit_max_tokens: (floor < MAX_OUTPUT_FLOOR && cap.is_none()) + .then(|| floor.max(configured.unwrap_or(0))), + } +} + +/// The `rate_limits` hook. +pub struct RateLimiter { + settings: RateLimitSettings, + store: Arc, + end_users: Option, +} + +impl RateLimiter { + /// `redis`, when given, holds the shared counters; `database` gives end users' limits. + pub fn new( + settings: RateLimitSettings, + redis: Option>, + database: Option, + ) -> Self { + Self { + store: Arc::new(Store { + redis, + local: Mutex::default(), + window: settings.window, + fail_closed: settings.fail_closed, + }), + settings, + end_users: database.map(EndUsers::new), + } + } + + pub fn for_gateway( + config: &Config, + redis: Option>, + database: Option, + ) -> Self { + Self::new(RateLimitSettings::from_config(config), redis, database) + } + + /// The scopes this request counts against, in the proxy's order. + async fn descriptors(&self, request: &RequestFacts) -> Vec { + let Some(key) = request.key.as_deref() else { + return Vec::new(); + }; + let limits = &key.limits; + let mut descriptors = Vec::new(); + let mut push = |key: &'static str, value: String, limits: Limits, parallel: Option| { + if limits.is_set() || parallel.is_some() { + descriptors.push(Descriptor { + key, + value, + limits, + max_parallel_requests: parallel, + }); + } + }; + if !key.hash.is_empty() { + push( + "api_key", + key.hash.clone(), + limits.key, + limits.max_parallel_requests, + ); + } + if let Some(user_id) = &key.user_id { + push("user", user_id.clone(), limits.user, None); + } + if let Some(team_id) = &key.team_id { + push("team", team_id.clone(), limits.team, None); + } + if let (Some(team_id), Some(user_id)) = (&key.team_id, &key.user_id) { + push( + "team_member", + format!("{team_id}:{user_id}"), + limits.team_member, + None, + ); + } + if let (Some(end_user_id), Some(end_users)) = (request.end_user.as_deref(), &self.end_users) + && let Some(end_user) = end_users + .get( + end_user_id, + self.settings.end_user_budget_id.as_deref(), + key.budgets.end_user_budget_id.as_deref(), + ) + .await + { + push( + "end_user", + end_user_id.to_owned(), + Limits { + rpm: end_user.limits.rpm, + tpm: end_user.limits.tpm, + }, + None, + ); + } + descriptors + } + + fn refusal(&self, request: &RequestFacts, descriptor: &Descriptor, status: Status) -> Refusal { + let reset_at = utc(epoch_seconds() + self.settings.window); + let message = format!( + "Rate limit exceeded for {}: {}. Limit type: {}. Current limit: {}, Remaining: {}. Limit resets at: {reset_at}", + descriptor.key, + descriptor.value, + status.kind.name(), + status.limit, + status.remaining.max(0), + ); + let mut refusal = match request.route { + Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(), + Route::ChatCompletions => { + let mut refusal = Refusal::new(429, "throttling_error", message).without_param(); + for (name, value) in [ + ("retry-after", self.settings.window.to_string()), + ("rate_limit_type", status.kind.name().to_owned()), + ("reset_at", reset_at), + ] { + if let Ok(value) = HeaderValue::from_str(&value) { + refusal = refusal.with_header(HeaderName::from_static(name), value); + } + } + refusal + } + }; + for (name, value) in key_limit_headers(request.key.as_deref()) { + refusal = refusal.with_header(name, value); + } + refusal + } + + fn unavailable(&self, request: &RequestFacts) -> Refusal { + match request.route { + Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(), + Route::ChatCompletions => { + Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param() + } + } + } +} + +/// `x-litellm-key-tpm-limit` and `x-litellm-key-rpm-limit`, which the proxy sends whenever the +/// key sets them. +fn key_limit_headers(key: Option<&KeyFacts>) -> Vec<(HeaderName, HeaderValue)> { + let Some(key) = key else { + return Vec::new(); + }; + [ + ("x-litellm-key-tpm-limit", key.limits.key.tpm), + ("x-litellm-key-rpm-limit", key.limits.key.rpm), + ] + .into_iter() + .filter_map(|(name, limit)| Some((HeaderName::from_static(name), HeaderValue::from(limit?)))) + .collect() +} + +impl Hook for RateLimiter { + fn name(&self) -> &'static str { + "rate_limits" + } + + fn reads(&self) -> Facts { + Facts::KEY | Facts::END_USER | Facts::REQUEST_BODY | Facts::USAGE + } + + fn writes(&self) -> Targets { + Targets::ADMISSION | Targets::UPSTREAM_BODY | Targets::RESPONSE_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + let descriptors = self.descriptors(request).await; + if descriptors.is_empty() { + return Ok(None); + } + let unavailable = |Unavailable| self.unavailable(request); + let mut statuses = Vec::new(); + + // Requests per window, counted whether or not the request goes through. + let windowed: Vec = (0..descriptors.len()) + .filter(|index| descriptors[*index].limits.rpm.is_some()) + .collect(); + if !windowed.is_empty() { + let prefixes: Vec = windowed + .iter() + .map(|index| descriptors[*index].prefix()) + .collect(); + let counts = self.store.requests(&prefixes).await.map_err(unavailable)?; + for (index, count) in windowed.iter().zip(counts) { + let limit = descriptors[*index].limits.rpm.unwrap_or_default(); + statuses.push(Status { + descriptor: *index, + kind: Kind::Requests, + limit, + remaining: limit - count, + }); + } + if let Some(over) = statuses.iter().find(|status| status.remaining < 0) { + return Err(self.refusal(request, &descriptors[over.descriptor], *over)); + } + } + + // Requests in flight, for the key only. + let mut slot = None; + if let Some(index) = descriptors + .iter() + .position(|descriptor| descriptor.max_parallel_requests.is_some()) + { + let descriptor = &descriptors[index]; + let limit = descriptor.max_parallel_requests.unwrap_or_default(); + let prefix = descriptor.prefix(); + let id = uuid::Uuid::new_v4().simple().to_string(); + let (taken, shared) = self + .store + .acquire(&prefix, limit, &id) + .await + .map_err(unavailable)?; + let in_flight = match taken { + Ok(in_flight) => in_flight, + Err(in_flight) => { + let status = Status { + descriptor: index, + kind: Kind::Parallel, + limit, + remaining: limit - in_flight, + }; + return Err(self.refusal(request, descriptor, status)); + } + }; + statuses.push(Status { + descriptor: index, + kind: Kind::Parallel, + limit, + remaining: (limit - in_flight).max(0), + }); + slot = Some(Slot { prefix, id, shared }); + } + + // Tokens per window, reserved from an estimate and settled when the response ends. + let limited: Vec<(usize, i64)> = descriptors + .iter() + .enumerate() + .filter_map(|(index, descriptor)| Some((index, descriptor.limits.tpm?))) + .collect(); + let mut tokens = None; + if let Some(min_tpm) = limited.iter().map(|(_, limit)| *limit).min() { + let body: Map = + serde_json::from_slice(request.body()).unwrap_or_default(); + let estimate = estimate( + &body, + request.key.as_deref(), + &request.requested_model, + min_tpm, + ); + if let Some(max_tokens) = estimate.implicit_max_tokens { + edit.set_upstream_body_field("max_tokens", Value::from(max_tokens)); + } + let scopes: Vec<(String, i64)> = limited + .iter() + .map(|(index, limit)| (descriptors[*index].prefix(), *limit)) + .collect(); + let reserved = self.store.reserve(&scopes, estimate.tokens).await; + let release = |slot: Option| async { + if let Some(slot) = slot { + self.store + .release(&slot.prefix, &slot.id, slot.shared) + .await; + } + }; + let (reserved, shared) = match reserved { + Ok(reserved) => reserved, + Err(Unavailable) => { + release(slot).await; + return Err(self.unavailable(request)); + } + }; + match reserved { + Reserved::Counts(counts) => { + for ((index, limit), count) in limited.iter().zip(counts) { + statuses.push(Status { + descriptor: *index, + kind: Kind::Tokens, + limit: *limit, + remaining: (limit - count).max(0), + }); + } + } + Reserved::Over { index, current } => { + release(slot).await; + let (descriptor, limit) = limited[index]; + let status = Status { + descriptor, + kind: Kind::Tokens, + limit, + remaining: limit - current, + }; + return Err(self.refusal(request, &descriptors[descriptor], status)); + } + } + tokens = Some(Tokens { + prefixes: scopes.into_iter().map(|(prefix, _)| prefix).collect(), + reserved: estimate.tokens, + shared, + }); + } + + // The proxy sends where each limit stands on complete responses only. + if !request.stream { + for status in &statuses { + let key = descriptors[status.descriptor].key; + let kind = status.kind.name(); + for (name, value) in [ + ( + format!("x-ratelimit-{key}-remaining-{kind}"), + status.remaining, + ), + (format!("x-ratelimit-{key}-limit-{kind}"), status.limit), + ] { + if let Ok(name) = HeaderName::from_bytes(name.as_bytes()) { + edit.add_response_header(name, HeaderValue::from(value)); + } + } + } + } + for (name, value) in key_limit_headers(request.key.as_deref()) { + edit.add_response_header(name, value); + } + if slot.is_none() && tokens.is_none() { + return Ok(None); + } + Ok(Some(Box::new(Held { + store: self.store.clone(), + slot, + tokens, + }) as Box)) + }) + } +} + +struct Slot { + prefix: String, + id: String, + shared: bool, +} + +struct Tokens { + prefixes: Vec, + reserved: i64, + shared: bool, +} + +/// What a request holds until it ends: an in-flight slot and reserved tokens. +struct Held { + store: Arc, + slot: Option, + tokens: Option, +} + +impl Settle for Held { + fn settle(self: Box, record: &Record) { + let Held { + store, + slot, + tokens, + } = *self; + let charge = tokens.and_then(|tokens| { + // The proxy counts every token but cache reads. Without usage, a success keeps its + // reservation and a failure returns it. + let actual = match (&record.outcome.usage, record.succeeded()) { + (Some(usage), _) => i64::try_from(usage.total_tokens() - usage.cache_read_tokens) + .unwrap_or(i64::MAX), + (None, true) => tokens.reserved, + (None, false) => 0, + }; + let delta = actual - tokens.reserved; + (delta != 0).then(|| { + ( + tokens + .prefixes + .into_iter() + .map(|prefix| (prefix, delta)) + .collect::>(), + tokens.shared, + ) + }) + }); + tokio::spawn(async move { + if let Some(slot) = slot { + store.release(&slot.prefix, &slot.id, slot.shared).await; + } + if let Some((deltas, shared)) = charge { + store.charge(deltas, shared).await; + } + }); + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn body(value: Value) -> Map { + match value { + Value::Object(object) => object, + _ => unreachable!(), + } + } + + #[rstest] + #[case::utc_epoch(0, "1970-01-01 00:00:00 UTC")] + #[case::utc_later(1_792_238_465, "2026-10-17 12:01:05 UTC")] + fn prints_reset_times_like_the_proxy(#[case] seconds: u64, #[case] expected: &str) { + assert_eq!(utc(seconds), expected); + } + + #[rstest] + // 40 characters: 10 input tokens; the small limit floors output at 100_000 / 4 = 1024. + #[case::uncapped(json!({"messages": [{"role": "user", "content": "a".repeat(40)}]}), 100_000, 10 + 1024, None)] + // A small limit floors output at 400 / 4 = 100 and caps the request there. + #[case::small_limit(json!({"messages": [{"role": "user", "content": "a".repeat(40)}]}), 400, 110, Some(100))] + #[case::capped(json!({"max_tokens": 50, "messages": [{"role": "user", "content": [{"type": "text", "text": "a".repeat(8)}]}]}), 400, 52, None)] + #[case::candidates(json!({"max_completion_tokens": 5, "n": 3, "messages": [{"role": "user", "content": "abcd"}]}), 100_000, 16, None)] + #[case::no_text(json!({"messages": []}), 100_000, 1, None)] + fn estimates_tokens_like_the_proxy( + #[case] request: Value, + #[case] min_tpm: i64, + #[case] tokens: i64, + #[case] implicit_max_tokens: Option, + ) { + assert_eq!( + estimate(&body(request), None, "m", min_tpm), + Estimate { + tokens, + implicit_max_tokens + } + ); + } + + #[rstest] + fn configured_output_estimates_come_from_the_key_then_the_team() { + let metadata = |value: Value| Arc::new(body(value)); + let key = KeyFacts { + metadata: metadata(json!({"default_estimated_output_tokens_per_model": {"m": 7}})), + team_metadata: metadata(json!({"default_estimated_output_tokens": 9})), + ..KeyFacts::default() + }; + let request = body(json!({"messages": [{"role": "user", "content": "abcd"}]})); + assert_eq!(estimate(&request, Some(&key), "m", 100_000).tokens, 8); + assert_eq!(estimate(&request, Some(&key), "other", 100_000).tokens, 10); + } + + #[rstest] + fn local_windows_count_reserve_and_reset() { + let mut local = Local::default(); + let prefixes = vec!["{api_key:h}".to_owned()]; + assert_eq!(local.requests(&prefixes, 100, 60), [1]); + assert_eq!(local.requests(&prefixes, 159, 60), [2]); + let scopes = vec![("{api_key:h}".to_owned(), 10)]; + assert_eq!( + local.reserve(&scopes, 6, 159, 60), + Reserved::Counts(vec![6]) + ); + assert_eq!( + local.reserve(&scopes, 6, 159, 60), + Reserved::Over { + index: 0, + current: 6 + } + ); + local.charge(&[("{api_key:h}".to_owned(), -4)]); + assert_eq!( + local.reserve(&scopes, 8, 159, 60), + Reserved::Counts(vec![10]) + ); + // The window started at 100 is over at 160: requests and tokens start again. + assert_eq!(local.requests(&prefixes, 160, 60), [1]); + assert_eq!( + local.reserve(&scopes, 10, 160, 60), + Reserved::Counts(vec![10]) + ); + } + + #[rstest] + fn local_slots_fill_and_release() { + let mut local = Local::default(); + assert_eq!(local.acquire("{api_key:h}", 2, "a"), Ok(1)); + assert_eq!(local.acquire("{api_key:h}", 2, "b"), Ok(2)); + assert_eq!(local.acquire("{api_key:h}", 2, "c"), Err(2)); + local.release("{api_key:h}", "a"); + assert_eq!(local.acquire("{api_key:h}", 2, "c"), Ok(2)); + } +} diff --git a/litellm-rust/crates/gateway/src/spend.rs b/litellm-rust/crates/gateway/src/spend.rs index fc373307bec..a2f0d077819 100644 --- a/litellm-rust/crates/gateway/src/spend.rs +++ b/litellm-rust/crates/gateway/src/spend.rs @@ -648,7 +648,7 @@ fn usage_values(usage: &Usage) -> (Value, Value) { } /// Days since 1970-01-01 as a UTC calendar date. -fn civil(days: i64) -> (i64, u32, u32) { +pub(crate) fn civil(days: i64) -> (i64, u32, u32) { let days = days + 719_468; let era = days.div_euclid(146_097); let day_of_era = days.rem_euclid(146_097); diff --git a/litellm-rust/crates/gateway/tests/rate_limits.rs b/litellm-rust/crates/gateway/tests/rate_limits.rs new file mode 100644 index 00000000000..2e7101bd89c --- /dev/null +++ b/litellm-rust/crates/gateway/tests/rate_limits.rs @@ -0,0 +1,515 @@ +//! The rate limits hook on real requests: refusals and headers as the proxy sends them, tokens +//! reserved and settled to usage, in-flight slots, and counters shared through Redis. Set +//! `LITELLM_TEST_REDIS_URL` to also run the shared cases against a live Redis. + +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{HeaderMap, Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway::{ + coordination::{Coordination, RedisSettings}, + rate_limits::{RateLimitSettings, RateLimiter}, +}; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLimits, KeyLookup, KeyModels, KeyProfile, KeyStatus, Limits, + LookupFuture, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +/// 3 input and 4 output tokens: 7 count against a token limit. +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; + +#[derive(Clone, Default)] +struct Upstream { + calls: Arc, + bodies: Arc>>, +} + +async fn upstream(delay: Duration) -> (String, Upstream) { + let upstream = Upstream::default(); + let seen = upstream.clone(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(move |request: Request| { + let seen = seen.clone(); + async move { + seen.calls.fetch_add(1, Ordering::SeqCst); + let body = to_bytes(request.into_body(), 1 << 20).await.unwrap(); + seen.bodies + .lock() + .unwrap() + .push(serde_json::from_slice(&body).unwrap_or(Value::Null)); + tokio::time::sleep(delay).await; + ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response() + } + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + (format!("http://{address}"), upstream) +} + +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 hash(token: &str) -> KeyHash { + KeyHash::from_token(&SecretValue::new(token)) +} + +fn key(team_id: Option<&str>, limits: KeyLimits) -> KeyStatus { + KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: team_id.map(str::to_owned), + profile: Arc::new(KeyProfile { + user_id: Some("user-1".into()), + limits, + ..KeyProfile::default() + }), + } +} + +fn rpm(limit: i64) -> Limits { + Limits { + rpm: Some(limit), + tpm: None, + } +} + +fn tpm(limit: i64) -> Limits { + Limits { + rpm: None, + tpm: Some(limit), + } +} + +struct Gateway { + app: Router, + upstream: Upstream, +} + +struct Setup { + redis: Option>, + database: Option, + delay: Duration, +} + +impl Default for Setup { + fn default() -> Self { + Self { + redis: None, + database: None, + delay: Duration::ZERO, + } + } +} + +async fn gateway(keys: Vec<(&str, KeyStatus)>, setup: Setup) -> Gateway { + let (address, upstream) = upstream(setup.delay).await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {address} + api_key: sk-upstream +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let limiter = RateLimiter::new(RateLimitSettings::default(), setup.redis, setup.database); + litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(limiter)], Vec::new()) + .await + .unwrap(); + let keys: Arc = Arc::new(Keys( + keys.into_iter() + .map(|(token, status)| (hash(token), status)) + .collect(), + )); + let app = litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys)); + Gateway { app, upstream } +} + +struct Reply { + status: u16, + headers: HeaderMap, + body: Value, +} + +async fn send(app: &Router, path: &str, body: Value) -> Reply { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer sk-key") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + Reply { + status, + headers, + body: serde_json::from_slice(&body).unwrap_or(Value::Null), + } +} + +fn message(extra: Value) -> Value { + let mut body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + for (field, value) in extra.as_object().unwrap() { + body[field] = value.clone(); + } + body +} + +async fn messages(app: &Router) -> Reply { + send(app, "/v1/messages", message(json!({}))).await +} + +async fn chat(app: &Router) -> Reply { + send(app, "/v1/chat/completions", message(json!({}))).await +} + +/// Settling runs after the response, off the request path. +async fn settled() { + tokio::time::sleep(Duration::from_millis(100)).await; +} + +fn header<'a>(reply: &'a Reply, name: &str) -> Option<&'a str> { + reply.headers.get(name).map(|value| value.to_str().unwrap()) +} + +/// The message up to the reset time, which moves with the clock. +fn refusal_message(reply: &Reply, path: &[&str]) -> String { + let mut value = &reply.body; + for part in path { + value = &value[*part]; + } + let message = value.as_str().unwrap(); + let (message, reset) = message.split_once(" Limit resets at: ").unwrap(); + assert!(reset.ends_with(" UTC"), "{reset}"); + message.to_owned() +} + +#[rstest] +#[tokio::test] +async fn requests_over_the_key_limit_are_refused_like_the_proxy() { + let limits = KeyLimits { + key: rpm(2), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let first = messages(&gateway.app).await; + assert_eq!(first.status, 200); + assert_eq!( + header(&first, "x-ratelimit-api_key-remaining-requests"), + Some("1") + ); + assert_eq!( + header(&first, "x-ratelimit-api_key-limit-requests"), + Some("2") + ); + assert_eq!(header(&first, "x-litellm-key-rpm-limit"), Some("2")); + assert_eq!(chat(&gateway.app).await.status, 200); + + let expected = format!( + "Rate limit exceeded for api_key: {}. Limit type: requests. Current limit: 2, Remaining: 0.", + hash("sk-key").as_str() + ); + // Chat Completions: the proxy's error body and its retry headers. + let refused = chat(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!(refusal_message(&refused, &["error", "message"]), expected); + assert_eq!(refused.body["error"]["type"], "throttling_error"); + assert_eq!(refused.body["error"]["param"], Value::Null); + assert_eq!(refused.body["error"]["code"], "429"); + assert_eq!(header(&refused, "retry-after"), Some("60")); + assert_eq!(header(&refused, "rate_limit_type"), Some("requests")); + assert!(header(&refused, "reset_at").unwrap().ends_with(" UTC")); + assert_eq!(header(&refused, "x-litellm-key-rpm-limit"), Some("2")); + // Messages: Anthropic's envelope, without the retry headers. + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!(refused.body["type"], "error"); + assert_eq!(refused.body["error"]["type"], "rate_limit_error"); + assert_eq!(refusal_message(&refused, &["error", "message"]), expected); + assert_eq!(header(&refused, "retry-after"), None); + assert_eq!(gateway.upstream.calls.load(Ordering::SeqCst), 2); +} + +#[rstest] +#[case::user(KeyLimits { user: rpm(1), ..KeyLimits::default() }, "user: user-1")] +#[case::team(KeyLimits { team: rpm(1), ..KeyLimits::default() }, "team: team-1")] +#[case::team_member(KeyLimits { team_member: rpm(1), ..KeyLimits::default() }, "team_member: team-1:user-1")] +#[tokio::test] +async fn owners_are_limited_under_their_own_names( + #[case] limits: KeyLimits, + #[case] descriptor: &str, +) { + let gateway = gateway( + vec![("sk-key", key(Some("team-1"), limits))], + Setup::default(), + ) + .await; + assert_eq!(messages(&gateway.app).await.status, 200); + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + format!( + "Rate limit exceeded for {descriptor}. Limit type: requests. Current limit: 1, Remaining: 0." + ) + ); +} + +#[rstest] +#[tokio::test] +async fn tokens_are_reserved_and_settled_to_usage() { + // Each request reserves 1 input + 8 output tokens and settles to 7. + let limits = KeyLimits { + key: tpm(20), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let first = messages(&gateway.app).await; + assert_eq!(first.status, 200); + assert_eq!( + header(&first, "x-ratelimit-api_key-remaining-tokens"), + Some("11") + ); + assert_eq!(header(&first, "x-litellm-key-tpm-limit"), Some("20")); + settled().await; + assert_eq!(messages(&gateway.app).await.status, 200); + settled().await; + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + format!( + "Rate limit exceeded for api_key: {}. Limit type: tokens. Current limit: 20, Remaining: 6.", + hash("sk-key").as_str() + ) + ); +} + +#[rstest] +#[tokio::test] +async fn small_token_limits_cap_uncapped_requests() { + let limits = KeyLimits { + key: tpm(400), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let body = json!({"model": "claude", "messages": [{"role": "user", "content": "hi"}]}); + assert_eq!(send(&gateway.app, "/v1/messages", body).await.status, 200); + assert_eq!( + gateway.upstream.bodies.lock().unwrap()[0]["max_tokens"], + 100 + ); +} + +#[rstest] +#[tokio::test] +async fn in_flight_requests_are_capped_per_key() { + let limits = KeyLimits { + max_parallel_requests: Some(1), + ..KeyLimits::default() + }; + let gateway = gateway( + vec![("sk-key", key(None, limits))], + Setup { + delay: Duration::from_millis(300), + ..Setup::default() + }, + ) + .await; + let (first, second) = tokio::join!(messages(&gateway.app), async { + tokio::time::sleep(Duration::from_millis(100)).await; + messages(&gateway.app).await + }); + assert_eq!(first.status, 200); + assert_eq!( + header( + &first, + "x-ratelimit-api_key-remaining-max_parallel_requests" + ), + Some("0") + ); + assert_eq!(second.status, 429); + assert_eq!( + refusal_message(&second, &["error", "message"]), + format!( + "Rate limit exceeded for api_key: {}. Limit type: max_parallel_requests. Current limit: 1, Remaining: 0.", + hash("sk-key").as_str() + ) + ); + settled().await; + assert_eq!(messages(&gateway.app).await.status, 200); +} + +fn redis_url() -> Option { + std::env::var("LITELLM_TEST_REDIS_URL").ok() +} + +async fn redis(url: &str, namespace: &str) -> Arc { + let coordination = + Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap(); + let router: Arc = Arc::new(NoRouter( + litellm_gateway_inference::ModelRouter::from_config( + &Config::from_yaml("model_list: []").unwrap(), + ) + .unwrap(), + )); + coordination.start(router, None, tokio_util::sync::CancellationToken::new()); + for _ in 0..100 { + if coordination.connected() { + return coordination; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("Redis did not connect"); +} + +struct NoRouter(litellm_gateway_inference::ModelRouter); + +impl litellm_gateway::coordination::RouterHandle for NoRouter { + fn router(&self) -> &litellm_gateway_inference::ModelRouter { + &self.0 + } +} + +#[rstest] +#[tokio::test] +async fn replicas_share_the_proxy_counters() { + let Some(url) = redis_url() else { return }; + let namespace = format!("rate-limits-{}", uuid::Uuid::new_v4()); + let limits = KeyLimits { + key: Limits { + rpm: Some(3), + tpm: Some(1_000), + }, + max_parallel_requests: Some(5), + ..KeyLimits::default() + }; + let mut replicas = Vec::new(); + for _ in 0..2 { + let setup = Setup { + redis: Some(redis(&url, &namespace).await), + ..Setup::default() + }; + replicas.push(gateway(vec![("sk-key", key(None, limits))], setup).await); + } + assert_eq!(messages(&replicas[0].app).await.status, 200); + assert_eq!(messages(&replicas[1].app).await.status, 200); + let third = messages(&replicas[0].app).await; + assert_eq!( + header(&third, "x-ratelimit-api_key-remaining-requests"), + Some("0") + ); + settled().await; + assert_eq!(messages(&replicas[1].app).await.status, 429); + + let mut connection = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let prefix = format!("{namespace}:{{api_key:{}}}", hash("sk-key").as_str()); + let get = async |connection: &mut redis::aio::MultiplexedConnection, suffix: &str| { + redis::cmd("GET") + .arg(format!("{prefix}:{suffix}")) + .query_async::>(connection) + .await + .unwrap() + }; + // The refused request still counted; three settled requests used 7 tokens each. + assert_eq!(get(&mut connection, "requests").await.as_deref(), Some("4")); + assert_eq!(get(&mut connection, "tokens").await.as_deref(), Some("21")); + assert!(get(&mut connection, "window").await.is_some()); + let in_flight: i64 = redis::cmd("ZCARD") + .arg(format!("{prefix}:max_parallel_requests")) + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(in_flight, 0); +} + +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows; the lookups under test use the checked macro" +)] +async fn end_users(pool: &litellm_db::PgPool) { + sqlx::raw_sql( + r#" +INSERT INTO "LiteLLM_BudgetTable" (budget_id, rpm_limit, created_by, updated_by) + VALUES ('one-a-minute', 1, 'x', 'x'); +INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id) + VALUES ('limited', 0, 'one-a-minute'), ('free', 0, NULL); +"#, + ) + .execute(pool) + .await + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn end_users_are_held_to_their_budget_limits() { + let postgres = litellm_db_testing::MigratedPostgres::start().await.unwrap(); + end_users(postgres.pool()).await; + let gateway = gateway( + vec![("sk-key", key(None, KeyLimits::default()))], + Setup { + database: Some(postgres.pool().clone()), + ..Setup::default() + }, + ) + .await; + let as_end_user = |user: &str| message(json!({"metadata": {"user_id": user}})); + for user in ["limited", "free", "free"] { + assert_eq!( + send(&gateway.app, "/v1/messages", as_end_user(user)) + .await + .status, + 200 + ); + } + let refused = send(&gateway.app, "/v1/messages", as_end_user("limited")).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + "Rate limit exceeded for end_user: limited. Limit type: requests. Current limit: 1, Remaining: 0." + ); +} diff --git a/litellm-rust/docs/customer-proxy/rate-limits.svg b/litellm-rust/docs/customer-proxy/rate-limits.svg new file mode 100644 index 00000000000..beff70d8482 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/rate-limits.svg @@ -0,0 +1,55 @@ + + + + + + + Rate limits hook: the proxy's limiter on the proxy's shared counters + + + Scopes, in the proxy's order + api_key: key rpm/tpm (else budget row's), + max_parallel_requests + user, team, team_member (team:user) + end_user: its budget row, or the default + read with the key; end users cached 60 s + no limits set: no Redis call at all + + + Three Redis scripts, in the proxy's phases + 1. requests: +1 on every {scope:id}:requests in one call, fixed window from the first hit; + over when count > limit; the count is never given back + 2. in flight (key only): ZSET of slots on Redis TIME, pruned at 3600 s, full at limit + 3. tokens: reserve chars/4 + output x candidates on every scope, all or none; + output = max_tokens, else the configured estimate, else max(input, min(1024, tpm/4)); + a limit under 4096 caps an uncapped request at that floor upstream + + + + Refusal: 429, first scope over + "Rate limit exceeded for api_key: <hash>. Limit type: requests. + Current limit: 10, Remaining: 0. Limit resets at: ... UTC" + chat: throttling_error + retry-after, rate_limit_type, reset_at + messages: Anthropic rate_limit_error envelope, no retry headers + a refused token reservation gives its slot back + + + + Admitted: request runs + complete responses carry where each limit stands: + x-ratelimit-{scope}-remaining-{type}, -limit-{type} + every response: x-litellm-key-rpm-limit / -tpm-limit + streams get no x-ratelimit-* headers, as in the proxy + + + + Settle (off the request path) + tokens: INCRBYFLOAT (actual - reserved) on each reserved scope, actual = all tokens but cache reads; a failure returns the reservation + in flight: ZREM the request's slot when the response ends, fails, or the client goes away + + + + Same keys, scripts and windows as the Python proxy, under its namespace + ns:{api_key:<hash>}:window | requests | tokens | max_parallel_requests: gateway replicas and Python proxies enforce one limit. + No Redis, or Redis unreachable: the same logic on this replica's counters (fail_closed_rate_limit_enforcement: 503 instead) + From 40d4c8b2d8f0ddb66c37ff5d038a7045edfd4e6d Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:45 -0700 Subject: [PATCH 05/10] feat(gateway): configurable header, spend metadata and provider metadata hooks (#47054) Adds four built-in hooks a config lists under litellm_settings.hooks and runs in that order: forward_headers copies named client headers to the provider, strip_headers removes headers from the provider request (including the gateway's default pass-through), spend_metadata adds header and key-metadata values to the spend log's spend_logs_metadata, and provider_metadata sets Messages body fields and headers from templates over request facts. Hooks' upstream header edits now reach the provider on both endpoints, and body edits can set nested fields. Credentials and signer headers are never forwarded or set, and a misconfigured hook stops the gateway at startup. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/crates/gateway-hooks/src/hook.rs | 58 +- .../gateway-inference/src/chat_completions.rs | 19 +- .../crates/gateway-inference/src/hooks.rs | 14 + .../crates/gateway-inference/src/messages.rs | 73 ++- litellm-rust/crates/gateway/src/error.rs | 2 + litellm-rust/crates/gateway/src/lib.rs | 1 + litellm-rust/crates/gateway/src/main.rs | 1 + .../crates/gateway/src/request_hooks.rs | 556 ++++++++++++++++++ .../crates/gateway/tests/request_hooks.rs | 339 +++++++++++ .../docs/customer-proxy/request-hooks.svg | 56 ++ 10 files changed, 1101 insertions(+), 18 deletions(-) create mode 100644 litellm-rust/crates/gateway/src/request_hooks.rs create mode 100644 litellm-rust/crates/gateway/tests/request_hooks.rs create mode 100644 litellm-rust/docs/customer-proxy/request-hooks.svg 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-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 118eb2cacad..e5fac1cafcc 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -93,6 +93,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,7 +132,7 @@ 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 call = ChatCompletionsCall { @@ -130,7 +142,7 @@ 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, }; Ok(litellm_host_http::serve_unary( @@ -166,6 +178,7 @@ async fn openai( gateway: &Gateway, deployment: &Deployment, body: Map, + extra_headers: Option>, ) -> Result { let machine = gateway.chat_completions.clone().openai_machine( OpenAiChatCall { @@ -174,7 +187,7 @@ 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, }, None, diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index 07833d8cc93..4893a91e835 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -154,6 +154,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) { 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/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index a2097417ce5..fbab39ef426 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -20,6 +20,8 @@ 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("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 57f05ec9322..c97c1d6515b 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -10,6 +10,7 @@ mod key_status; mod keys; pub mod metrics; pub mod rate_limits; +pub mod request_hooks; 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 3e648f5a758..b8a35eddbcd 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -128,6 +128,7 @@ async fn main() -> Result<(), Box> { ), )); } + 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(); diff --git a/litellm-rust/crates/gateway/src/request_hooks.rs b/litellm-rust/crates/gateway/src/request_hooks.rs new file mode 100644 index 00000000000..348a6852504 --- /dev/null +++ b/litellm-rust/crates/gateway/src/request_hooks.rs @@ -0,0 +1,556 @@ +//! Built-in hooks a config turns on and configures under `litellm_settings.hooks`, run in the +//! order listed: +//! +//! ```yaml +//! litellm_settings: +//! hooks: +//! - forward_headers: +//! headers: ["x-request-origin", "x-trace-*"] +//! - spend_metadata: +//! from_headers: +//! x-litellm-spend-logs-metadata: json +//! x-cost-center: cost_center +//! from_key_metadata: [cost_center] +//! - strip_headers: +//! headers: ["x-litellm-spend-logs-metadata", "x-cost-center"] +//! - provider_metadata: +//! body: +//! metadata.user_id: "{end_user|key.user_id}" +//! headers: +//! x-caller-team: "{key.team_alias}" +//! ``` + +use std::sync::Arc; + +use futures_util::future::BoxFuture; +use http::{HeaderName, HeaderValue}; +use litellm_config::Config; +use litellm_gateway_hooks::{ + Facts, HeaderPattern, Hook, Refusal, RequestEdit, RequestFacts, Route, Settle, Targets, +}; +use serde::Deserialize; +use serde_json::{Map, Value}; + +use crate::{Error, spend::SPEND_LOGS_METADATA}; + +/// Headers a hook never forwards: credentials, and those the HTTP client or a request signer +/// sets itself. +const PROTECTED: [&str; 13] = [ + "authorization", + "proxy-authorization", + "x-api-key", + "api-key", + "x-litellm-api-key", + "x-goog-api-key", + "cookie", + "host", + "content-length", + "content-type", + "transfer-encoding", + "connection", + "date", +]; + +fn protected(name: &str) -> bool { + PROTECTED + .iter() + .any(|protected| protected.eq_ignore_ascii_case(name)) + || name + .get(..6) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("x-amz-")) +} + +/// The hooks `litellm_settings.hooks` lists, in its order. +pub fn from_config(config: &Config) -> Result>, Error> { + let Some(entries) = config.litellm_settings.additional_fields.get("hooks") else { + return Ok(Vec::new()); + }; + let litellm_config::Value::Sequence(entries) = entries else { + return Err(Error::Hooks("must be a list".into())); + }; + entries.iter().map(hook).collect() +} + +fn hook(entry: &litellm_config::Value) -> Result, Error> { + let (name, settings) = match entry { + litellm_config::Value::Mapping(mapping) if mapping.len() == 1 => { + let (name, settings) = mapping.iter().next().expect("one entry"); + (name.as_str().unwrap_or_default(), settings.clone()) + } + litellm_config::Value::String(name) => (name.as_str(), litellm_config::Value::Null), + _ => { + return Err(Error::Hooks( + "each entry is a hook name with its settings".into(), + )); + } + }; + let built: Result, String> = match name { + "forward_headers" => parse(settings) + .and_then(patterns) + .map(|patterns| Arc::new(ForwardHeaders { patterns }) as Arc), + "strip_headers" => parse(settings) + .and_then(patterns) + .map(|patterns| Arc::new(StripHeaders { patterns }) as Arc), + "spend_metadata" => parse(settings) + .and_then(SpendMetadata::new) + .map(|hook| Arc::new(hook) as Arc), + "provider_metadata" => parse(settings) + .and_then(ProviderMetadata::new) + .map(|hook| Arc::new(hook) as Arc), + other => return Err(Error::Hooks(format!("unknown hook {other}"))), + }; + built.map_err(|error| Error::Hooks(format!("{name}: {error}"))) +} + +fn parse(settings: litellm_config::Value) -> Result { + T::deserialize(settings).map_err(|error| error.to_string()) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct HeaderList { + headers: Vec, +} + +fn patterns(list: HeaderList) -> Result, String> { + list.headers + .iter() + .map(|pattern| { + HeaderPattern::parse(pattern).map_err(|_| format!("{pattern} is not a header name")) + }) + .collect() +} + +/// `forward_headers`: copies the client headers it names onto the provider request. +struct ForwardHeaders { + patterns: Vec, +} + +impl Hook for ForwardHeaders { + fn name(&self) -> &'static str { + "forward_headers" + } + + fn reads(&self) -> Facts { + Facts::CLIENT_HEADERS + } + + fn writes(&self) -> Targets { + Targets::UPSTREAM_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + for (name, value) in request.client_headers.iter() { + if !protected(name.as_str()) + && self + .patterns + .iter() + .any(|pattern| pattern.matches(name.as_str())) + { + edit.set_upstream_header(name.clone(), value.clone()); + } + } + Box::pin(async { Ok(None) }) + } +} + +/// `strip_headers`: removes the headers it names from the provider request, whether a hook set +/// them or the gateway forwards them by default. +struct StripHeaders { + patterns: Vec, +} + +impl Hook for StripHeaders { + fn name(&self) -> &'static str { + "strip_headers" + } + + fn reads(&self) -> Facts { + Facts::NONE + } + + fn writes(&self) -> Targets { + Targets::UPSTREAM_HEADERS + } + + fn run<'a>( + &'a self, + _request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + for pattern in &self.patterns { + edit.remove_upstream_headers(pattern.clone()); + } + Box::pin(async { Ok(None) }) + } +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct SpendMetadataSettings { + /// Header name to `json`, merged as an object, or to the key it is stored under. + #[serde(default)] + from_headers: litellm_config::Value, + /// Fields of the key's metadata copied under the same names. + #[serde(default)] + from_key_metadata: Vec, +} + +enum HeaderRule { + Json, + Field(String), +} + +/// `spend_metadata`: adds caller-supplied values to the request's `spend_logs_metadata` +/// without changing the request. +struct SpendMetadata { + headers: Vec<(HeaderName, HeaderRule)>, + key_fields: Vec, +} + +impl SpendMetadata { + fn new(settings: SpendMetadataSettings) -> Result { + let headers = match settings.from_headers { + litellm_config::Value::Null => Vec::new(), + litellm_config::Value::Mapping(mapping) => mapping + .iter() + .map(|(name, rule)| { + let (Some(name), Some(rule)) = (name.as_str(), rule.as_str()) else { + return Err("from_headers maps header names to json or a key".to_owned()); + }; + let name = HeaderName::from_bytes(name.as_bytes()) + .map_err(|_| format!("{name} is not a header name"))?; + let rule = match rule { + "json" => HeaderRule::Json, + field => HeaderRule::Field(field.to_owned()), + }; + Ok((name, rule)) + }) + .collect::>()?, + _ => return Err("from_headers maps header names to json or a key".into()), + }; + Ok(Self { + headers, + key_fields: settings.from_key_metadata, + }) + } +} + +impl Hook for SpendMetadata { + fn name(&self) -> &'static str { + "spend_metadata" + } + + fn reads(&self) -> Facts { + Facts::CLIENT_HEADERS | Facts::KEY + } + + fn writes(&self) -> Targets { + Targets::ANNOTATIONS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + let mut metadata = Map::new(); + for (name, rule) in &self.headers { + let Some(value) = request + .client_headers + .get(name) + .and_then(|value| value.to_str().ok()) + else { + continue; + }; + match rule { + HeaderRule::Json => match serde_json::from_str::(value) { + Ok(Value::Object(fields)) => metadata.extend(fields), + _ => tracing::warn!( + header = name.as_str(), + "spend metadata header is not a JSON object; ignored" + ), + }, + HeaderRule::Field(field) => { + metadata.insert(field.clone(), Value::from(value)); + } + } + } + // The key's own values win over what the caller sent. + if let Some(key) = request.key.as_deref() { + for field in &self.key_fields { + if let Some(value) = key.metadata.get(field) { + metadata.insert(field.clone(), value.clone()); + } + } + } + if !metadata.is_empty() { + edit.annotate(SPEND_LOGS_METADATA, Value::Object(metadata)); + } + Box::pin(async { Ok(None) }) + } +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct ProviderMetadataSettings { + /// Dotted body paths, such as `metadata.user_id`, to templates. + #[serde(default)] + body: std::collections::BTreeMap, + #[serde(default)] + headers: std::collections::BTreeMap, +} + +/// A fact a template can name. +#[derive(Clone, Debug, PartialEq)] +enum Fact { + EndUser, + Model, + RequestId, + KeyAlias, + KeyUserId, + KeyTeamId, + KeyTeamAlias, + KeyMetadata(String), + TeamMetadata(String), + Header(HeaderName), +} + +impl Fact { + fn parse(name: &str) -> Result { + Ok(match name { + "end_user" => Self::EndUser, + "model" => Self::Model, + "request_id" => Self::RequestId, + "key.alias" => Self::KeyAlias, + "key.user_id" => Self::KeyUserId, + "key.team_id" => Self::KeyTeamId, + "key.team_alias" => Self::KeyTeamAlias, + other => { + if let Some(field) = other.strip_prefix("key.metadata.") { + Self::KeyMetadata(field.to_owned()) + } else if let Some(field) = other.strip_prefix("team.metadata.") { + Self::TeamMetadata(field.to_owned()) + } else if let Some(header) = other.strip_prefix("header.") { + Self::Header( + HeaderName::from_bytes(header.as_bytes()) + .map_err(|_| format!("{header} is not a header name"))?, + ) + } else { + return Err(format!("unknown fact {other}")); + } + } + }) + } + + fn value(&self, request: &RequestFacts) -> Option { + let key = request.key.as_deref(); + let scalar = |value: &Value| match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + }; + match self { + Self::EndUser => request.end_user.as_deref().map(str::to_owned), + Self::Model => Some(request.requested_model.to_string()), + Self::RequestId => Some(request.request_id.to_string()), + Self::KeyAlias => key?.alias.clone(), + Self::KeyUserId => key?.user_id.clone(), + Self::KeyTeamId => key?.team_id.clone(), + Self::KeyTeamAlias => key?.team_alias.clone(), + Self::KeyMetadata(field) => key?.metadata.get(field).and_then(scalar), + Self::TeamMetadata(field) => key?.team_metadata.get(field).and_then(scalar), + Self::Header(name) => request + .client_headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned), + } + .filter(|value| !value.is_empty()) + } +} + +#[derive(Debug, PartialEq)] +enum Part { + Text(String), + /// The first of these facts that is set. + Facts(Vec), +} + +/// Text with `{fact|fallback}` placeholders, parsed once. A placeholder none of whose facts are +/// set leaves the whole value unset. +#[derive(Debug, PartialEq)] +struct Template(Vec); + +impl Template { + fn parse(source: &str) -> Result { + let mut parts = Vec::new(); + let mut rest = source; + while let Some(open) = rest.find('{') { + if open > 0 { + parts.push(Part::Text(rest[..open].to_owned())); + } + let close = rest[open..] + .find('}') + .ok_or_else(|| format!("unclosed placeholder in {source}"))?; + let facts = rest[open + 1..open + close] + .split('|') + .map(|fact| Fact::parse(fact.trim())) + .collect::>()?; + parts.push(Part::Facts(facts)); + rest = &rest[open + close + 1..]; + } + if !rest.is_empty() { + parts.push(Part::Text(rest.to_owned())); + } + Ok(Self(parts)) + } + + fn reads(&self) -> Facts { + let mut reads = Facts::NONE; + for part in &self.0 { + let Part::Facts(facts) = part else { continue }; + for fact in facts { + reads |= match fact { + Fact::EndUser => Facts::END_USER, + Fact::Header(_) => Facts::CLIENT_HEADERS, + Fact::Model | Fact::RequestId => Facts::NONE, + _ => Facts::KEY, + }; + } + } + reads + } + + fn render(&self, request: &RequestFacts) -> Option { + let mut rendered = String::new(); + for part in &self.0 { + match part { + Part::Text(text) => rendered.push_str(text), + Part::Facts(facts) => { + rendered.push_str(&facts.iter().find_map(|fact| fact.value(request))?) + } + } + } + Some(rendered) + } +} + +/// `provider_metadata`: sets body fields and headers on the provider request from facts about +/// the caller. Body fields are set on `/v1/messages` requests only. +struct ProviderMetadata { + body: Vec<(Vec, Template)>, + headers: Vec<(HeaderName, Template)>, + reads: Facts, +} + +impl ProviderMetadata { + fn new(settings: ProviderMetadataSettings) -> Result { + let body = settings + .body + .iter() + .map(|(path, template)| { + let path: Vec = path.split('.').map(str::to_owned).collect(); + if path.iter().any(String::is_empty) { + return Err(format!("{} is not a body path", path.join("."))); + } + Ok((path, Template::parse(template)?)) + }) + .collect::, String>>()?; + let headers = settings + .headers + .iter() + .map(|(name, template)| { + let header = HeaderName::from_bytes(name.as_bytes()) + .map_err(|_| format!("{name} is not a header name"))?; + if protected(header.as_str()) { + return Err(format!("{name} cannot be set by a hook")); + } + Ok((header, Template::parse(template)?)) + }) + .collect::, String>>()?; + let reads = body + .iter() + .map(|(_, template)| template) + .chain(headers.iter().map(|(_, template)| template)) + .fold(Facts::NONE, |reads, template| reads | template.reads()); + Ok(Self { + body, + headers, + reads, + }) + } +} + +impl Hook for ProviderMetadata { + fn name(&self) -> &'static str { + "provider_metadata" + } + + fn reads(&self) -> Facts { + self.reads + } + + fn writes(&self) -> Targets { + Targets::UPSTREAM_BODY | Targets::UPSTREAM_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + // Body paths name Messages API fields; Chat Completions takes the headers only. + let body = match request.route { + Route::Messages => self.body.as_slice(), + Route::ChatCompletions => &[], + }; + for (path, template) in body { + if let Some(value) = template.render(request) { + edit.set_upstream_body_path(path.clone(), Value::from(value)); + } + } + for (name, template) in &self.headers { + if let Some(value) = template + .render(request) + .and_then(|value| HeaderValue::from_str(&value).ok()) + { + edit.set_upstream_header(name.clone(), value); + } + } + Box::pin(async { Ok(None) }) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + fn templates_name_facts_with_fallbacks() { + assert_eq!( + Template::parse("team-{key.team_alias|key.team_id}").unwrap(), + Template(vec![ + Part::Text("team-".into()), + Part::Facts(vec![Fact::KeyTeamAlias, Fact::KeyTeamId]), + ]) + ); + assert!(Template::parse("{nope}").is_err()); + assert!(Template::parse("{end_user").is_err()); + } + + #[rstest] + #[case::credential("Authorization", true)] + #[case::signed("X-Amz-Date", true)] + #[case::trace("x-trace-id", false)] + fn credentials_are_never_forwarded(#[case] name: &str, #[case] expected: bool) { + assert_eq!(protected(name), expected); + } +} diff --git a/litellm-rust/crates/gateway/tests/request_hooks.rs b/litellm-rust/crates/gateway/tests/request_hooks.rs new file mode 100644 index 00000000000..af8a74b3079 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/request_hooks.rs @@ -0,0 +1,339 @@ +//! The configurable request hooks on real requests: headers forwarded and stripped, provider +//! metadata set from the caller, and spend metadata handed to observers. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +use axum::{ + Router, + body::{Body, Bytes, to_bytes}, + http::{HeaderMap, Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; + +use futures_util::future::BoxFuture; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, +}; +use litellm_gateway_hooks::{ + Facts, KeyFacts, Observer, Observers, Pipeline, Record, RequestFacts, Route, +}; +use rstest::rstest; +use serde_json::{Map, Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; + +type Seen = Arc>>; + +async fn upstream() -> (String, Seen) { + let seen = Seen::default(); + let recorded = seen.clone(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(move |request: Request| { + let recorded = recorded.clone(); + async move { + let headers = request.headers().clone(); + let body = to_bytes(request.into_body(), 1 << 20).await.unwrap(); + recorded.lock().unwrap().push(( + headers, + serde_json::from_slice(&body).unwrap_or(Value::Null), + )); + ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response() + } + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + (format!("http://{address}"), seen) +} + +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()) }) + } +} + +const HOOKS: &str = r#" + hooks: + - forward_headers: + headers: ["x-trace-*", "x-cost-center", "authorization"] + - strip_headers: + headers: ["x-cost-center", "anthropic-beta"] + - provider_metadata: + body: + metadata.user_id: "{end_user|key.user_id}" + headers: + x-caller-team: "team-{key.team_alias}" +"#; + +async fn gateway(hooks: &str) -> (Router, Seen) { + let (address, seen) = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {address} + api_key: sk-upstream +general_settings: + master_key: sk-master +litellm_settings:{hooks}" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let hooks = litellm_gateway::request_hooks::from_config(&config).unwrap(); + litellm_gateway::hooks::install(&inference, &config, hooks, Vec::new()) + .await + .unwrap(); + let key = KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: Some("team-1".into()), + profile: Arc::new(KeyProfile { + user_id: Some("user-1".into()), + team_alias: Some("research".into()), + ..KeyProfile::default() + }), + }; + let keys: Arc = Arc::new(Keys(HashMap::from([( + KeyHash::from_token(&SecretValue::new("sk-key")), + key, + )]))); + let app = litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys)); + (app, seen) +} + +async fn send(app: &Router, path: &str, body: Value) -> u16 { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer sk-key") + .header("x-trace-id", "trace-1") + .header("x-cost-center", "cc-9") + .header("anthropic-beta", "some-beta") + .header("x-other", "kept-out") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + response.status().as_u16() +} + +fn header<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> { + headers.get(name).map(|value| value.to_str().unwrap()) +} + +#[rstest] +#[tokio::test] +async fn messages_carry_the_configured_headers_and_metadata() { + let (app, seen) = gateway(HOOKS).await; + let body = json!({ + "model": "claude", + "max_tokens": 8, + "metadata": {"tag": "t"}, + "messages": [{"role": "user", "content": "hi"}], + }); + assert_eq!(send(&app, "/v1/messages", body).await, 200); + let (headers, body) = seen.lock().unwrap().pop().unwrap(); + assert_eq!(header(&headers, "x-trace-id"), Some("trace-1")); + assert_eq!(header(&headers, "x-caller-team"), Some("team-research")); + // Stripped after forwarding, and the gateway's own pass-through is stripped too. + assert_eq!(header(&headers, "x-cost-center"), None); + assert_eq!(header(&headers, "anthropic-beta"), None); + assert_eq!(header(&headers, "x-other"), None); + // Credentials are never forwarded, even when listed. + assert_ne!(header(&headers, "authorization"), Some("Bearer sk-key")); + assert_eq!(header(&headers, "x-api-key"), Some("sk-upstream")); + // The provider keeps only the `metadata` fields it knows. + assert_eq!(body["metadata"]["user_id"], "user-1"); +} + +#[rstest] +#[tokio::test] +async fn the_end_user_comes_before_the_key_user() { + let (app, seen) = gateway(HOOKS).await; + let body = json!({ + "model": "claude", + "max_tokens": 8, + "litellm_metadata": {"user": "end-7"}, + "messages": [{"role": "user", "content": "hi"}], + }); + assert_eq!(send(&app, "/v1/messages", body).await, 200); + let (_, body) = seen.lock().unwrap().pop().unwrap(); + assert_eq!(body["metadata"]["user_id"], "end-7"); +} + +#[rstest] +#[tokio::test] +async fn chat_completions_take_the_headers_only() { + let (app, seen) = gateway(HOOKS).await; + let body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + assert_eq!(send(&app, "/v1/chat/completions", body).await, 200); + let (headers, body) = seen.lock().unwrap().pop().unwrap(); + assert_eq!(header(&headers, "x-trace-id"), Some("trace-1")); + assert_eq!(header(&headers, "x-caller-team"), Some("team-research")); + assert_eq!(header(&headers, "x-cost-center"), None); + assert_eq!(body.get("metadata"), None); +} + +#[rstest] +#[tokio::test] +async fn without_hooks_the_default_pass_through_is_unchanged() { + let (app, seen) = gateway("\n drop_params: true\n").await; + let body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + assert_eq!(send(&app, "/v1/messages", body).await, 200); + let (headers, _) = seen.lock().unwrap().pop().unwrap(); + assert_eq!(header(&headers, "anthropic-beta"), Some("some-beta")); + assert_eq!(header(&headers, "x-trace-id"), None); +} + +struct Annotations(Mutex>>); + +impl Observer for Annotations { + fn name(&self) -> &'static str { + "annotations" + } + + fn reads(&self) -> Facts { + Facts::ANNOTATIONS + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + Box::pin(async move { + self.0.lock().unwrap().push( + record + .annotations + .get(litellm_gateway::spend::SPEND_LOGS_METADATA) + .cloned(), + ); + }) + } +} + +#[rstest] +#[tokio::test] +async fn spend_metadata_reaches_observers() { + let config = Config::from_yaml( + r#" +model_list: [] +litellm_settings: + hooks: + - spend_metadata: + from_headers: + x-litellm-spend-logs-metadata: json + x-cost-center: cost_center + from_key_metadata: [owner] +"#, + ) + .unwrap(); + let observed = Arc::new(Annotations(Mutex::default())); + let pipeline = Pipeline::new( + litellm_gateway::request_hooks::from_config(&config).unwrap(), + Observers::start(vec![observed.clone()], 8), + Arc::default(), + ); + let run = async |headers: &[(&'static str, &'static str)]| { + let mut map = HeaderMap::new(); + for (name, value) in headers { + map.insert(*name, value.parse().unwrap()); + } + let Value::Object(metadata) = json!({"owner": "platform"}) else { + unreachable!() + }; + let key = KeyFacts { + metadata: Arc::new(metadata), + ..KeyFacts::default() + }; + let request = RequestFacts::new( + "r".into(), + Route::Messages, + Some(Arc::new(key)), + None, + "claude".into(), + Arc::new(map), + false, + None, + Bytes::from_static(b"{}"), + ); + let admission = pipeline.admit(Arc::new(request)).await.unwrap(); + pipeline.finish( + admission, + litellm_gateway_hooks::Outcome { + deployment: None, + status: 200, + usage: None, + cost: None, + time_to_first_byte: None, + ended_at: std::time::SystemTime::now(), + interrupted: false, + }, + None, + ); + }; + run(&[ + ( + "x-litellm-spend-logs-metadata", + r#"{"project": "p1", "owner": "caller"}"#, + ), + ("x-cost-center", "cc-9"), + ]) + .await; + run(&[("x-litellm-spend-logs-metadata", "not json")]).await; + pipeline.close().await; + let mut expected = Map::new(); + expected.insert("project".into(), json!("p1")); + expected.insert("owner".into(), json!("platform")); + expected.insert("cost_center".into(), json!("cc-9")); + assert_eq!( + *observed.0.lock().unwrap(), + [ + Some(Value::Object(expected)), + Some(json!({"owner": "platform"})) + ] + ); +} + +#[rstest] +#[case::unknown_hook(" - rewrite_everything: {}\n", "unknown hook rewrite_everything")] +#[case::bad_template( + " - provider_metadata:\n headers:\n x-a: \"{key.secret}\"\n", + "provider_metadata: unknown fact key.secret" +)] +#[case::credential_header( + " - provider_metadata:\n headers:\n authorization: \"{key.user_id}\"\n", + "provider_metadata: authorization cannot be set by a hook" +)] +#[case::unknown_setting( + " - strip_headers:\n names: [x-a]\n", + "strip_headers: unknown field `names`" +)] +fn misconfigured_hooks_stop_the_gateway(#[case] hooks: &str, #[case] message: &str) { + let config = Config::from_yaml(&format!( + "model_list: []\nlitellm_settings:\n hooks:\n{hooks}" + )) + .unwrap(); + let error = litellm_gateway::request_hooks::from_config(&config) + .err() + .unwrap() + .to_string(); + assert!(error.contains(message), "{error}"); +} diff --git a/litellm-rust/docs/customer-proxy/request-hooks.svg b/litellm-rust/docs/customer-proxy/request-hooks.svg new file mode 100644 index 00000000000..74e96af77e9 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/request-hooks.svg @@ -0,0 +1,56 @@ + + + + + + Configurable request hooks: litellm_settings.hooks, run top to bottom + + + Client request + x-trace-id: trace-1 + x-cost-center: cc-9 + x-litellm-spend-logs- + metadata: {...} + anthropic-beta: ... + authorization: Bearer sk-... + key: user, team alias, + metadata; end user + + + + forward_headers copies listed client headers (exact or prefix-*) + + spend_metadata headers (json or one key) + key metadata + + strip_headers removes forwarded and pass-through headers + + provider_metadata "{end_user|key.user_id}" templates + + + + + + Provider request + x-trace-id: trace-1 + x-caller-team: research + metadata.user_id: user-1 + (messages only) + no x-cost-center, no + anthropic-beta, never + client credentials + + + + Spend log + the spend_metadata annotation lands in + LiteLLM_SpendLogs.metadata.spend_logs_metadata + + + + Rules + Each hook declares what it reads and writes; a later hook in the list wins on the same header or field. + Header edits reach the provider on both endpoints and every provider, Bedrock's signed requests included. + Credentials, signer and framing headers (authorization, x-api-key, cookie, x-amz-*, ...) are never forwarded or set. + An unknown hook, setting or template fact stops the gateway at startup; a bad JSON header is dropped with a warning. + Templates name facts only: end_user, model, request_id, key.alias|user_id|team_id|team_alias, key.metadata.*, team.metadata.*, header.* + From 16dacb02b79224db83ba6f0f294e1d720b2f23cb Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:45 -0700 Subject: [PATCH 06/10] feat(gateway): model capacity hook, the proxy's dynamic rate limiter (#47055) Adds the model_capacity hook, turned on by listing dynamic_rate_limiter_v3 (or model_capacity) in litellm_settings.callbacks. It caps each model group at the sum of its deployments' rpm and tpm across every replica, and with priority_reservation holds each priority's pool to its share once the group's saturation reaches the threshold, counting it below. It uses the proxy's model_saturation_check and priority_model keys and scripts, its refusal messages and headers, and adds a successful response's tokens to the group and the pool admission chose. The rate limit counters gain a general check-and-add over request and token pairs, and the throttling refusal shape is shared. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/crates/gateway/src/lib.rs | 1 + litellm-rust/crates/gateway/src/main.rs | 5 + .../crates/gateway/src/model_capacity.rs | 654 ++++++++++++++++++ .../crates/gateway/src/rate_limits.rs | 232 +++++-- .../crates/gateway/tests/model_capacity.rs | 306 ++++++++ .../docs/customer-proxy/model-capacity.svg | 52 ++ 6 files changed, 1177 insertions(+), 73 deletions(-) create mode 100644 litellm-rust/crates/gateway/src/model_capacity.rs create mode 100644 litellm-rust/crates/gateway/tests/model_capacity.rs create mode 100644 litellm-rust/docs/customer-proxy/model-capacity.svg diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index c97c1d6515b..c21352cde86 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -9,6 +9,7 @@ pub mod hooks; mod key_status; mod keys; pub mod metrics; +pub mod model_capacity; pub mod rate_limits; pub mod request_hooks; mod secrets; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index b8a35eddbcd..b242a3b88ab 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -128,6 +128,11 @@ 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)); 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/rate_limits.rs b/litellm-rust/crates/gateway/src/rate_limits.rs index 2b9ca57cbbc..2948e3bc842 100644 --- a/litellm-rust/crates/gateway/src/rate_limits.rs +++ b/litellm-rust/crates/gateway/src/rate_limits.rs @@ -225,14 +225,14 @@ impl Descriptor { } #[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum Kind { +pub(crate) enum Kind { Requests, Tokens, Parallel, } impl Kind { - fn name(self) -> &'static str { + pub(crate) fn name(self) -> &'static str { match self { Self::Requests => "requests", Self::Tokens => "tokens", @@ -250,15 +250,25 @@ struct Status { remaining: i64, } -/// The token reservation went over on the descriptor at `index`, which held `current`. +/// A counter to check and add to: `{prefix}:requests` or `{prefix}:tokens` in its window. +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct Pair { + pub(crate) prefix: String, + pub(crate) kind: Kind, + pub(crate) limit: i64, + pub(crate) increment: i64, +} + +/// Every pair's count after adding, or the pair at `index` that would go over, which held +/// `current`. #[derive(Debug, PartialEq)] -enum Reserved { +pub(crate) enum Reserved { Counts(Vec), Over { index: usize, current: i64 }, } /// Redis could not be used and the settings say to refuse rather than fall back. -struct Unavailable; +pub(crate) struct Unavailable; #[derive(Debug)] struct Window { @@ -307,27 +317,24 @@ impl Local { .collect() } - fn reserve( - &mut self, - scopes: &[(String, i64)], - increment: i64, - now: u64, - size: u64, - ) -> Reserved { - for (index, (prefix, limit)) in scopes.iter().enumerate() { - let current = match self.windows.get(prefix) { - Some(window) if now.saturating_sub(window.start) < size => window.tokens, + fn check_and_add(&mut self, pairs: &[Pair], now: u64, size: u64) -> Reserved { + for (index, pair) in pairs.iter().enumerate() { + let current = match self.windows.get(&pair.prefix) { + Some(window) if now.saturating_sub(window.start) < size => match pair.kind { + Kind::Requests => window.requests, + _ => window.tokens, + }, _ => 0, }; - if blocked(current, increment, *limit) { + if blocked(current, pair.increment, pair.limit) { return Reserved::Over { index, current }; } } Reserved::Counts( - scopes + pairs .iter() - .map(|(prefix, _)| { - let (window, expired) = self.window(prefix, now, size); + .map(|pair| { + let (window, expired) = self.window(&pair.prefix, now, size); if expired { *window = Window { start: now, @@ -335,13 +342,26 @@ impl Local { tokens: 0, }; } - window.tokens += increment; - window.tokens + let count = match pair.kind { + Kind::Requests => &mut window.requests, + _ => &mut window.tokens, + }; + *count += pair.increment; + *count }) .collect(), ) } + /// A counter's value in its current window. + fn read(&self, prefix: &str, kind: Kind, now: u64, size: u64) -> Option { + let window = self.windows.get(prefix)?; + (now.saturating_sub(window.start) < size).then_some(match kind { + Kind::Requests => window.requests, + _ => window.tokens, + }) + } + fn charge(&mut self, deltas: &[(String, i64)]) { for (prefix, delta) in deltas { if let Some(window) = self.windows.get_mut(prefix) { @@ -380,7 +400,7 @@ fn blocked(current: i64, increment: i64, limit: i64) -> bool { } /// Where the counters live. -struct Store { +pub(crate) struct Store { redis: Option>, local: Mutex, window: u64, @@ -388,6 +408,15 @@ struct Store { } impl Store { + pub(crate) fn new(redis: Option>, window: u64, fail_closed: bool) -> Self { + Self { + redis, + local: Mutex::default(), + window, + fail_closed, + } + } + fn connection(&self) -> Option<(ConnectionManager, &Coordination)> { let redis = self.redis.as_deref()?; Some((redis.connection()?, redis)) @@ -414,7 +443,7 @@ impl Store { } /// Counts the request on each prefix's window and returns the counts. - async fn requests(&self, prefixes: &[String]) -> Result, Unavailable> { + pub(crate) async fn requests(&self, prefixes: &[String]) -> Result, Unavailable> { let now = epoch_seconds(); if self.redis.is_some() { let error = match self.connection() { @@ -500,21 +529,43 @@ impl Store { &self, scopes: &[(String, i64)], increment: i64, + ) -> Result<(Reserved, bool), Unavailable> { + let pairs: Vec = scopes + .iter() + .map(|(prefix, limit)| Pair { + prefix: prefix.clone(), + kind: Kind::Tokens, + limit: *limit, + increment, + }) + .collect(); + self.check_and_add(&pairs).await + } + + /// Adds every pair's increment, or none when one would go over its limit. `true` when the + /// counters are in Redis. + pub(crate) async fn check_and_add( + &self, + pairs: &[Pair], ) -> Result<(Reserved, bool), Unavailable> { if self.redis.is_some() { let error = match self.connection() { Some((mut connection, redis)) => { let mut command = redis::cmd("EVAL"); - command.arg(RESERVE).arg(scopes.len() * 2); - for (prefix, _) in scopes { + command.arg(RESERVE).arg(pairs.len() * 2); + for pair in pairs { command - .arg(redis.namespaced(&format!("{prefix}:window"))) - .arg(redis.namespaced(&format!("{prefix}:tokens"))); + .arg(redis.namespaced(&format!("{}:window", pair.prefix))) + .arg(redis.namespaced(&format!( + "{}:{}", + pair.prefix, + pair.kind.name() + ))); } - for (_, limit) in scopes { + for pair in pairs { command - .arg(*limit) - .arg(increment) + .arg(pair.limit) + .arg(pair.increment) .arg(self.window) .arg(self.window); } @@ -528,7 +579,7 @@ impl Store { )); } Ok(reply) => match reply.as_slice() { - [1, index, current, _] if (1..=scopes.len() as i64).contains(index) => { + [1, index, current, _] if (1..=pairs.len() as i64).contains(index) => { return Ok(( Reserved::Over { index: (*index - 1) as usize, @@ -548,13 +599,34 @@ impl Store { } Ok(( self.local() - .reserve(scopes, increment, epoch_seconds(), self.window), + .check_and_add(pairs, epoch_seconds(), self.window), false, )) } + /// A counter's current value, `None` when it is not set. + pub(crate) async fn read(&self, prefix: &str, kind: Kind) -> Option { + if let Some((mut connection, redis)) = self.connection() { + let key = redis.namespaced(&format!("{prefix}:{}", kind.name())); + match redis::cmd("GET") + .arg(key) + .query_async::>(&mut connection) + .await + { + Ok(value) => { + return value + .and_then(|value| value.parse::().ok()) + .map(|value| value as i64); + } + Err(error) => tracing::warn!(%error, "reading a rate limit counter failed"), + } + } + self.local() + .read(prefix, kind, epoch_seconds(), self.window) + } + /// Adds each delta to its prefix's tokens. - async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) { + pub(crate) async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) { if shared && let Some((mut connection, redis)) = self.connection() { let mut command = redis::cmd("EVAL"); command.arg(CHARGE).arg(deltas.len()); @@ -711,12 +783,7 @@ impl RateLimiter { database: Option, ) -> Self { Self { - store: Arc::new(Store { - redis, - local: Mutex::default(), - window: settings.window, - fail_closed: settings.fail_closed, - }), + store: Arc::new(Store::new(redis, settings.window, settings.fail_closed)), settings, end_users: database.map(EndUsers::new), } @@ -801,22 +868,15 @@ impl RateLimiter { status.limit, status.remaining.max(0), ); - let mut refusal = match request.route { - Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(), - Route::ChatCompletions => { - let mut refusal = Refusal::new(429, "throttling_error", message).without_param(); - for (name, value) in [ - ("retry-after", self.settings.window.to_string()), - ("rate_limit_type", status.kind.name().to_owned()), - ("reset_at", reset_at), - ] { - if let Ok(value) = HeaderValue::from_str(&value) { - refusal = refusal.with_header(HeaderName::from_static(name), value); - } - } - refusal - } - }; + let mut refusal = throttled( + request.route, + message, + vec![ + ("retry-after", self.settings.window.to_string()), + ("rate_limit_type", status.kind.name().to_owned()), + ("reset_at", reset_at), + ], + ); for (name, value) in key_limit_headers(request.key.as_deref()) { refusal = refusal.with_header(name, value); } @@ -824,11 +884,37 @@ impl RateLimiter { } fn unavailable(&self, request: &RequestFacts) -> Refusal { - match request.route { - Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(), - Route::ChatCompletions => { - Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param() + unavailable(request.route) + } +} + +/// The proxy's 429 from one of its own limiters: on Chat Completions its error body with the +/// limiter's headers, on `/v1/messages` Anthropic's envelope, where the proxy drops them. +pub(crate) fn throttled( + route: Route, + message: String, + headers: Vec<(&'static str, String)>, +) -> Refusal { + match route { + Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(), + Route::ChatCompletions => { + let mut refusal = Refusal::new(429, "throttling_error", message).without_param(); + for (name, value) in headers { + if let Ok(value) = HeaderValue::from_str(&value) { + refusal = refusal.with_header(HeaderName::from_static(name), value); + } } + refusal + } + } +} + +/// The proxy's 503 when counters cannot be verified and enforcement fails closed. +pub(crate) fn unavailable(route: Route) -> Refusal { + match route { + Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(), + Route::ChatCompletions => { + Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param() } } } @@ -1150,35 +1236,35 @@ mod tests { assert_eq!(estimate(&request, Some(&key), "other", 100_000).tokens, 10); } + fn reserve(local: &mut Local, increment: i64, now: u64) -> Reserved { + let pair = Pair { + prefix: "{api_key:h}".into(), + kind: Kind::Tokens, + limit: 10, + increment, + }; + local.check_and_add(&[pair], now, 60) + } + #[rstest] fn local_windows_count_reserve_and_reset() { let mut local = Local::default(); let prefixes = vec!["{api_key:h}".to_owned()]; assert_eq!(local.requests(&prefixes, 100, 60), [1]); assert_eq!(local.requests(&prefixes, 159, 60), [2]); - let scopes = vec![("{api_key:h}".to_owned(), 10)]; + assert_eq!(reserve(&mut local, 6, 159), Reserved::Counts(vec![6])); assert_eq!( - local.reserve(&scopes, 6, 159, 60), - Reserved::Counts(vec![6]) - ); - assert_eq!( - local.reserve(&scopes, 6, 159, 60), + reserve(&mut local, 6, 159), Reserved::Over { index: 0, current: 6 } ); local.charge(&[("{api_key:h}".to_owned(), -4)]); - assert_eq!( - local.reserve(&scopes, 8, 159, 60), - Reserved::Counts(vec![10]) - ); + assert_eq!(reserve(&mut local, 8, 159), Reserved::Counts(vec![10])); // The window started at 100 is over at 160: requests and tokens start again. assert_eq!(local.requests(&prefixes, 160, 60), [1]); - assert_eq!( - local.reserve(&scopes, 10, 160, 60), - Reserved::Counts(vec![10]) - ); + assert_eq!(reserve(&mut local, 10, 160), Reserved::Counts(vec![10])); } #[rstest] diff --git a/litellm-rust/crates/gateway/tests/model_capacity.rs b/litellm-rust/crates/gateway/tests/model_capacity.rs new file mode 100644 index 00000000000..0dddd6ad901 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/model_capacity.rs @@ -0,0 +1,306 @@ +//! The model capacity hook on real requests: a group held to its deployments' summed capacity, +//! priorities held to their share, refusals and headers as the proxy sends them, and counters +//! shared through Redis. Set `LITELLM_TEST_REDIS_URL` to also run the shared case. + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{HeaderMap, Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway::{ + coordination::{Coordination, RedisSettings}, + model_capacity::{CapacitySettings, ModelCapacity}, +}; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +/// 3 input and 4 output tokens: 7 count against the group. +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; + +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| async { + ([(CONTENT_TYPE, "application/json")], MESSAGE).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()) }) + } +} + +/// Two deployments of `claude`, each with `limits` (such as `rpm: 1`). +async fn gateway(limits: &str, settings: &str, redis: Option>) -> Router { + let address = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: {{model: anthropic/claude-x, api_base: {address}, api_key: sk-a, {limits}}} + - model_name: claude + litellm_params: {{model: anthropic/claude-y, api_base: {address}, api_key: sk-b, {limits}}} +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +litellm_settings: + callbacks: [dynamic_rate_limiter_v3] +{settings}" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let hook = ModelCapacity::for_gateway(&config, redis).unwrap(); + litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(hook)], Vec::new()) + .await + .unwrap(); + let mut metadata = serde_json::Map::new(); + metadata.insert("priority".into(), json!("prod")); + let key = KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: None, + profile: Arc::new(KeyProfile { + metadata: Arc::new(metadata), + ..KeyProfile::default() + }), + }; + let keys: Arc = Arc::new(Keys(HashMap::from([( + KeyHash::from_token(&SecretValue::new("sk-prod")), + key, + )]))); + litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys)) +} + +struct Reply { + status: u16, + headers: HeaderMap, + body: Value, +} + +async fn send(app: &Router, path: &str, token: &str) -> Reply { + let body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {token}")) + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + Reply { + status, + headers, + body: serde_json::from_slice(&body).unwrap_or(Value::Null), + } +} + +fn header<'a>(reply: &'a Reply, name: &str) -> Option<&'a str> { + reply.headers.get(name).map(|value| value.to_str().unwrap()) +} + +async fn settled() { + tokio::time::sleep(Duration::from_millis(100)).await; +} + +#[rstest] +#[tokio::test] +async fn a_group_is_held_to_its_deployments_capacity() { + let app = gateway("rpm: 1", "", None).await; + let first = send(&app, "/v1/messages", "sk-master").await; + assert_eq!(first.status, 200); + assert_eq!( + header( + &first, + "x-ratelimit-model_saturation_check-remaining-requests" + ), + Some("1") + ); + assert_eq!( + header(&first, "x-ratelimit-model_saturation_check-limit-requests"), + Some("2") + ); + assert_eq!(header(&first, "x-litellm-rate-limiter-version"), Some("v3")); + assert_eq!(header(&first, "x-litellm-priority"), Some("default")); + assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200); + + let message = "Model capacity reached for claude. Priority: None, Rate limit type: requests, Model TPM: not configured, Model RPM: 2, Remaining: 0"; + let refused = send(&app, "/v1/chat/completions", "sk-master").await; + assert_eq!(refused.status, 429); + assert_eq!( + refused.body, + json!({"error": {"message": message, "type": "throttling_error", "param": null, "code": "429"}}) + ); + assert_eq!(header(&refused, "retry-after"), Some("60")); + assert_eq!(header(&refused, "rate_limit_type"), Some("requests")); + assert_eq!(header(&refused, "x-litellm-priority"), Some("default")); + let refused = send(&app, "/v1/messages", "sk-master").await; + assert_eq!( + refused.body, + json!({"type": "error", "error": {"type": "rate_limit_error", "message": message}}) + ); + assert_eq!(header(&refused, "retry-after"), None); +} + +#[rstest] +#[tokio::test] +async fn tokens_count_after_the_response() { + // 2 x 5 tokens of capacity; each response uses 7. + let app = gateway("tpm: 5", "", None).await; + assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200); + settled().await; + assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200); + settled().await; + let refused = send(&app, "/v1/messages", "sk-master").await; + assert_eq!(refused.status, 429); + assert_eq!( + refused.body["error"]["message"], + "Model capacity reached for claude. Priority: None, Rate limit type: tokens, Model TPM: 10, Model RPM: not configured, Remaining: -4" + ); +} + +#[rstest] +#[tokio::test] +async fn busy_groups_hold_priorities_to_their_share() { + // 20 requests a window; `prod` gets a fifth once the group is busy, here from the start. + let app = gateway( + "rpm: 10", + " priority_reservation: {prod: 0.1}\n priority_reservation_settings: {saturation_threshold: 0.0}\n", + None, + ) + .await; + assert_eq!( + send(&app, "/v1/chat/completions", "sk-prod").await.status, + 200 + ); + let second = send(&app, "/v1/chat/completions", "sk-prod").await; + assert_eq!(second.status, 200); + assert_eq!(header(&second, "x-litellm-priority"), Some("prod")); + assert_eq!( + header(&second, "x-ratelimit-priority_model-remaining-requests"), + Some("0") + ); + let refused = send(&app, "/v1/chat/completions", "sk-prod").await; + assert_eq!(refused.status, 429); + assert_eq!( + refused.body["error"]["message"], + "Priority-based rate limit exceeded. Model: claude, Priority: prod, Rate limit type: requests, Model TPM: not configured, Model RPM: 20, Remaining: 0, Model saturation: 0.0%" + ); + assert_eq!(header(&refused, "x-litellm-saturation"), Some("0.00%")); + // Other callers share the default pool and the rest of the group. + assert_eq!( + send(&app, "/v1/chat/completions", "sk-master").await.status, + 200 + ); +} + +#[rstest] +fn the_hook_is_on_only_when_listed() { + let config = Config::from_yaml("model_list: []").unwrap(); + assert_eq!(CapacitySettings::from_config(&config), None); + let config = + Config::from_yaml("model_list: []\nlitellm_settings:\n callbacks: model_capacity\n") + .unwrap(); + assert_eq!( + CapacitySettings::from_config(&config), + Some(CapacitySettings::default()) + ); +} + +async fn redis(url: &str, namespace: &str) -> Arc { + let coordination = + Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap(); + let router: Arc = Arc::new(NoRouter( + litellm_gateway_inference::ModelRouter::from_config( + &Config::from_yaml("model_list: []").unwrap(), + ) + .unwrap(), + )); + coordination.start(router, None, tokio_util::sync::CancellationToken::new()); + for _ in 0..100 { + if coordination.connected() { + return coordination; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("Redis did not connect"); +} + +struct NoRouter(litellm_gateway_inference::ModelRouter); + +impl litellm_gateway::coordination::RouterHandle for NoRouter { + fn router(&self) -> &litellm_gateway_inference::ModelRouter { + &self.0 + } +} + +#[rstest] +#[tokio::test] +async fn replicas_share_the_group_counters() { + let Ok(url) = std::env::var("LITELLM_TEST_REDIS_URL") else { + return; + }; + let namespace = format!("model-capacity-{}", uuid::Uuid::new_v4()); + let limits = "rpm: 2, tpm: 1000"; + let first = gateway(limits, "", Some(redis(&url, &namespace).await)).await; + let second = gateway(limits, "", Some(redis(&url, &namespace).await)).await; + for app in [&first, &second, &first, &second] { + assert_eq!(send(app, "/v1/messages", "sk-master").await.status, 200); + } + settled().await; + assert_eq!(send(&second, "/v1/messages", "sk-master").await.status, 429); + + let mut connection = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let get = async |connection: &mut redis::aio::MultiplexedConnection, key: String| { + redis::cmd("GET") + .arg(key) + .query_async::>(connection) + .await + .unwrap() + }; + let group = format!("{namespace}:{{model_saturation_check:claude}}"); + // The refused request is not counted; four responses used 7 tokens each. + assert_eq!( + get(&mut connection, format!("{group}:requests")) + .await + .as_deref(), + Some("4") + ); + assert_eq!( + get(&mut connection, format!("{group}:tokens")) + .await + .as_deref(), + Some("28") + ); + let pool = format!("{namespace}:{{priority_model:claude:default_pool}}:tokens"); + assert_eq!(get(&mut connection, pool).await.as_deref(), Some("28")); +} diff --git a/litellm-rust/docs/customer-proxy/model-capacity.svg b/litellm-rust/docs/customer-proxy/model-capacity.svg new file mode 100644 index 00000000000..32a512c05bc --- /dev/null +++ b/litellm-rust/docs/customer-proxy/model-capacity.svg @@ -0,0 +1,52 @@ + + + + + + + Model capacity hook (dynamic_rate_limiter_v3): fleet-wide group capacity and priority shares + + + At startup + group capacity = sum of deployments' + rpm / tpm (deployment, litellm_params, + model_info) + priority_reservation shares: percent, + rpm or tpm; normalized when over 1 + on when callbacks list the hook + + + Per request (one Redis script, all or nothing) + priority: team metadata, else key metadata; pool <group>:<priority> or <group>:default_pool + saturation = max(requests / rpm, tokens / tpm), read and reused for saturation_check_cache_ttl + {model_saturation_check:<group>}: +1 request, tokens checked; always enforced + {priority_model:<pool>}: share x capacity, enforced at saturation >= threshold, + only counted below it + + + + Refusal: 429, the proxy's messages + "Model capacity reached for <group>. Priority: ..., Rate limit + type: ..., Model TPM: ..., Model RPM: ..., Remaining: ..." + "Priority-based rate limit exceeded. ... Model saturation: 62.5%" + chat: retry-after, rate_limit_type, x-litellm-priority, + x-litellm-saturation; messages: Anthropic envelope only + + + + Admitted: request runs + complete responses carry x-litellm-priority, + x-litellm-rate-limiter-version: v3 and + x-ratelimit-model_saturation_check-* / + x-ratelimit-priority_model-* + + + + Settle: a successful response's tokens (token_rate_limit_type: total, input or output) + added to the group's and the admitted pool's :tokens, with the window TTL + + + + Same keys and scripts as the Python proxy, under its namespace: Rust and Python replicas share one count + No Redis, or Redis unreachable: this replica's counters (fail_closed_rate_limit_enforcement: 503 instead) + From 96881793357a0b6fe8c4a13f6304592a6a4b93b9 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:46 -0700 Subject: [PATCH 07/10] feat(gateway): prometheus observer with the proxy's metric and label names (#47060) Adds the prometheus observer, turned on by listing prometheus in litellm_settings.callbacks, success_callback or failure_callback. Each finished request updates the proxy's request, spend, token, latency and deployment metrics under its names and labels, and /metrics serves them after the gateway's own. prometheus_metrics_config, the exclude lists and prometheus_latency_buckets are read as the proxy reads them. The metrics the gateway already served under the same names come from the observer alone when it is on. Outcomes now carry when the attempt was sent upstream, for the LLM API latency and time to first token. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) --- .../crates/gateway-hooks/src/facts.rs | 2 + .../crates/gateway-hooks/src/observer.rs | 1 + .../crates/gateway-hooks/src/pipeline.rs | 2 + .../crates/gateway-inference/src/hooks.rs | 15 +- litellm-rust/crates/gateway/src/lib.rs | 1 + litellm-rust/crates/gateway/src/main.rs | 7 +- litellm-rust/crates/gateway/src/metrics.rs | 94 +- litellm-rust/crates/gateway/src/prometheus.rs | 1246 +++++++++++++++++ litellm-rust/crates/gateway/src/spend.rs | 2 +- .../crates/gateway/tests/request_hooks.rs | 1 + .../docs/customer-proxy/prometheus.svg | 54 + 11 files changed, 1384 insertions(+), 41 deletions(-) create mode 100644 litellm-rust/crates/gateway/src/prometheus.rs create mode 100644 litellm-rust/docs/customer-proxy/prometheus.svg 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/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-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index 4893a91e835..f6f41210300 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -52,7 +52,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, } @@ -181,7 +182,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 @@ -194,7 +196,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) => { @@ -204,6 +209,7 @@ impl Tracking { usage: None, cost: None, time_to_first_byte: None, + upstream_started_at, ended_at: SystemTime::now(), interrupted: false, }; @@ -237,6 +243,7 @@ impl Tracking { route, admission: Some(admission), deployment, + upstream_started_at, status: response.status().as_u16(), meter, started, @@ -276,6 +283,7 @@ struct Finish { route: Route, admission: Option, deployment: Option, + upstream_started_at: Option, status: u16, meter: Meter, body: Option>, @@ -318,6 +326,7 @@ impl Drop for Finish { usage: self.usage(), cost: None, time_to_first_byte: self.first_byte, + upstream_started_at: self.upstream_started_at, ended_at: SystemTime::now(), interrupted: !self.ended, }; diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index c21352cde86..e7bc29372e4 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -10,6 +10,7 @@ mod key_status; mod keys; pub mod metrics; pub mod model_capacity; +pub mod prometheus; pub mod rate_limits; pub mod request_hooks; mod secrets; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index b242a3b88ab..bdf8ffef1da 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -140,6 +140,10 @@ async fn main() -> Result<(), Box> { 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()); + } 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()?; @@ -154,7 +158,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/prometheus.rs b/litellm-rust/crates/gateway/src/prometheus.rs new file mode 100644 index 00000000000..1a397351fdb --- /dev/null +++ b/litellm-rust/crates/gateway/src/prometheus.rs @@ -0,0 +1,1246 @@ +//! 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