mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
feat(gateway): hooks and observers around inference requests, with usage and cost (#47032)
* feat(gateway): hooks and observers around inference requests, with usage and cost Adds the litellm-gateway-hooks crate: request facts built only when something reads them, on-path hooks that declare what they read and what they may change and run in config order, off-path observers on bounded queues, a usage tap for SSE and JSON responses, and per-deployment prices from the model cost map overlaid by config prices. The messages and chat handlers run through it when anything is installed; with nothing installed requests take the old path. The key store now also returns the key alias, user, team alias and key and team metadata, and router attempts carry the deployment id. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * fix(model-catalog): read token counts the cost map writes as whole floats The published cost map now has prompt_cache_min_tokens: 4096.0, which the typed u64 fields rejected, so the gateway dropped the whole map at startup and priced every request at zero. Counts accept any whole non-negative number, as Python reads any JSON number, and a test keeps the checked-in map parsing. --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Nate Armstrong
GitHub
parent
ee03f153bd
commit
e2691a7031
34 changed files with 3442 additions and 52 deletions
22
litellm-rust/Cargo.lock
generated
22
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
@ -32,6 +32,7 @@ pub struct VerificationToken {
|
|||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Team {
|
||||
pub alias: Option<String>,
|
||||
pub blocked: bool,
|
||||
pub models: Vec<String>,
|
||||
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),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
pub profile: Arc<crate::keys::KeyProfile>,
|
||||
}
|
||||
|
||||
impl KeyGrant {
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub team_alias: Option<String>,
|
||||
pub metadata: Arc<Map<String, Value>>,
|
||||
pub team_metadata: Arc<Map<String, Value>>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
profile: Arc<KeyProfile>,
|
||||
},
|
||||
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<String>,
|
||||
pub profile: Arc<KeyProfile>,
|
||||
}
|
||||
|
||||
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 {
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
22
litellm-rust/crates/gateway-hooks/Cargo.toml
Normal file
22
litellm-rust/crates/gateway-hooks/Cargo.toml
Normal file
|
|
@ -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"] }
|
||||
293
litellm-rust/crates/gateway-hooks/src/facts.rs
Normal file
293
litellm-rust/crates/gateway-hooks/src/facts.rs
Normal file
|
|
@ -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<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub team_id: Option<String>,
|
||||
pub team_alias: Option<String>,
|
||||
pub org_id: Option<String>,
|
||||
pub metadata: Arc<Map<String, Value>>,
|
||||
pub team_metadata: Arc<Map<String, Value>>,
|
||||
}
|
||||
|
||||
/// 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<str>,
|
||||
pub route: Route,
|
||||
pub key: Option<Arc<KeyFacts>>,
|
||||
pub end_user: Option<Arc<str>>,
|
||||
/// The public model name the client asked for.
|
||||
pub requested_model: Arc<str>,
|
||||
pub client_headers: Arc<HeaderMap>,
|
||||
pub stream: bool,
|
||||
/// The output budget the request set, if any.
|
||||
pub max_tokens: Option<u64>,
|
||||
pub started_at: SystemTime,
|
||||
body: Bytes,
|
||||
input_tokens: OnceLock<u64>,
|
||||
}
|
||||
|
||||
impl RequestFacts {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
request_id: Arc<str>,
|
||||
route: Route,
|
||||
key: Option<Arc<KeyFacts>>,
|
||||
end_user: Option<Arc<str>>,
|
||||
requested_model: Arc<str>,
|
||||
client_headers: Arc<HeaderMap>,
|
||||
stream: bool,
|
||||
max_tokens: Option<u64>,
|
||||
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<str>,
|
||||
pub model_group: Arc<str>,
|
||||
pub model: String,
|
||||
pub provider: Option<String>,
|
||||
pub api_base: Option<String>,
|
||||
}
|
||||
|
||||
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<DeploymentFacts>,
|
||||
/// The HTTP status the client got.
|
||||
pub status: u16,
|
||||
pub usage: Option<Usage>,
|
||||
/// In dollars, from usage and the deployment's prices; `None` when either is missing.
|
||||
pub cost: Option<f64>,
|
||||
pub time_to_first_byte: Option<Duration>,
|
||||
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<RequestFacts>,
|
||||
pub outcome: Outcome,
|
||||
pub annotations: Annotations,
|
||||
/// Present only when an observer reads [`Facts::RESPONSE_BODY`].
|
||||
pub response_body: Option<Bytes>,
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
320
litellm-rust/crates/gateway-hooks/src/hook.rs
Normal file
320
litellm-rust/crates/gateway-hooks/src/hook.rs
Normal file
|
|
@ -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<Self, http::header::InvalidHeaderName> {
|
||||
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<Cow<'static, str>>,
|
||||
message: impl Into<String>,
|
||||
) -> 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<HeaderOp>,
|
||||
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<String>, 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<String, Value>) {
|
||||
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<Self>, 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<RequestFacts>,
|
||||
edit: &'a mut RequestEdit,
|
||||
) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, 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)));
|
||||
}
|
||||
}
|
||||
22
litellm-rust/crates/gateway-hooks/src/lib.rs
Normal file
22
litellm-rust/crates/gateway-hooks/src/lib.rs
Normal file
|
|
@ -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};
|
||||
211
litellm-rust/crates/gateway-hooks/src/observer.rs
Normal file
211
litellm-rust/crates/gateway-hooks/src/observer.rs
Normal file
|
|
@ -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<Arc<Record>>;
|
||||
|
||||
struct Lane {
|
||||
name: &'static str,
|
||||
sender: mpsc::Sender<Message>,
|
||||
dropped: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
/// Hands records to every observer.
|
||||
pub struct Observers {
|
||||
lanes: Vec<Lane>,
|
||||
tasks: Mutex<Vec<JoinHandle<()>>>,
|
||||
reads: Facts,
|
||||
}
|
||||
|
||||
impl Observers {
|
||||
/// Starts a task per observer. Must be called inside a tokio runtime.
|
||||
pub fn start(observers: Vec<Arc<dyn Observer>>, 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::<Message>(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<Record>) {
|
||||
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<Item = (&'static str, u64)> + '_ {
|
||||
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<Record> {
|
||||
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<Vec<u16>>,
|
||||
flushed: Mutex<bool>,
|
||||
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;
|
||||
}
|
||||
}
|
||||
320
litellm-rust/crates/gateway-hooks/src/pipeline.rs
Normal file
320
litellm-rust/crates/gateway-hooks/src/pipeline.rs
Normal file
|
|
@ -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<Arc<dyn Hook>>,
|
||||
observers: Observers,
|
||||
prices: Arc<Prices>,
|
||||
reads: Facts,
|
||||
}
|
||||
|
||||
impl Pipeline {
|
||||
/// `hooks` run in the order given, which is the config's.
|
||||
pub fn new(hooks: Vec<Arc<dyn Hook>>, observers: Observers, prices: Arc<Prices>) -> 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<RequestFacts>) -> Result<Admission, Refusal> {
|
||||
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<Bytes>) {
|
||||
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<RequestFacts>,
|
||||
edit: RequestEdit,
|
||||
settles: Vec<Box<dyn Settle>>,
|
||||
}
|
||||
|
||||
impl Admission {
|
||||
pub fn request(&self) -> &Arc<RequestFacts> {
|
||||
&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<RequestFacts> {
|
||||
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<Mutex<Vec<String>>>;
|
||||
|
||||
struct Step {
|
||||
name: &'static str,
|
||||
refuse: bool,
|
||||
log: Log,
|
||||
}
|
||||
|
||||
struct Settled(&'static str, Log);
|
||||
|
||||
impl Settle for Settled {
|
||||
fn settle(self: Box<Self>, 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<RequestFacts>,
|
||||
edit: &'a mut RequestEdit,
|
||||
) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, 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<dyn Settle>
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Costs(Mutex<Vec<(Option<f64>, Option<Value>)>>);
|
||||
|
||||
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<dyn Hook> {
|
||||
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"]);
|
||||
}
|
||||
}
|
||||
376
litellm-rust/crates/gateway-hooks/src/pricing.rs
Normal file
376
litellm-rust/crates/gateway-hooks/src/pricing.rs
Normal file
|
|
@ -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_<N>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<Arc<str>, 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<str>, fields: &Map<String, Value>) {
|
||||
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<f64> {
|
||||
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<String, Value>],
|
||||
) -> Map<String, Value> {
|
||||
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<u64>,
|
||||
tier: CostTier,
|
||||
}
|
||||
|
||||
fn price_key(name: &str) -> Option<PriceKey> {
|
||||
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<TierRates>, 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<TierRates>,
|
||||
}
|
||||
|
||||
struct ParsedRates {
|
||||
standard: Rates,
|
||||
tiers: Vec<TierRates>,
|
||||
thresholds: Vec<ParsedThreshold>,
|
||||
}
|
||||
|
||||
impl ParsedRates {
|
||||
fn parse(fields: &Map<String, Value>) -> 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<PricingPlan, litellm_cost::PricingError> {
|
||||
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"));
|
||||
}
|
||||
}
|
||||
452
litellm-rust/crates/gateway-hooks/src/usage.rs
Normal file
452
litellm-rust/crates/gateway-hooks/src/usage.rs
Normal file
|
|
@ -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<Self> {
|
||||
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<Self> {
|
||||
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<u64>,
|
||||
output_tokens: Option<u64>,
|
||||
cache_creation_input_tokens: Option<u64>,
|
||||
cache_read_input_tokens: Option<u64>,
|
||||
cache_creation: Option<CacheCreation>,
|
||||
server_tool_use: Option<ServerToolUse>,
|
||||
service_tier: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct CacheCreation {
|
||||
ephemeral_1h_input_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct ServerToolUse {
|
||||
web_search_requests: Option<u64>,
|
||||
}
|
||||
|
||||
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<OpenAiPromptDetails>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct OpenAiPromptDetails {
|
||||
cached_tokens: Option<u64>,
|
||||
cache_creation_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
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<u8>,
|
||||
state: Line,
|
||||
usage: Option<Usage>,
|
||||
}
|
||||
|
||||
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<Usage> {
|
||||
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::<AnthropicEvent>(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<AnthropicMessage>,
|
||||
usage: Option<AnthropicUsage>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct AnthropicMessage {
|
||||
usage: Option<AnthropicUsage>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct OpenAiChunk {
|
||||
usage: Option<OpenAiUsage>,
|
||||
}
|
||||
|
||||
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()
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
headers: HeaderMap,
|
||||
body: JsonBody,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
handle(&gateway, &identity, body, None).await
|
||||
handle(&gateway, &identity, &headers, body, None).await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_from_model_path(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
Path(model): Path<String>,
|
||||
headers: HeaderMap,
|
||||
body: JsonBody,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
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<String, Value>, model: Option<&str>) -> Map<String, Valu
|
|||
async fn handle(
|
||||
gateway: &Gateway,
|
||||
identity: &AuthenticatedRequest,
|
||||
headers: &HeaderMap,
|
||||
JsonBody { object, raw }: JsonBody,
|
||||
model: Option<&str>,
|
||||
) -> Result<Response, Error> {
|
||||
|
|
@ -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<String, Value>| {
|
||||
if let Some(tracking) = &tracking {
|
||||
tracking.edit_body(body);
|
||||
}
|
||||
};
|
||||
edit(&mut body);
|
||||
let reparse = || -> Result<Map<String, Value>, 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.
|
||||
|
|
|
|||
|
|
@ -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<Refusal> 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",
|
||||
|
|
|
|||
379
litellm-rust/crates/gateway-inference/src/hooks.rs
Normal file
379
litellm-rust/crates/gateway-inference/src/hooks.rs
Normal file
|
|
@ -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<Pipeline>,
|
||||
/// The hash of the master key, which the proxy records as the key of its requests.
|
||||
master_key_hash: Option<Arc<str>>,
|
||||
}
|
||||
|
||||
impl Hooks {
|
||||
pub fn new(pipeline: Arc<Pipeline>, master_key_hash: Option<String>) -> Self {
|
||||
Self {
|
||||
pipeline,
|
||||
master_key_hash: master_key_hash.map(Into::into),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn pipeline(&self) -> &Arc<Pipeline> {
|
||||
&self.pipeline
|
||||
}
|
||||
}
|
||||
|
||||
/// One request's admission, and the deployment its last attempt went to.
|
||||
pub(crate) struct Tracking {
|
||||
pipeline: Arc<Pipeline>,
|
||||
route: Route,
|
||||
admission: Admission,
|
||||
deployment: Mutex<Option<DeploymentFacts>>,
|
||||
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<String, Value>,
|
||||
name: &str,
|
||||
) -> Result<Option<Tracking>, 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<Arc<KeyFacts>> {
|
||||
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<String, Value>) -> 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<String, Value>) {
|
||||
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<Response, Error>) -> Result<Response, Error> {
|
||||
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<Bytes>),
|
||||
}
|
||||
|
||||
/// Records the request when the response body ends or is dropped.
|
||||
struct Finish {
|
||||
pipeline: Arc<Pipeline>,
|
||||
route: Route,
|
||||
admission: Option<Admission>,
|
||||
deployment: Option<DeploymentFacts>,
|
||||
status: u16,
|
||||
meter: Meter,
|
||||
body: Option<Vec<Bytes>>,
|
||||
started: Instant,
|
||||
first_byte: Option<Duration>,
|
||||
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<Usage> {
|
||||
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<Usage> {
|
||||
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<Bytes, axum::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
|
@ -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<Arc<dyn litellm_cache_response::ResponseCacheService>>,
|
||||
hooks: OnceLock<Hooks>,
|
||||
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(),
|
||||
|
|
|
|||
|
|
@ -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<String, Value>| {
|
||||
if let Some(tracking) = &tracking {
|
||||
tracking.edit_body(body);
|
||||
}
|
||||
};
|
||||
edit(&mut body);
|
||||
let reparse = || -> Result<Map<String, Value>, 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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
217
litellm-rust/crates/gateway/src/hooks.rs
Normal file
217
litellm-rust/crates/gateway/src/hooks.rs
Normal file
|
|
@ -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<Arc<dyn Hook>>,
|
||||
observers: Vec<Arc<dyn Observer>>,
|
||||
) -> 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<Option<String>, 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<Catalog> {
|
||||
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<Vec<u8>, Box<dyn std::error::Error>> {
|
||||
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<Item = (&'a String, V)>,
|
||||
) -> Map<String, Value> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
|
@ -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<f64>, spend: f64) -> Option<Team> {
|
||||
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
|
||||
}
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -108,6 +108,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
|
|||
)));
|
||||
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<dyn Error>> {
|
|||
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<dyn Error>> {
|
|||
)
|
||||
.await;
|
||||
shutdown.cancel();
|
||||
litellm_gateway::hooks::close(&inference).await;
|
||||
if let Some(mcp) = mcp {
|
||||
mcp.close().await;
|
||||
}
|
||||
|
|
|
|||
395
litellm-rust/crates/gateway/tests/hooks.rs
Normal file
395
litellm-rust/crates/gateway/tests/hooks.rs
Normal file
|
|
@ -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<Mutex<Vec<Value>>>;
|
||||
|
||||
/// 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<Body>| {
|
||||
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<KeyHash, KeyStatus>);
|
||||
|
||||
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<dyn KeyLookup> {
|
||||
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<Mutex<Vec<u16>>>);
|
||||
|
||||
impl Settle for Released {
|
||||
fn settle(self: Box<Self>, 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<Mutex<Vec<u16>>>);
|
||||
|
||||
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<RequestFacts>,
|
||||
edit: &'a mut RequestEdit,
|
||||
) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, 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<dyn Settle>))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
struct Seen {
|
||||
status: u16,
|
||||
key_alias: Option<String>,
|
||||
team_id: Option<String>,
|
||||
deployment: Option<String>,
|
||||
usage: Option<(u64, u64, u64)>,
|
||||
cost: Option<f64>,
|
||||
interrupted: bool,
|
||||
}
|
||||
|
||||
struct Collect(Mutex<Vec<Seen>>);
|
||||
|
||||
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<litellm_gateway_inference::Gateway>,
|
||||
bodies: Bodies,
|
||||
settled: Arc<Mutex<Vec<u16>>>,
|
||||
seen: Arc<Collect>,
|
||||
}
|
||||
|
||||
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]);
|
||||
}
|
||||
|
|
@ -191,6 +191,7 @@ fn on_team() -> KeyStatus {
|
|||
models,
|
||||
allowed_routes,
|
||||
team_id: Some("team-1".into()),
|
||||
profile: Default::default(),
|
||||
},
|
||||
status => status,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<dyn KeyLookup> {
|
|||
},
|
||||
allowed_routes: AllowedRoutes::default(),
|
||||
team_id: Some("team-1".into()),
|
||||
profile: Default::default(),
|
||||
},
|
||||
),
|
||||
(
|
||||
|
|
@ -86,6 +88,7 @@ fn keys() -> Arc<dyn KeyLookup> {
|
|||
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<dyn KeyLookup> {
|
|||
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.
|
||||
|
|
|
|||
|
|
@ -283,13 +283,25 @@ pub struct ModelInfo {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub litellm_provider: Option<String>,
|
||||
/// 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<u64>,
|
||||
/// 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<u64>,
|
||||
/// 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<u64>,
|
||||
/// 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<f64>,
|
||||
/// 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<u64>,
|
||||
/// 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<u64>,
|
||||
/// 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<f64>,
|
||||
/// 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<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub rules: Option<Vec<Value>>,
|
||||
|
|
@ -566,7 +590,11 @@ pub struct ModelInfo {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tiered_pricing: Option<Vec<TieredRate>>,
|
||||
/// 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<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub use_openai_responses_path: Option<bool>,
|
||||
|
|
@ -580,3 +608,42 @@ pub struct ModelInfo {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub web_search_billing_unit: Option<WebSearchBillingUnit>,
|
||||
}
|
||||
|
||||
/// 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<Option<u64>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let Some(number) = Option::<serde_json::Number>::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::<ModelInfo>(serde_json::json!({"rpm": 1.5})).is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
75
litellm-rust/docs/customer-proxy/hooks-pipeline.svg
Normal file
75
litellm-rust/docs/customer-proxy/hooks-pipeline.svg
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 640" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
|
||||
<defs>
|
||||
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
|
||||
<marker id="r" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#dc2626"/></marker>
|
||||
</defs>
|
||||
<rect width="1000" height="640" fill="#ffffff"/>
|
||||
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">Hooks and observers around one /v1/messages or chat request</text>
|
||||
|
||||
<!-- auth -->
|
||||
<rect x="24" y="60" width="210" height="104" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
|
||||
<text x="38" y="84" font-weight="600" fill="#1e3a8a">Auth</text>
|
||||
<text x="38" y="106" fill="#1e293b" font-size="13">key store row now also gives</text>
|
||||
<text x="38" y="124" fill="#1e293b" font-size="13">alias, user, team alias,</text>
|
||||
<text x="38" y="142" fill="#1e293b" font-size="13">key and team metadata</text>
|
||||
|
||||
<!-- facts -->
|
||||
<rect x="270" y="60" width="300" height="104" rx="8" fill="#f8fafc" stroke="#64748b"/>
|
||||
<text x="284" y="84" font-weight="600" fill="#0f172a">Request facts (built only if read)</text>
|
||||
<text x="284" y="106" fill="#334155" font-size="13">key, end user, client headers, model,</text>
|
||||
<text x="284" y="124" fill="#334155" font-size="13">stream, max_tokens, raw body (shared),</text>
|
||||
<text x="284" y="142" fill="#334155" font-size="13">input estimate (bytes / 4, made once)</text>
|
||||
<line x1="234" y1="112" x2="266" y2="112" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<!-- hooks -->
|
||||
<rect x="606" y="60" width="370" height="140" rx="8" fill="#ecfdf5" stroke="#10b981"/>
|
||||
<text x="620" y="84" font-weight="600" fill="#065f46">Hooks, in config order</text>
|
||||
<text x="620" y="106" fill="#1e293b" font-size="13">each declares reads (facts) and writes (targets):</text>
|
||||
<text x="620" y="124" fill="#1e293b" font-size="13">admission, upstream headers, upstream body,</text>
|
||||
<text x="620" y="142" fill="#1e293b" font-size="13">response headers, annotations</text>
|
||||
<text x="620" y="160" fill="#1e293b" font-size="13">same target twice: later hook wins</text>
|
||||
<text x="620" y="178" fill="#1e293b" font-size="13">undeclared write: dropped with a warning</text>
|
||||
<line x1="570" y1="112" x2="602" y2="112" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<!-- refusal -->
|
||||
<rect x="606" y="230" width="370" height="76" rx="8" fill="#fef2f2" stroke="#dc2626"/>
|
||||
<text x="620" y="254" font-weight="600" fill="#991b1b">Refusal</text>
|
||||
<text x="620" y="274" fill="#1e293b" font-size="13">status, proxy error body, hook headers (retry-after)</text>
|
||||
<text x="620" y="292" fill="#1e293b" font-size="13">earlier hooks settle; observers still get a record</text>
|
||||
<line x1="790" y1="200" x2="790" y2="226" stroke="#dc2626" stroke-width="1.5" marker-end="url(#r)"/>
|
||||
|
||||
<!-- routing -->
|
||||
<rect x="24" y="230" width="546" height="76" rx="8" fill="#f8fafc" stroke="#64748b"/>
|
||||
<text x="38" y="254" font-weight="600" fill="#0f172a">Routing and provider call (unchanged)</text>
|
||||
<text x="38" y="274" fill="#334155" font-size="13">body edits applied to every attempt; each attempt notes its deployment id</text>
|
||||
<text x="38" y="292" fill="#334155" font-size="13">(model_info.id, or the router's derived id)</text>
|
||||
<path d="M606 180 L300 180 L300 226" fill="none" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<!-- tap -->
|
||||
<rect x="24" y="340" width="546" height="122" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
|
||||
<text x="38" y="364" font-weight="600" fill="#92400e">Response tap</text>
|
||||
<text x="38" y="386" fill="#1e293b" font-size="13">stream: scans SSE lines; Anthropic decides from the first 96 bytes,</text>
|
||||
<text x="38" y="404" fill="#1e293b" font-size="13">so content lines are never copied; OpenAI keeps only usage chunks</text>
|
||||
<text x="38" y="422" fill="#1e293b" font-size="13">JSON: parses only the top-level "usage" object; Content-Length kept</text>
|
||||
<text x="38" y="440" fill="#1e293b" font-size="13">ends on last chunk or client drop (recorded as interrupted)</text>
|
||||
<line x1="297" y1="306" x2="297" y2="336" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<!-- cost -->
|
||||
<rect x="606" y="340" width="370" height="122" rx="8" fill="#f8fafc" stroke="#64748b"/>
|
||||
<text x="620" y="364" font-weight="600" fill="#0f172a">Cost (litellm-cost)</text>
|
||||
<text x="620" y="386" fill="#334155" font-size="13">prices per deployment id, built at startup:</text>
|
||||
<text x="620" y="404" fill="#334155" font-size="13">cost map entry overlaid by config prices</text>
|
||||
<text x="620" y="422" fill="#334155" font-size="13">_above_Nk_tokens thresholds, _priority/_flex,</text>
|
||||
<text x="620" y="440" fill="#334155" font-size="13">cache read / 5m / 1h write rates</text>
|
||||
<line x1="570" y1="400" x2="602" y2="400" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<!-- observers -->
|
||||
<rect x="24" y="496" width="952" height="116" rx="8" fill="#f5f3ff" stroke="#8b5cf6"/>
|
||||
<text x="38" y="520" font-weight="600" fill="#5b21b6">Settle, then observers (off the request path)</text>
|
||||
<text x="38" y="542" fill="#1e293b" font-size="13">hooks that asked are told how the request ended (to release reservations)</text>
|
||||
<text x="38" y="560" fill="#1e293b" font-size="13">one record (facts, deployment, status, usage, cost, timings, annotations) per request</text>
|
||||
<text x="38" y="578" fill="#1e293b" font-size="13">each observer: own tokio task behind a bounded queue; full queue drops and counts; shutdown drains then flushes</text>
|
||||
<text x="38" y="596" fill="#64748b" font-size="13">nothing configured: none of this runs and requests take the old path</text>
|
||||
<line x1="297" y1="462" x2="297" y2="492" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<line x1="790" y1="462" x2="790" y2="492" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 5.9 KiB |
Loading…
Add table
Reference in a new issue