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 + + +