diff --git a/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json b/litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json similarity index 55% rename from litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json rename to litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json index c9805198db5..de185711b56 100644 --- a/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json +++ b/litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n (SELECT d.max_budget FROM \"LiteLLM_BudgetTable\" AS d\n WHERE d.budget_id = $2) AS \"default_max_budget?\"\n FROM (SELECT 1) AS one\n LEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\n LEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = e.budget_id", + "query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n b.rpm_limit AS \"rpm_limit?\",\n b.tpm_limit AS \"tpm_limit?\"\n FROM (SELECT 1) AS one\n LEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\n LEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = COALESCE($3, e.budget_id, $2)", "describe": { "columns": [ { @@ -38,13 +38,30 @@ }, { "ordinal": 3, - "name": "default_max_budget?", - "type_info": "Float8", - "origin": "Expression" + "name": "rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 4, + "name": "tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "tpm_limit" + } + } } ], "parameters": { "Left": [ + "Text", "Text", "Text" ] @@ -53,8 +70,9 @@ true, true, true, - null + true, + true ] }, - "hash": "a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe" + "hash": "623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6" } diff --git a/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json b/litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json similarity index 64% rename from litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json rename to litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json index 1a9ad366cfa..d8f743ae315 100644 --- a/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json +++ b/litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n COALESCE(v.max_budget, kb.max_budget) AS max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL\n OR v.max_parallel_requests IS NOT NULL\n OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL\n OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL\n OR t.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL\n OR mb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n tm.spend AS \"member_spend?\",\n mb.max_budget AS \"member_max_budget?\",\n CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC')\n THEN mb.temp_budget_increase END AS \"member_budget_increase?\",\n tdb.max_budget AS \"member_default_budget?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n u.spend AS \"user_spend?\",\n u.max_budget AS \"user_max_budget?\",\n (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL\n OR u.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS tdb\n ON tdb.budget_id = t.metadata ->> 'team_member_budget_id'\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", + "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n COALESCE(v.max_budget, kb.max_budget) AS max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n COALESCE(v.rpm_limit, kb.rpm_limit) AS rpm_limit,\n COALESCE(v.tpm_limit, kb.tpm_limit) AS tpm_limit,\n v.max_parallel_requests,\n (COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n t.rpm_limit AS \"team_rpm_limit?\",\n t.tpm_limit AS \"team_tpm_limit?\",\n (COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n mb.rpm_limit AS \"member_rpm_limit?\",\n mb.tpm_limit AS \"member_tpm_limit?\",\n (COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n tm.spend AS \"member_spend?\",\n mb.max_budget AS \"member_max_budget?\",\n CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC')\n THEN mb.temp_budget_increase END AS \"member_budget_increase?\",\n tdb.max_budget AS \"member_default_budget?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n u.spend AS \"user_spend?\",\n u.max_budget AS \"user_max_budget?\",\n u.rpm_limit AS \"user_rpm_limit?\",\n u.tpm_limit AS \"user_tpm_limit?\",\n (COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS tdb\n ON tdb.budget_id = t.metadata ->> 'team_member_budget_id'\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", "describe": { "columns": [ { @@ -105,54 +105,77 @@ }, { "ordinal": 10, + "name": "rpm_limit", + "type_info": "Int8", + "origin": "Expression" + }, + { + "ordinal": 11, + "name": "tpm_limit", + "type_info": "Int8", + "origin": "Expression" + }, + { + "ordinal": 12, + "name": "max_parallel_requests", + "type_info": "Int4", + "origin": { + "Table": { + "table": "\"LiteLLM_VerificationToken\"", + "name": "max_parallel_requests" + } + } + }, + { + "ordinal": 13, "name": "key_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 11, + "ordinal": 14, "name": "key_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 12, + "ordinal": 15, "name": "key_model_aliases!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 13, + "ordinal": 16, "name": "key_object_permissions!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 14, + "ordinal": 17, "name": "key_cache_controls!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 15, + "ordinal": 18, "name": "key_project!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 16, + "ordinal": 19, "name": "key_organization!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 17, + "ordinal": 20, "name": "key_agent!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 18, + "ordinal": 21, "name": "team_id", "type_info": "Text", "origin": { @@ -163,7 +186,7 @@ } }, { - "ordinal": 19, + "ordinal": 22, "name": "team_row_id?", "type_info": "Text", "origin": { @@ -174,7 +197,7 @@ } }, { - "ordinal": 20, + "ordinal": 23, "name": "team_alias?", "type_info": "Text", "origin": { @@ -185,7 +208,7 @@ } }, { - "ordinal": 21, + "ordinal": 24, "name": "team_blocked?", "type_info": "Bool", "origin": { @@ -196,7 +219,7 @@ } }, { - "ordinal": 22, + "ordinal": 25, "name": "team_models?", "type_info": "TextArray", "origin": { @@ -207,7 +230,7 @@ } }, { - "ordinal": 23, + "ordinal": 26, "name": "team_spend?", "type_info": "Float8", "origin": { @@ -218,7 +241,7 @@ } }, { - "ordinal": 24, + "ordinal": 27, "name": "team_max_budget?", "type_info": "Float8", "origin": { @@ -229,7 +252,7 @@ } }, { - "ordinal": 25, + "ordinal": 28, "name": "team_metadata?", "type_info": "Jsonb", "origin": { @@ -240,43 +263,87 @@ } }, { - "ordinal": 26, + "ordinal": 29, + "name": "team_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 30, + "name": "team_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 31, "name": "team_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 27, + "ordinal": 32, "name": "team_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 28, + "ordinal": 33, "name": "team_model_aliases!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 29, + "ordinal": 34, "name": "team_object_permissions!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 30, + "ordinal": 35, "name": "team_organization!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 31, + "ordinal": 36, + "name": "member_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 37, + "name": "member_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 38, "name": "member_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 32, + "ordinal": 39, "name": "member_allowed_models?", "type_info": "TextArray", "origin": { @@ -287,7 +354,7 @@ } }, { - "ordinal": 33, + "ordinal": 40, "name": "member_spend?", "type_info": "Float8", "origin": { @@ -298,7 +365,7 @@ } }, { - "ordinal": 34, + "ordinal": 41, "name": "member_max_budget?", "type_info": "Float8", "origin": { @@ -309,13 +376,13 @@ } }, { - "ordinal": 35, + "ordinal": 42, "name": "member_budget_increase?", "type_info": "Float8", "origin": "Expression" }, { - "ordinal": 36, + "ordinal": 43, "name": "member_default_budget?", "type_info": "Float8", "origin": { @@ -326,7 +393,7 @@ } }, { - "ordinal": 37, + "ordinal": 44, "name": "user_id", "type_info": "Text", "origin": { @@ -337,7 +404,7 @@ } }, { - "ordinal": 38, + "ordinal": 45, "name": "user_row_id?", "type_info": "Text", "origin": { @@ -348,7 +415,7 @@ } }, { - "ordinal": 39, + "ordinal": 46, "name": "user_models?", "type_info": "TextArray", "origin": { @@ -359,7 +426,7 @@ } }, { - "ordinal": 40, + "ordinal": 47, "name": "user_spend?", "type_info": "Float8", "origin": { @@ -370,7 +437,7 @@ } }, { - "ordinal": 41, + "ordinal": 48, "name": "user_max_budget?", "type_info": "Float8", "origin": { @@ -381,19 +448,41 @@ } }, { - "ordinal": 42, + "ordinal": 49, + "name": "user_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 50, + "name": "user_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 51, "name": "user_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 43, + "ordinal": 52, "name": "user_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 44, + "ordinal": 53, "name": "user_object_permissions!", "type_info": "Bool", "origin": "Expression" @@ -417,6 +506,9 @@ false, null, null, + true, + null, + null, null, null, null, @@ -431,7 +523,8 @@ true, true, true, - null, + true, + true, null, null, null, @@ -439,6 +532,9 @@ null, true, true, + null, + true, + true, true, null, true, @@ -447,10 +543,12 @@ true, true, true, + true, + true, null, null, null ] }, - "hash": "08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9" + "hash": "77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773" } diff --git a/litellm-rust/crates/db/src/end_user.rs b/litellm-rust/crates/db/src/end_user.rs index a4d98eef2aa..4f7d05a32ec 100644 --- a/litellm-rust/crates/db/src/end_user.rs +++ b/litellm-rust/crates/db/src/end_user.rs @@ -1,6 +1,6 @@ use sqlx::PgPool; -use crate::Error; +use crate::{Error, RateLimits}; /// An end user's spend and budget, as the proxy's end-user budget checks read them. #[derive(Clone, Debug, Default, PartialEq)] @@ -8,33 +8,42 @@ pub struct EndUser { /// `LiteLLM_EndUserTable.spend`, when the end user has a row. pub spend: Option, pub blocked: bool, - /// The row's budget, else the default budget's. + /// The budget's `max_budget`: the override's, else the row's, else the default's. pub max_budget: Option, + /// That budget's `rpm_limit` and `tpm_limit`. + pub limits: RateLimits, } -/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own. +/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own +/// and `override_budget_id`, when given, as its budget whatever the row says. pub async fn find_end_user( pool: &PgPool, end_user_id: &str, default_budget_id: Option<&str>, + override_budget_id: Option<&str>, ) -> Result { let row = sqlx::query!( r#"SELECT e.spend AS "spend?", e.blocked AS "blocked?", b.max_budget AS "max_budget?", - (SELECT d.max_budget FROM "LiteLLM_BudgetTable" AS d - WHERE d.budget_id = $2) AS "default_max_budget?" + b.rpm_limit AS "rpm_limit?", + b.tpm_limit AS "tpm_limit?" FROM (SELECT 1) AS one LEFT JOIN "LiteLLM_EndUserTable" AS e ON e.user_id = $1 - LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = e.budget_id"#, + LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = COALESCE($3, e.budget_id, $2)"#, end_user_id, default_budget_id, + override_budget_id, ) .fetch_one(pool) .await?; Ok(EndUser { spend: row.spend, blocked: row.blocked.unwrap_or(false), - max_budget: row.max_budget.or(row.default_max_budget), + max_budget: row.max_budget, + limits: RateLimits { + rpm: row.rpm_limit, + tpm: row.tpm_limit, + }, }) } diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index b41ee5ace78..20d7d7616bb 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -16,5 +16,6 @@ pub use spend::{ }; pub use sqlx::PgPool; pub use verification_token::{ - Restrictions, Team, TeamMembership, User, VerificationToken, find_verification_token, + RateLimits, Restrictions, Team, TeamMembership, User, VerificationToken, + find_verification_token, }; diff --git a/litellm-rust/crates/db/src/verification_token.rs b/litellm-rust/crates/db/src/verification_token.rs index 3ff008b0cff..5db37ae5c9b 100644 --- a/litellm-rust/crates/db/src/verification_token.rs +++ b/litellm-rust/crates/db/src/verification_token.rs @@ -19,6 +19,9 @@ pub struct VerificationToken { pub key_alias: Option, pub key_name: Option, pub metadata: Value, + /// The key's `rpm_limit` and `tpm_limit`, else its budget table's. + pub limits: RateLimits, + pub max_parallel_requests: Option, pub restrictions: Restrictions, pub team_id: Option, /// `None` when the key has no team or the team row is missing. @@ -38,9 +41,17 @@ pub struct Team { pub spend: f64, pub max_budget: Option, pub metadata: Value, + pub limits: RateLimits, pub restrictions: Restrictions, } +/// Requests and tokens per rate limit window. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct RateLimits { + pub rpm: Option, + pub tpm: Option, +} + /// The key's team membership: its spend and budget, and what its budget table sets beyond them. #[derive(Clone, Debug, Default, PartialEq)] pub struct TeamMembership { @@ -50,7 +61,9 @@ pub struct TeamMembership { pub max_budget: Option, /// The team's default member budget, `metadata.team_member_budget_id`. pub default_max_budget: Option, - /// `rpm_limit`, `tpm_limit`, `tpd_limit` or `max_parallel_requests`. + /// The membership budget table's `rpm_limit` and `tpm_limit`. + pub limits: RateLimits, + /// An `rpm_limit` or `tpm_limit` inside the budget table's `model_max_budget`. pub rate_limits: bool, /// The member's own model scope; empty inherits the team's. pub allowed_models: Vec, @@ -61,6 +74,7 @@ pub struct User { pub spend: f64, pub max_budget: Option, pub models: Vec, + pub limits: RateLimits, pub restrictions: Restrictions, } @@ -68,8 +82,8 @@ pub struct User { /// `allowed_routes`. #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct Restrictions { - /// `rpm_limit`, `tpm_limit`, `tpd_limit` or `max_parallel_requests`, on the row or on a key's - /// budget table, or an `rpm_limit` or `tpm_limit` inside `model_max_budget`. + /// Per-model limits: an `rpm_limit` or `tpm_limit` inside `model_max_budget`, on the row or + /// on a key's budget table. pub rate_limits: bool, pub policies: bool, /// `aliases`, `router_settings`, or a team's `model_id`. @@ -98,11 +112,10 @@ pub async fn find_verification_token( v.key_alias, v.key_name, v.metadata, - (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL - OR v.max_parallel_requests IS NOT NULL - OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL - OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(v.model_max_budget, + COALESCE(v.rpm_limit, kb.rpm_limit) AS rpm_limit, + COALESCE(v.tpm_limit, kb.tpm_limit) AS tpm_limit, + v.max_parallel_requests, + (COALESCE(jsonb_path_exists(v.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false) OR COALESCE(jsonb_path_exists(kb.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) @@ -123,9 +136,9 @@ pub async fn find_verification_token( t.spend AS "team_spend?", t.max_budget AS "team_max_budget?", t.metadata AS "team_metadata?", - (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL - OR t.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(t.model_max_budget, + t.rpm_limit AS "team_rpm_limit?", + t.tpm_limit AS "team_tpm_limit?", + (COALESCE(jsonb_path_exists(t.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "team_rate_limits!", COALESCE(cardinality(t.policies) > 0, false) AS "team_policies!", @@ -133,9 +146,9 @@ pub async fn find_verification_token( OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS "team_model_aliases!", (t.object_permission_id IS NOT NULL) AS "team_object_permissions!", (t.organization_id IS NOT NULL) AS "team_organization!", - (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL - OR mb.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(mb.model_max_budget, + mb.rpm_limit AS "member_rpm_limit?", + mb.tpm_limit AS "member_tpm_limit?", + (COALESCE(jsonb_path_exists(mb.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "member_rate_limits!", mb.allowed_models AS "member_allowed_models?", @@ -149,9 +162,9 @@ pub async fn find_verification_token( u.models AS "user_models?", u.spend AS "user_spend?", u.max_budget AS "user_max_budget?", - (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL - OR u.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(u.model_max_budget, + u.rpm_limit AS "user_rpm_limit?", + u.tpm_limit AS "user_tpm_limit?", + (COALESCE(jsonb_path_exists(u.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "user_rate_limits!", COALESCE(cardinality(u.policies) > 0, false) AS "user_policies!", @@ -181,6 +194,11 @@ pub async fn find_verification_token( key_alias: row.key_alias, key_name: row.key_name, metadata: row.metadata, + limits: RateLimits { + rpm: row.rpm_limit, + tpm: row.tpm_limit, + }, + max_parallel_requests: row.max_parallel_requests.map(i64::from), restrictions: Restrictions { rate_limits: row.key_rate_limits, policies: row.key_policies, @@ -199,6 +217,10 @@ pub async fn find_verification_token( spend: row.team_spend.unwrap_or(0.0), max_budget: row.team_max_budget, metadata: row.team_metadata.unwrap_or(Value::Null), + limits: RateLimits { + rpm: row.team_rpm_limit, + tpm: row.team_tpm_limit, + }, restrictions: Restrictions { rate_limits: row.team_rate_limits, policies: row.team_policies, @@ -214,6 +236,10 @@ pub async fn find_verification_token( .member_max_budget .map(|budget| budget + row.member_budget_increase.unwrap_or(0.0)), default_max_budget: row.member_default_budget, + limits: RateLimits { + rpm: row.member_rpm_limit, + tpm: row.member_tpm_limit, + }, rate_limits: row.member_rate_limits, allowed_models: row.member_allowed_models.unwrap_or_default(), }, @@ -222,6 +248,10 @@ pub async fn find_verification_token( spend: row.user_spend.unwrap_or(0.0), max_budget: row.user_max_budget, models: row.user_models.unwrap_or_default(), + limits: RateLimits { + rpm: row.user_rpm_limit, + tpm: row.user_tpm_limit, + }, restrictions: Restrictions { rate_limits: row.user_rate_limits, policies: row.user_policies, diff --git a/litellm-rust/crates/db/tests/verification_token.rs b/litellm-rust/crates/db/tests/verification_token.rs index ffe0b4d66f6..fc4cb83e04f 100644 --- a/litellm-rust/crates/db/tests/verification_token.rs +++ b/litellm-rust/crates/db/tests/verification_token.rs @@ -1,8 +1,8 @@ use std::time::{Duration, SystemTime}; use litellm_db::{ - EndUser, Restrictions, Team, TeamMembership, User, VerificationToken, find_end_user, - find_verification_token, + EndUser, RateLimits, Restrictions, Team, TeamMembership, User, VerificationToken, + find_end_user, find_verification_token, }; use litellm_db_testing::MigratedPostgres; use rstest::rstest; @@ -93,6 +93,8 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { key_alias: None, key_name: None, metadata: json!({}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: None, team: None, @@ -115,6 +117,8 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { key_alias: Some("prod".into()), key_name: Some("sk-...abcd".into()), metadata: json!({"tags": ["a"]}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: Some("team-1".into()), team: Some(Team { @@ -124,16 +128,21 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: 2.5, max_budget: Some(10.0), metadata: json!({"require_trace_id": true}), - restrictions: Restrictions { - rate_limits: true, - ..Restrictions::default() + limits: RateLimits { + rpm: None, + tpm: Some(100), }, + restrictions: Restrictions::default(), }), team_membership: TeamMembership { spend: Some(0.0), max_budget: None, default_max_budget: None, - rate_limits: true, + limits: RateLimits { + rpm: Some(3), + tpm: None, + }, + rate_limits: false, allowed_models: vec![], }, user_id: Some("user-1".into()), @@ -141,8 +150,11 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: 0.0, max_budget: Some(1.0), models: vec!["gpt-b".into()], + limits: RateLimits { + rpm: Some(5), + tpm: None, + }, restrictions: Restrictions { - rate_limits: true, policies: true, ..Restrictions::default() }, @@ -152,7 +164,7 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { assert_eq!( find("restricted").await.restrictions, Restrictions { - rate_limits: true, + rate_limits: false, policies: true, model_aliases: true, object_permissions: true, @@ -162,16 +174,33 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { agent: true, } ); + let limited = find("limited-only").await; assert_eq!( - find("limited-only").await.restrictions, - Restrictions { - rate_limits: true, - ..Restrictions::default() + (limited.max_parallel_requests, limited.restrictions), + (Some(2), Restrictions::default()) + ); + // A key's own limits win; its budget row's fill the gaps. + assert_eq!( + find("restricted").await.limits, + RateLimits { + rpm: Some(10), + tpm: None } ); - // Budgets are not restrictions; limits held in a budget row or per model are. + let budget_row = find("budget-row-limits").await; + assert_eq!( + (budget_row.limits, budget_row.restrictions), + ( + RateLimits { + rpm: Some(3), + tpm: None + }, + Restrictions::default() + ) + ); + // Budgets and plain limits are not restrictions; limits per model are. assert_eq!(find("budgeted").await.restrictions, Restrictions::default()); - for token in ["budget-row-limits", "budget-row-model-limits", "model-rpm"] { + for token in ["budget-row-model-limits", "model-rpm"] { assert_eq!( find(token).await.restrictions, Restrictions { @@ -188,6 +217,7 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: Some(0.75), max_budget: Some(5.0), default_max_budget: None, + limits: RateLimits::default(), rate_limits: false, allowed_models: vec!["gpt-a".into()], } @@ -216,15 +246,17 @@ async fn reads_end_user_budgets_with_their_default() { let postgres = MigratedPostgres::start().await.unwrap(); let pool = postgres.pool(); setup(pool).await; - let find = - async |id: &str, default: Option<&str>| find_end_user(pool, id, default).await.unwrap(); + let find = async |id: &str, default: Option<&str>| { + find_end_user(pool, id, default, None).await.unwrap() + }; assert_eq!( find("end-1", Some("b-2")).await, EndUser { spend: Some(1.5), blocked: false, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); assert_eq!( @@ -232,16 +264,37 @@ async fn reads_end_user_budgets_with_their_default() { EndUser { spend: Some(0.25), blocked: true, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); assert_eq!(find("end-2", None).await.max_budget, None); + // An override replaces the row's own budget. + assert_eq!( + find_end_user(pool, "end-1", None, Some("b-2")) + .await + .unwrap() + .limits, + RateLimits { + rpm: Some(3), + tpm: None + } + ); + // A row without its own budget takes the default's limits too. + assert_eq!( + find("end-2", Some("b-2")).await.limits, + RateLimits { + rpm: Some(3), + tpm: None + } + ); assert_eq!( find("unknown", Some("b-1")).await, EndUser { spend: None, blocked: false, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); } diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs index 6310e2eda90..30494a7e1a3 100644 --- a/litellm-rust/crates/gateway-auth/src/error.rs +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -269,7 +269,7 @@ mod tests { restriction: Restriction::RateLimits, owner: Owner::TeamMembership, }), - "Authentication Error, keys with rate limits set on its team membership are not supported by this gateway yet." + "Authentication Error, keys with per-model, per-tag or dynamic rate limits set on its team membership are not supported by this gateway yet." )] #[case::banned_param( Error::BannedParam("aws_bedrock_runtime_endpoint"), diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index ff99572a1ee..654b8503d3e 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -241,7 +241,7 @@ pub enum Refusal { #[derive(Clone, Copy, Debug, Eq, PartialEq, strum::Display, strum::VariantArray)] pub enum Restriction { - #[strum(serialize = "rate limits")] + #[strum(serialize = "per-model, per-tag or dynamic rate limits")] RateLimits, #[strum(serialize = "guardrails or policies")] Guardrails, @@ -298,6 +298,34 @@ pub struct KeyProfile { pub metadata: Arc>, pub team_metadata: Arc>, pub budgets: KeyBudgets, + pub limits: KeyLimits, +} + +/// Requests and tokens per rate limit window. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Limits { + pub rpm: Option, + pub tpm: Option, +} + +impl Limits { + pub fn is_set(&self) -> bool { + self.rpm.is_some() || self.tpm.is_some() + } +} + +/// The rate limits of a key and of who it belongs to, read with the key and refreshed with it. +/// End users' limits are read on their own. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct KeyLimits { + /// The key's own limits, else its budget table's. + pub key: Limits, + /// Requests the key may have in flight at once. + pub max_parallel_requests: Option, + pub user: Limits, + pub team: Limits, + /// The key's user within its team, from the membership's budget table. + pub team_member: Limits, } /// Stored spend against an optional budget, as the key store last read it. diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs index 437768707d0..d696427fd8d 100644 --- a/litellm-rust/crates/gateway-hooks/src/facts.rs +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -8,7 +8,7 @@ use bytes::Bytes; use http::HeaderMap; use serde_json::{Map, Value}; -use litellm_gateway_auth::keys::KeyBudgets; +use litellm_gateway_auth::keys::{KeyBudgets, KeyLimits}; use crate::usage::Usage; @@ -104,6 +104,7 @@ pub struct KeyFacts { pub metadata: Arc>, pub team_metadata: Arc>, pub budgets: KeyBudgets, + pub limits: KeyLimits, } /// What the gateway knows about a request before it is sent, shared by every hook and observer. diff --git a/litellm-rust/crates/gateway-hooks/src/hook.rs b/litellm-rust/crates/gateway-hooks/src/hook.rs index 3e57e6b12c8..9246f77d1a2 100644 --- a/litellm-rust/crates/gateway-hooks/src/hook.rs +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -100,6 +100,8 @@ pub struct Refusal { pub param: Option<&'static str>, pub message: String, pub headers: Vec<(HeaderName, HeaderValue)>, + /// Send Anthropic's error envelope, with `kind` as its error type, instead of the proxy's. + pub anthropic: bool, } impl Refusal { @@ -114,9 +116,17 @@ impl Refusal { param: Some("None"), message: message.into(), headers: Vec::new(), + anthropic: false, } } + /// Sends `{"type": "error", "error": {"type": kind, "message": message}}`, as the proxy's + /// `/v1/messages` route does for some errors. + pub fn anthropic(mut self) -> Self { + self.anthropic = true; + self + } + /// Sends a `null` `param`, as the proxy's budget errors do. pub fn without_param(mut self) -> Self { self.param = None; diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 05009604fba..3aa4b44f9aa 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -177,12 +177,18 @@ impl Error { /// A hook's refusal, in the proxy's error shape and with the hook's headers. fn refused(refusal: Refusal) -> Response { let status = StatusCode::from_u16(refusal.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); - let body = json!({"error": { - "message": refusal.message, - "type": refusal.kind, - "param": refusal.param, - "code": status.as_u16().to_string(), - }}); + let body = match refusal.anthropic { + true => json!({"type": "error", "error": { + "type": refusal.kind, + "message": refusal.message, + }}), + false => json!({"error": { + "message": refusal.message, + "type": refusal.kind, + "param": refusal.param, + "code": status.as_u16().to_string(), + }}), + }; let mut response = (status, Json(body)).into_response(); response.headers_mut().extend(refusal.headers); response diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index e60bcd140bc..07833d8cc93 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -121,6 +121,7 @@ fn key_facts(hooks: &Hooks, identity: &AuthenticatedRequest) -> Option Some(Arc::new(KeyFacts { diff --git a/litellm-rust/crates/gateway/src/budgets.rs b/litellm-rust/crates/gateway/src/budgets.rs index be04ae028be..7196bf132da 100644 --- a/litellm-rust/crates/gateway/src/budgets.rs +++ b/litellm-rust/crates/gateway/src/budgets.rs @@ -17,22 +17,18 @@ use std::{ use futures_util::future::BoxFuture; use litellm_config::Config; -use litellm_db::{EndUser, PgPool}; +use litellm_db::PgPool; use litellm_gateway_hooks::{ Facts, Hook, Prices, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets, Usage, }; use redis::aio::ConnectionManager; -use crate::coordination::Coordination; +use crate::{coordination::Coordination, end_users::EndUsers}; /// The output tokens the proxy assumes when neither the request nor the model caps them. const DEFAULT_OUTPUT_TOKENS: u64 = 16_384; /// The proxy's slack when a reservation exactly meets a budget. const EPSILON: f64 = 1e-12; -/// How long an end user's budget is reused before it is read again. -const END_USER_TTL: Duration = Duration::from_secs(60); -/// End users held in memory; the cache starts over past this. -const MAX_END_USERS: usize = 10_000; const FAIL_CLOSED_READ: &str = "{'error': 'Budget enforcement unavailable: current spend could not be verified against Redis or the database, and fail_closed_budget_enforcement is enabled, so the request was rejected to avoid exceeding the configured budget. Retry shortly.'}"; /// Seeds or raises each counter to its floor, runs the budget checks in the order given and @@ -600,45 +596,6 @@ async fn admit_redis( } } -/// An end user and the default budget it was read with. -type EndUserKey = (String, Option); - -/// End users' budgets, read from the proxy's database and reused for a minute. -struct EndUsers { - pool: PgPool, - cache: Mutex>, -} - -impl EndUsers { - async fn get(&self, id: &str, default_budget_id: Option<&str>) -> Option { - let key = (id.to_owned(), default_budget_id.map(str::to_owned)); - let now = Instant::now(); - if let Some((end_user, read_at)) = self - .cache - .lock() - .unwrap_or_else(|e| e.into_inner()) - .get(&key) - && now.duration_since(*read_at) < END_USER_TTL - { - return Some(end_user.clone()); - } - match litellm_db::find_end_user(&self.pool, id, default_budget_id).await { - Ok(end_user) => { - let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); - if cache.len() >= MAX_END_USERS { - cache.clear(); - } - cache.insert(key, (end_user.clone(), now)); - Some(end_user) - } - Err(error) => { - tracing::warn!(%error, "reading an end user's budget failed; not checking it"); - None - } - } - } -} - /// The `budgets` hook. pub struct Budgets { settings: BudgetSettings, @@ -670,10 +627,7 @@ impl Budgets { fail_closed: settings.fail_closed, }), settings, - end_users: database.map(|pool| EndUsers { - pool, - cache: Mutex::default(), - }), + end_users: database.map(EndUsers::new), groups, prices: OnceLock::new(), } @@ -817,7 +771,7 @@ impl Budgets { .as_deref() .or(self.settings.end_user_budget_id.as_deref()); let end_user = match &self.end_users { - Some(end_users) => end_users.get(end_user_id, default_budget_id).await, + Some(end_users) => end_users.get(end_user_id, default_budget_id, None).await, None => None, } .unwrap_or_default(); diff --git a/litellm-rust/crates/gateway/src/end_users.rs b/litellm-rust/crates/gateway/src/end_users.rs new file mode 100644 index 00000000000..5c23e4f6328 --- /dev/null +++ b/litellm-rust/crates/gateway/src/end_users.rs @@ -0,0 +1,70 @@ +//! End users' budgets and limits, read from the proxy's database and reused for a minute. + +use std::{ + collections::HashMap, + sync::Mutex, + time::{Duration, Instant}, +}; + +use litellm_db::{EndUser, PgPool}; + +/// How long an end user's budget is reused before it is read again. +const END_USER_TTL: Duration = Duration::from_secs(60); +/// End users held in memory; the cache starts over past this. +const MAX_END_USERS: usize = 10_000; + +/// An end user and the default and override budgets it was read with. +type EndUserKey = (String, Option, Option); + +pub(crate) struct EndUsers { + pool: PgPool, + cache: Mutex>, +} + +impl EndUsers { + pub(crate) fn new(pool: PgPool) -> Self { + Self { + pool, + cache: Mutex::default(), + } + } + + /// `None` when the database could not be read. + pub(crate) async fn get( + &self, + id: &str, + default_budget_id: Option<&str>, + override_budget_id: Option<&str>, + ) -> Option { + let key = ( + id.to_owned(), + default_budget_id.map(str::to_owned), + override_budget_id.map(str::to_owned), + ); + let now = Instant::now(); + if let Some((end_user, read_at)) = self + .cache + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&key) + && now.duration_since(*read_at) < END_USER_TTL + { + return Some(end_user.clone()); + } + match litellm_db::find_end_user(&self.pool, id, default_budget_id, override_budget_id).await + { + Ok(end_user) => { + let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); + if cache.len() >= MAX_END_USERS { + cache.clear(); + } + cache.insert(key, (end_user.clone(), now)); + Some(end_user) + } + Err(error) => { + tracing::warn!(%error, "reading an end user's budget failed; not checking it"); + None + } + } + } +} diff --git a/litellm-rust/crates/gateway/src/key_status.rs b/litellm-rust/crates/gateway/src/key_status.rs index 50d8bf65567..20e2e795336 100644 --- a/litellm-rust/crates/gateway/src/key_status.rs +++ b/litellm-rust/crates/gateway/src/key_status.rs @@ -1,20 +1,21 @@ -use litellm_db::{Restrictions, VerificationToken}; +use litellm_db::{RateLimits, Restrictions, VerificationToken}; use std::sync::Arc; use litellm_gateway_auth::keys::{ - AllowedRoutes, KeyBudgets, KeyModels, KeyProfile, KeyStatus, MemberSpend, ModelAccess, Owner, - Refusal, Restriction, Spend, + AllowedRoutes, KeyBudgets, KeyLimits, KeyModels, KeyProfile, KeyStatus, Limits, MemberSpend, + ModelAccess, Owner, Refusal, Restriction, Spend, }; use serde_json::Value; /// Key and team `metadata` fields the proxy enforces and this gateway does not. -const METADATA_RESTRICTIONS: [(&str, Restriction); 11] = [ +const METADATA_RESTRICTIONS: [(&str, Restriction); 12] = [ ("guardrails", Restriction::Guardrails), ("policies", Restriction::Guardrails), ("tags", Restriction::Tags), ("model_rpm_limit", Restriction::RateLimits), ("model_tpm_limit", Restriction::RateLimits), ("tag_rpm_limit", Restriction::RateLimits), + ("throttle_on_budget_exceeded", Restriction::RateLimits), ("allowed_tools", Restriction::ToolAllowlists), ("enforced_params", Restriction::EnforcedParams), ("priority", Restriction::Priority), @@ -39,6 +40,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K _ => Arc::default(), }; let budgets = budgets(&row); + let limits = limits(&row); let (team_models, team_alias, team_metadata) = match row.team { Some(team) => ( Some(ModelAccess::new(team.models)), @@ -62,6 +64,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K metadata: object(row.metadata), team_metadata, budgets, + limits, }), } } @@ -96,6 +99,29 @@ fn budgets(row: &VerificationToken) -> KeyBudgets { } } +/// The rate limits the rate limits hook enforces for this key. +fn limits(row: &VerificationToken) -> KeyLimits { + let limits = |limits: RateLimits| Limits { + rpm: limits.rpm, + tpm: limits.tpm, + }; + KeyLimits { + key: limits(row.limits), + max_parallel_requests: row.max_parallel_requests, + user: row + .user + .as_ref() + .map(|user| limits(user.limits)) + .unwrap_or_default(), + team: row + .team + .as_ref() + .map(|team| limits(team.limits)) + .unwrap_or_default(), + team_member: limits(row.team_membership.limits), + } +} + fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option { let team = match (&row.team_id, &row.team) { (Some(team_id), None) => { @@ -145,7 +171,7 @@ fn unsupported( }) } -/// What the membership's budget table sets: rate limits, or an `allowed_models` list the proxy's +/// What the membership's budget table sets: per-model rate limits, or an `allowed_models` list the proxy's /// `check_team_member_model_access` can deny with (it skips an empty list, and `*` and /// `all-proxy-models` grant every model). fn membership(row: &VerificationToken) -> Option { @@ -195,6 +221,14 @@ fn metadata_restriction(metadata: &Value) -> Option { .iter() .find(|(field, _)| metadata.get(field).is_some_and(is_set)) .map(|(_, restriction)| *restriction) + .or_else(|| { + // Dynamic limits apply only while a deployment is failing, which the gateway does + // not track. + ["rpm_limit_type", "tpm_limit_type"] + .iter() + .any(|field| metadata.get(field).and_then(Value::as_str) == Some("dynamic")) + .then_some(Restriction::RateLimits) + }) } fn is_set(value: &Value) -> bool { @@ -236,6 +270,8 @@ mod tests { key_alias: None, key_name: None, metadata: json!({}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: None, team: None, @@ -253,6 +289,7 @@ mod tests { spend, max_budget, metadata: json!({}), + limits: RateLimits::default(), restrictions: Restrictions::default(), }) } @@ -297,6 +334,46 @@ mod tests { assert_eq!(status(row, false), active(true)); } + #[rstest] + fn carries_the_key_and_owner_rate_limits(row: VerificationToken) { + let limits = |rpm, tpm| RateLimits { rpm, tpm }; + let row = VerificationToken { + limits: limits(Some(10), None), + max_parallel_requests: Some(2), + team_id: Some("team-1".into()), + team: team(false, None, 0.0).map(|team| Team { + limits: limits(None, Some(1_000)), + ..team + }), + team_membership: TeamMembership { + limits: limits(Some(3), Some(30)), + ..TeamMembership::default() + }, + ..on_user(Some("team-1"), &[], NONE, row) + }; + let row = VerificationToken { + user: row.user.clone().map(|user| User { + limits: limits(Some(5), None), + ..user + }), + ..row + }; + let KeyStatus::Active { profile, .. } = status(row, false) else { + panic!("expected an active key"); + }; + let limits = |rpm, tpm| Limits { rpm, tpm }; + assert_eq!( + profile.limits, + KeyLimits { + key: limits(Some(10), None), + max_parallel_requests: Some(2), + user: limits(Some(5), None), + team: limits(None, Some(1_000)), + team_member: limits(Some(3), Some(30)), + } + ); + } + #[rstest] #[case::blocked_key(true, None, None, None, 0.0, KeyStatus::Blocked)] #[case::team_row_missing( @@ -376,6 +453,7 @@ mod tests { spend: 0.0, max_budget: None, models: strings(models), + limits: RateLimits::default(), restrictions, }), ..row @@ -473,6 +551,18 @@ mod tests { |row| key_metadata(json!({"tag_rpm_limit": {"a": 1}}), row), Some((Restriction::RateLimits, Owner::Key)) )] + #[case::metadata_dynamic_rpm( + |row| key_metadata(json!({"rpm_limit_type": "dynamic"}), row), + Some((Restriction::RateLimits, Owner::Key)) + )] + #[case::metadata_guaranteed_tpm( + |row| key_metadata(json!({"tpm_limit_type": "guaranteed_throughput"}), row), + None + )] + #[case::metadata_budget_throttle( + |row| key_metadata(json!({"throttle_on_budget_exceeded": true}), row), + Some((Restriction::RateLimits, Owner::Key)) + )] #[case::metadata_allowed_tools( |row| key_metadata(json!({"allowed_tools": ["search"]}), row), Some((Restriction::ToolAllowlists, Owner::Key)) diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index dcd3aa994a3..57f05ec9322 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -2,12 +2,14 @@ mod auth; pub mod budgets; pub mod coordination; pub mod drain; +mod end_users; mod error; pub mod health; pub mod hooks; mod key_status; mod keys; pub mod metrics; +pub mod rate_limits; mod secrets; pub mod spend; pub use error::Error; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 0954fd7634b..3e648f5a758 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -120,6 +120,13 @@ async fn main() -> Result<(), Box> { coordination.clone(), database.clone(), ))); + hooks.push(Arc::new( + litellm_gateway::rate_limits::RateLimiter::for_gateway( + &config, + coordination.clone(), + database.clone(), + ), + )); } let mut observers: Vec> = Vec::new(); let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config)); diff --git a/litellm-rust/crates/gateway/src/rate_limits.rs b/litellm-rust/crates/gateway/src/rate_limits.rs new file mode 100644 index 00000000000..2b9ca57cbbc --- /dev/null +++ b/litellm-rust/crates/gateway/src/rate_limits.rs @@ -0,0 +1,1193 @@ +//! The rate limits hook: refuses a request when its key, user, team, team member or end user is +//! over its requests or tokens per window, or its key has too many requests in flight, with the +//! proxy's counters, order, messages and headers. +//! +//! The counters are the proxy's own Redis keys, `{:}:window|requests|tokens` and +//! `{api_key:}:max_parallel_requests`, changed by scripts with the proxy's semantics, so +//! gateway replicas and Python proxies sharing that Redis enforce one limit. Windows are fixed and +//! start on their first request. A request counts against requests whether or not it is let +//! through; tokens are reserved up front from an estimate and settled to the real count when the +//! response ends. Without Redis, or while it is unreachable, the counters live in this replica's +//! memory. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use futures_util::future::BoxFuture; +use http::{HeaderName, HeaderValue}; +use litellm_config::Config; +use litellm_db::PgPool; +use litellm_gateway_auth::keys::Limits; +use litellm_gateway_hooks::{ + Facts, Hook, KeyFacts, Record, Refusal, RequestEdit, RequestFacts, Route, Settle, Targets, +}; +use redis::aio::ConnectionManager; +use serde_json::{Map, Value}; + +use crate::{coordination::Coordination, end_users::EndUsers, spend::civil}; + +const DEFAULT_WINDOW: u64 = 60; +/// How long an in-flight slot counts before it is presumed leaked. +const SLOT_TTL: u64 = 3600; +/// The output tokens assumed at most when a request does not cap them. +const MAX_OUTPUT_FLOOR: i64 = 1024; +/// Windows and slots held in memory before expired ones are dropped. +const MAX_LOCAL: usize = 100_000; +const FAIL_CLOSED: &str = "Rate limit enforcement unavailable: request counters could not be verified against Redis, and fail_closed_rate_limit_enforcement is enabled, so the request was rejected to avoid exceeding the configured rate limit. Retry shortly."; + +/// Counts a request on every `(window, requests)` pair, starting a window that is missing or +/// over, and returns each window's start and count. The caller compares; the count always rises. +const REQUESTS: &str = r#" +local results = {} +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local window_start = redis.call('GET', window_key) + if not window_start or (now - tonumber(window_start)) >= window_size then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, 1) + redis.call('EXPIRE', window_key, window_size) + redis.call('EXPIRE', counter_key, window_size) + table.insert(results, tostring(now)) + table.insert(results, 1) + else + local counter = redis.call('INCR', counter_key) + if redis.call('TTL', counter_key) == -1 then + redis.call('EXPIRE', counter_key, window_size) + end + table.insert(results, window_start) + table.insert(results, counter) + end +end +return results +"#; + +/// Adds the increment to every `(window, tokens)` pair, or to none when one would go over its +/// limit: `{0, count, start, ...}` or `{1, index, current, limit}`. +const RESERVE: &str = r#" +local now = tonumber(redis.call('TIME')[1]) +local count = #KEYS / 2 +local state = {} +for i = 1, count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[base]) + local increment = tonumber(ARGV[base + 1]) + local window_size = tonumber(ARGV[base + 3]) + local window_start = redis.call('GET', window_key) + local expired = (not window_start) or ((now - tonumber(window_start)) >= window_size) + local current = 0 + if not expired then + current = tonumber(redis.call('GET', counter_key) or 0) + end + local blocked + if increment > 0 then + blocked = current + increment > limit + else + blocked = current >= limit + end + if blocked then + return { 1, i, current, limit } + end + state[i] = { expired, window_start } +end +local results = { 0 } +local reset = {} +for i = 1, count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[base + 1]) + local ttl = tonumber(ARGV[base + 2]) + local window_size = tonumber(ARGV[base + 3]) + if state[i][1] then + if not reset[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset[window_key] = true + end + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + table.insert(results, now) + else + local counter = redis.call('INCRBY', counter_key, increment) + if redis.call('TTL', counter_key) == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, counter) + table.insert(results, tonumber(state[i][2])) + end +end +return results +"#; + +/// Takes an in-flight slot unless the gauge is full: `{0, in_flight}` or `{1, 1, in_flight, limit}`. +const ACQUIRE: &str = r#" +local now = tonumber(redis.call('TIME')[1]) +local limit = tonumber(ARGV[1]) +local slot_ttl = tonumber(ARGV[2]) +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now - slot_ttl) +local in_flight = redis.call('ZCARD', KEYS[1]) +if in_flight + 1 > limit then + return { 1, 1, in_flight, limit } +end +redis.call('ZADD', KEYS[1], now, ARGV[3]) +redis.call('EXPIRE', KEYS[1], slot_ttl) +return { 0, redis.call('ZCARD', KEYS[1]) } +"#; + +/// Adds each token delta, giving a counter without an expiry the window's. +const CHARGE: &str = r#" +local results = {} +for i = 1, #KEYS do + local value = redis.call('INCRBYFLOAT', KEYS[i], tonumber(ARGV[i * 2 - 1])) + local ttl = tonumber(ARGV[i * 2]) + if ttl and ttl > 0 and redis.call('TTL', KEYS[i]) == -1 then + redis.call('EXPIRE', KEYS[i], ttl) + end + table.insert(results, value) +end +return results +"#; + +#[derive(Clone, Debug, PartialEq)] +pub struct RateLimitSettings { + /// `LITELLM_RATE_LIMIT_WINDOW_SIZE`, in seconds; 60 by default. + pub window: u64, + /// `general_settings.fail_closed_rate_limit_enforcement`: refuse with a 503 rather than fall + /// back to this replica's counters while Redis is unreachable. + pub fail_closed: bool, + /// `litellm_settings.max_end_user_budget_id`: the limits of end users without their own. + pub end_user_budget_id: Option, +} + +impl Default for RateLimitSettings { + fn default() -> Self { + Self { + window: DEFAULT_WINDOW, + fail_closed: false, + end_user_budget_id: None, + } + } +} + +impl RateLimitSettings { + pub fn from_config(config: &Config) -> Self { + let general = &config.general_settings.additional_fields; + Self { + window: std::env::var("LITELLM_RATE_LIMIT_WINDOW_SIZE") + .ok() + .and_then(|window| window.trim().parse().ok()) + .filter(|window| *window > 0) + .unwrap_or(DEFAULT_WINDOW), + fail_closed: match general.get("fail_closed_rate_limit_enforcement") { + Some(litellm_config::Value::Bool(flag)) => *flag, + Some(litellm_config::Value::String(text)) => text.eq_ignore_ascii_case("true"), + _ => false, + }, + end_user_budget_id: config + .litellm_settings + .additional_fields + .get("max_end_user_budget_id") + .and_then(litellm_config::Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned), + } + } +} + +/// One scope a request counts against: `{key}:{value}` with its limits. +#[derive(Clone, Debug, PartialEq)] +struct Descriptor { + key: &'static str, + value: String, + limits: Limits, + max_parallel_requests: Option, +} + +impl Descriptor { + fn prefix(&self) -> String { + format!("{{{}:{}}}", self.key, self.value) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Kind { + Requests, + Tokens, + Parallel, +} + +impl Kind { + fn name(self) -> &'static str { + match self { + Self::Requests => "requests", + Self::Tokens => "tokens", + Self::Parallel => "max_parallel_requests", + } + } +} + +/// Where a descriptor stands after this request. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct Status { + descriptor: usize, + kind: Kind, + limit: i64, + remaining: i64, +} + +/// The token reservation went over on the descriptor at `index`, which held `current`. +#[derive(Debug, PartialEq)] +enum Reserved { + Counts(Vec), + Over { index: usize, current: i64 }, +} + +/// Redis could not be used and the settings say to refuse rather than fall back. +struct Unavailable; + +#[derive(Debug)] +struct Window { + start: u64, + requests: i64, + tokens: i64, +} + +/// This replica's counters, used without Redis or while it is unreachable. +#[derive(Default)] +struct Local { + windows: HashMap, + slots: HashMap>, +} + +impl Local { + fn window(&mut self, prefix: &str, now: u64, size: u64) -> (&mut Window, bool) { + if self.windows.len() >= MAX_LOCAL { + self.windows + .retain(|_, window| now.saturating_sub(window.start) < size); + } + let window = self.windows.entry(prefix.to_owned()).or_insert(Window { + start: now, + requests: 0, + tokens: 0, + }); + let expired = now.saturating_sub(window.start) >= size; + (window, expired) + } + + fn requests(&mut self, prefixes: &[String], now: u64, size: u64) -> Vec { + prefixes + .iter() + .map(|prefix| { + let (window, expired) = self.window(prefix, now, size); + if expired { + *window = Window { + start: now, + requests: 0, + tokens: 0, + }; + } + window.requests += 1; + window.requests + }) + .collect() + } + + fn reserve( + &mut self, + scopes: &[(String, i64)], + increment: i64, + now: u64, + size: u64, + ) -> Reserved { + for (index, (prefix, limit)) in scopes.iter().enumerate() { + let current = match self.windows.get(prefix) { + Some(window) if now.saturating_sub(window.start) < size => window.tokens, + _ => 0, + }; + if blocked(current, increment, *limit) { + return Reserved::Over { index, current }; + } + } + Reserved::Counts( + scopes + .iter() + .map(|(prefix, _)| { + let (window, expired) = self.window(prefix, now, size); + if expired { + *window = Window { + start: now, + requests: 0, + tokens: 0, + }; + } + window.tokens += increment; + window.tokens + }) + .collect(), + ) + } + + fn charge(&mut self, deltas: &[(String, i64)]) { + for (prefix, delta) in deltas { + if let Some(window) = self.windows.get_mut(prefix) { + window.tokens += delta; + } + } + } + + fn acquire(&mut self, prefix: &str, limit: i64, slot: &str) -> Result { + let ttl = Duration::from_secs(SLOT_TTL); + if self.slots.len() >= MAX_LOCAL { + self.slots.retain(|_, slots| !slots.is_empty()); + } + let slots = self.slots.entry(prefix.to_owned()).or_default(); + slots.retain(|_, taken| taken.elapsed() < ttl); + let in_flight = slots.len() as i64; + if in_flight + 1 > limit { + return Err(in_flight); + } + slots.insert(slot.to_owned(), Instant::now()); + Ok(in_flight + 1) + } + + fn release(&mut self, prefix: &str, slot: &str) { + if let Some(slots) = self.slots.get_mut(prefix) { + slots.remove(slot); + } + } +} + +fn blocked(current: i64, increment: i64, limit: i64) -> bool { + match increment > 0 { + true => current + increment > limit, + false => current >= limit, + } +} + +/// Where the counters live. +struct Store { + redis: Option>, + local: Mutex, + window: u64, + fail_closed: bool, +} + +impl Store { + fn connection(&self) -> Option<(ConnectionManager, &Coordination)> { + let redis = self.redis.as_deref()?; + Some((redis.connection()?, redis)) + } + + fn local(&self) -> std::sync::MutexGuard<'_, Local> { + self.local.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Logs why Redis was not used; refuses when the settings say to. + fn fall_back(&self, error: Option) -> Result<(), Unavailable> { + match error { + Some(error) => { + tracing::warn!(%error, "rate limit counters failed; using this replica's") + } + None => { + tracing::debug!("Redis not connected; using this replica's rate limit counters") + } + } + match self.fail_closed { + true => Err(Unavailable), + false => Ok(()), + } + } + + /// Counts the request on each prefix's window and returns the counts. + async fn requests(&self, prefixes: &[String]) -> Result, Unavailable> { + let now = epoch_seconds(); + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let mut command = redis::cmd("EVAL"); + command.arg(REQUESTS).arg(prefixes.len() * 2); + for prefix in prefixes { + command + .arg(redis.namespaced(&format!("{prefix}:window"))) + .arg(redis.namespaced(&format!("{prefix}:requests"))); + } + command.arg(now).arg(self.window); + match integers(command.query_async(&mut connection).await) { + Ok(reply) if reply.len() == prefixes.len() * 2 => { + return Ok(reply.iter().skip(1).step_by(2).copied().collect()); + } + Ok(_) => Some(malformed()), + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok(self.local().requests(prefixes, now, self.window)) + } + + /// Takes an in-flight slot: the count with it, or the full count without. `true` when the + /// slot is in Redis. + async fn acquire( + &self, + prefix: &str, + limit: i64, + slot: &str, + ) -> Result<(Result, bool), Unavailable> { + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let reply = redis::cmd("EVAL") + .arg(ACQUIRE) + .arg(1) + .arg(redis.namespaced(&format!("{prefix}:max_parallel_requests"))) + .arg(limit) + .arg(SLOT_TTL) + .arg(slot) + .query_async(&mut connection) + .await; + match integers(reply) { + Ok(reply) => match reply.as_slice() { + [0, in_flight] => return Ok((Ok(*in_flight), true)), + [1, _, in_flight, _] => return Ok((Err(*in_flight), true)), + _ => Some(malformed()), + }, + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok((self.local().acquire(prefix, limit, slot), false)) + } + + async fn release(&self, prefix: &str, slot: &str, shared: bool) { + if shared && let Some((mut connection, redis)) = self.connection() { + let key = redis.namespaced(&format!("{prefix}:max_parallel_requests")); + if let Err(error) = redis::cmd("ZREM") + .arg(key) + .arg(slot) + .query_async::(&mut connection) + .await + { + tracing::warn!(%error, "releasing an in-flight slot failed; it expires on its own"); + } + return; + } + self.local().release(prefix, slot); + } + + /// Reserves `increment` tokens on every `(prefix, limit)`, or on none. `true` when the + /// reservation is in Redis. + async fn reserve( + &self, + scopes: &[(String, i64)], + increment: i64, + ) -> Result<(Reserved, bool), Unavailable> { + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let mut command = redis::cmd("EVAL"); + command.arg(RESERVE).arg(scopes.len() * 2); + for (prefix, _) in scopes { + command + .arg(redis.namespaced(&format!("{prefix}:window"))) + .arg(redis.namespaced(&format!("{prefix}:tokens"))); + } + for (_, limit) in scopes { + command + .arg(*limit) + .arg(increment) + .arg(self.window) + .arg(self.window); + } + match integers(command.query_async(&mut connection).await) { + Ok(reply) if reply.first() == Some(&0) => { + return Ok(( + Reserved::Counts( + reply.iter().skip(1).step_by(2).copied().collect(), + ), + true, + )); + } + Ok(reply) => match reply.as_slice() { + [1, index, current, _] if (1..=scopes.len() as i64).contains(index) => { + return Ok(( + Reserved::Over { + index: (*index - 1) as usize, + current: *current, + }, + true, + )); + } + _ => Some(malformed()), + }, + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok(( + self.local() + .reserve(scopes, increment, epoch_seconds(), self.window), + false, + )) + } + + /// Adds each delta to its prefix's tokens. + async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) { + if shared && let Some((mut connection, redis)) = self.connection() { + let mut command = redis::cmd("EVAL"); + command.arg(CHARGE).arg(deltas.len()); + for (prefix, _) in &deltas { + command.arg(redis.namespaced(&format!("{prefix}:tokens"))); + } + for (_, delta) in &deltas { + command.arg(*delta).arg(self.window); + } + if let Err(error) = command.query_async::(&mut connection).await { + tracing::warn!(%error, "settling reserved tokens failed"); + } + return; + } + self.local().charge(&deltas); + } +} + +fn malformed() -> redis::RedisError { + redis::RedisError::from(( + redis::ErrorKind::UnexpectedReturnType, + "unexpected rate limit script reply", + )) +} + +/// A script's array reply as integers; counters may come back as strings. +fn integers(reply: redis::RedisResult>) -> redis::RedisResult> { + reply? + .iter() + .map(|value| { + let text = match value { + redis::Value::Int(number) => return Ok(*number), + redis::Value::BulkString(bytes) => std::str::from_utf8(bytes).ok(), + redis::Value::SimpleString(text) => Some(text.as_str()), + _ => None, + }; + text.and_then(|text| text.parse::().ok()) + .map(|number| number as i64) + .ok_or_else(malformed) + }) + .collect() +} + +fn epoch_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +/// `%Y-%m-%d %H:%M:%S UTC`, as the proxy prints when a limit resets. +fn utc(seconds: u64) -> String { + let (year, month, day) = civil((seconds / 86_400) as i64); + let time = seconds % 86_400; + format!( + "{year:04}-{month:02}-{day:02} {:02}:{:02}:{:02} UTC", + time / 3600, + time % 3600 / 60, + time % 60 + ) +} + +/// The proxy's token estimate for a request: a token per four characters of message text, +/// plus the output it may produce for every candidate. +#[derive(Debug, PartialEq)] +struct Estimate { + tokens: i64, + /// The output cap to send upstream when the request has none and the limits are small. + implicit_max_tokens: Option, +} + +fn estimate( + body: &Map, + key: Option<&KeyFacts>, + model: &str, + min_tpm: i64, +) -> Estimate { + let chars: usize = body + .get("messages") + .and_then(Value::as_array) + .into_iter() + .flatten() + .map(|message| match message.get("content") { + Some(Value::String(text)) => text.chars().count(), + Some(Value::Array(blocks)) => blocks + .iter() + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .map(|text| text.chars().count()) + .sum(), + _ => 0, + }) + .sum(); + let chars = i64::try_from(chars).unwrap_or(i64::MAX); + let input = match chars { + 0 => 0, + chars => (chars / 4).max(1), + }; + let integer = |value: &Value| match value { + Value::Number(number) => number.as_f64().map(|number| number as i64), + Value::String(text) => text.trim().parse::().ok().map(|number| number as i64), + _ => None, + }; + let cap = ["max_tokens", "max_completion_tokens"] + .iter() + .filter_map(|field| body.get(*field).and_then(integer)) + .max(); + let configured = key.and_then(|key| { + [&key.metadata, &key.team_metadata] + .into_iter() + .find_map(|metadata| { + metadata + .get("default_estimated_output_tokens_per_model") + .and_then(|models| models.get(model)) + .and_then(integer) + .or_else(|| { + metadata + .get("default_estimated_output_tokens") + .and_then(integer) + }) + }) + }); + let floor = (min_tpm / 4).clamp(1, MAX_OUTPUT_FLOOR); + let output = match (cap, configured) { + (Some(cap), _) => cap, + (None, Some(configured)) => configured, + (None, None) if chars == 0 => 0, + (None, None) => input.max(floor), + }; + let candidates = ["n", "best_of"] + .iter() + .filter_map(|field| body.get(*field).and_then(integer)) + .fold(1, i64::max); + Estimate { + tokens: input + .saturating_add(output.saturating_mul(candidates)) + .max(1), + implicit_max_tokens: (floor < MAX_OUTPUT_FLOOR && cap.is_none()) + .then(|| floor.max(configured.unwrap_or(0))), + } +} + +/// The `rate_limits` hook. +pub struct RateLimiter { + settings: RateLimitSettings, + store: Arc, + end_users: Option, +} + +impl RateLimiter { + /// `redis`, when given, holds the shared counters; `database` gives end users' limits. + pub fn new( + settings: RateLimitSettings, + redis: Option>, + database: Option, + ) -> Self { + Self { + store: Arc::new(Store { + redis, + local: Mutex::default(), + window: settings.window, + fail_closed: settings.fail_closed, + }), + settings, + end_users: database.map(EndUsers::new), + } + } + + pub fn for_gateway( + config: &Config, + redis: Option>, + database: Option, + ) -> Self { + Self::new(RateLimitSettings::from_config(config), redis, database) + } + + /// The scopes this request counts against, in the proxy's order. + async fn descriptors(&self, request: &RequestFacts) -> Vec { + let Some(key) = request.key.as_deref() else { + return Vec::new(); + }; + let limits = &key.limits; + let mut descriptors = Vec::new(); + let mut push = |key: &'static str, value: String, limits: Limits, parallel: Option| { + if limits.is_set() || parallel.is_some() { + descriptors.push(Descriptor { + key, + value, + limits, + max_parallel_requests: parallel, + }); + } + }; + if !key.hash.is_empty() { + push( + "api_key", + key.hash.clone(), + limits.key, + limits.max_parallel_requests, + ); + } + if let Some(user_id) = &key.user_id { + push("user", user_id.clone(), limits.user, None); + } + if let Some(team_id) = &key.team_id { + push("team", team_id.clone(), limits.team, None); + } + if let (Some(team_id), Some(user_id)) = (&key.team_id, &key.user_id) { + push( + "team_member", + format!("{team_id}:{user_id}"), + limits.team_member, + None, + ); + } + if let (Some(end_user_id), Some(end_users)) = (request.end_user.as_deref(), &self.end_users) + && let Some(end_user) = end_users + .get( + end_user_id, + self.settings.end_user_budget_id.as_deref(), + key.budgets.end_user_budget_id.as_deref(), + ) + .await + { + push( + "end_user", + end_user_id.to_owned(), + Limits { + rpm: end_user.limits.rpm, + tpm: end_user.limits.tpm, + }, + None, + ); + } + descriptors + } + + fn refusal(&self, request: &RequestFacts, descriptor: &Descriptor, status: Status) -> Refusal { + let reset_at = utc(epoch_seconds() + self.settings.window); + let message = format!( + "Rate limit exceeded for {}: {}. Limit type: {}. Current limit: {}, Remaining: {}. Limit resets at: {reset_at}", + descriptor.key, + descriptor.value, + status.kind.name(), + status.limit, + status.remaining.max(0), + ); + let mut refusal = match request.route { + Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(), + Route::ChatCompletions => { + let mut refusal = Refusal::new(429, "throttling_error", message).without_param(); + for (name, value) in [ + ("retry-after", self.settings.window.to_string()), + ("rate_limit_type", status.kind.name().to_owned()), + ("reset_at", reset_at), + ] { + if let Ok(value) = HeaderValue::from_str(&value) { + refusal = refusal.with_header(HeaderName::from_static(name), value); + } + } + refusal + } + }; + for (name, value) in key_limit_headers(request.key.as_deref()) { + refusal = refusal.with_header(name, value); + } + refusal + } + + fn unavailable(&self, request: &RequestFacts) -> Refusal { + match request.route { + Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(), + Route::ChatCompletions => { + Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param() + } + } + } +} + +/// `x-litellm-key-tpm-limit` and `x-litellm-key-rpm-limit`, which the proxy sends whenever the +/// key sets them. +fn key_limit_headers(key: Option<&KeyFacts>) -> Vec<(HeaderName, HeaderValue)> { + let Some(key) = key else { + return Vec::new(); + }; + [ + ("x-litellm-key-tpm-limit", key.limits.key.tpm), + ("x-litellm-key-rpm-limit", key.limits.key.rpm), + ] + .into_iter() + .filter_map(|(name, limit)| Some((HeaderName::from_static(name), HeaderValue::from(limit?)))) + .collect() +} + +impl Hook for RateLimiter { + fn name(&self) -> &'static str { + "rate_limits" + } + + fn reads(&self) -> Facts { + Facts::KEY | Facts::END_USER | Facts::REQUEST_BODY | Facts::USAGE + } + + fn writes(&self) -> Targets { + Targets::ADMISSION | Targets::UPSTREAM_BODY | Targets::RESPONSE_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + let descriptors = self.descriptors(request).await; + if descriptors.is_empty() { + return Ok(None); + } + let unavailable = |Unavailable| self.unavailable(request); + let mut statuses = Vec::new(); + + // Requests per window, counted whether or not the request goes through. + let windowed: Vec = (0..descriptors.len()) + .filter(|index| descriptors[*index].limits.rpm.is_some()) + .collect(); + if !windowed.is_empty() { + let prefixes: Vec = windowed + .iter() + .map(|index| descriptors[*index].prefix()) + .collect(); + let counts = self.store.requests(&prefixes).await.map_err(unavailable)?; + for (index, count) in windowed.iter().zip(counts) { + let limit = descriptors[*index].limits.rpm.unwrap_or_default(); + statuses.push(Status { + descriptor: *index, + kind: Kind::Requests, + limit, + remaining: limit - count, + }); + } + if let Some(over) = statuses.iter().find(|status| status.remaining < 0) { + return Err(self.refusal(request, &descriptors[over.descriptor], *over)); + } + } + + // Requests in flight, for the key only. + let mut slot = None; + if let Some(index) = descriptors + .iter() + .position(|descriptor| descriptor.max_parallel_requests.is_some()) + { + let descriptor = &descriptors[index]; + let limit = descriptor.max_parallel_requests.unwrap_or_default(); + let prefix = descriptor.prefix(); + let id = uuid::Uuid::new_v4().simple().to_string(); + let (taken, shared) = self + .store + .acquire(&prefix, limit, &id) + .await + .map_err(unavailable)?; + let in_flight = match taken { + Ok(in_flight) => in_flight, + Err(in_flight) => { + let status = Status { + descriptor: index, + kind: Kind::Parallel, + limit, + remaining: limit - in_flight, + }; + return Err(self.refusal(request, descriptor, status)); + } + }; + statuses.push(Status { + descriptor: index, + kind: Kind::Parallel, + limit, + remaining: (limit - in_flight).max(0), + }); + slot = Some(Slot { prefix, id, shared }); + } + + // Tokens per window, reserved from an estimate and settled when the response ends. + let limited: Vec<(usize, i64)> = descriptors + .iter() + .enumerate() + .filter_map(|(index, descriptor)| Some((index, descriptor.limits.tpm?))) + .collect(); + let mut tokens = None; + if let Some(min_tpm) = limited.iter().map(|(_, limit)| *limit).min() { + let body: Map = + serde_json::from_slice(request.body()).unwrap_or_default(); + let estimate = estimate( + &body, + request.key.as_deref(), + &request.requested_model, + min_tpm, + ); + if let Some(max_tokens) = estimate.implicit_max_tokens { + edit.set_upstream_body_field("max_tokens", Value::from(max_tokens)); + } + let scopes: Vec<(String, i64)> = limited + .iter() + .map(|(index, limit)| (descriptors[*index].prefix(), *limit)) + .collect(); + let reserved = self.store.reserve(&scopes, estimate.tokens).await; + let release = |slot: Option| async { + if let Some(slot) = slot { + self.store + .release(&slot.prefix, &slot.id, slot.shared) + .await; + } + }; + let (reserved, shared) = match reserved { + Ok(reserved) => reserved, + Err(Unavailable) => { + release(slot).await; + return Err(self.unavailable(request)); + } + }; + match reserved { + Reserved::Counts(counts) => { + for ((index, limit), count) in limited.iter().zip(counts) { + statuses.push(Status { + descriptor: *index, + kind: Kind::Tokens, + limit: *limit, + remaining: (limit - count).max(0), + }); + } + } + Reserved::Over { index, current } => { + release(slot).await; + let (descriptor, limit) = limited[index]; + let status = Status { + descriptor, + kind: Kind::Tokens, + limit, + remaining: limit - current, + }; + return Err(self.refusal(request, &descriptors[descriptor], status)); + } + } + tokens = Some(Tokens { + prefixes: scopes.into_iter().map(|(prefix, _)| prefix).collect(), + reserved: estimate.tokens, + shared, + }); + } + + // The proxy sends where each limit stands on complete responses only. + if !request.stream { + for status in &statuses { + let key = descriptors[status.descriptor].key; + let kind = status.kind.name(); + for (name, value) in [ + ( + format!("x-ratelimit-{key}-remaining-{kind}"), + status.remaining, + ), + (format!("x-ratelimit-{key}-limit-{kind}"), status.limit), + ] { + if let Ok(name) = HeaderName::from_bytes(name.as_bytes()) { + edit.add_response_header(name, HeaderValue::from(value)); + } + } + } + } + for (name, value) in key_limit_headers(request.key.as_deref()) { + edit.add_response_header(name, value); + } + if slot.is_none() && tokens.is_none() { + return Ok(None); + } + Ok(Some(Box::new(Held { + store: self.store.clone(), + slot, + tokens, + }) as Box)) + }) + } +} + +struct Slot { + prefix: String, + id: String, + shared: bool, +} + +struct Tokens { + prefixes: Vec, + reserved: i64, + shared: bool, +} + +/// What a request holds until it ends: an in-flight slot and reserved tokens. +struct Held { + store: Arc, + slot: Option, + tokens: Option, +} + +impl Settle for Held { + fn settle(self: Box, record: &Record) { + let Held { + store, + slot, + tokens, + } = *self; + let charge = tokens.and_then(|tokens| { + // The proxy counts every token but cache reads. Without usage, a success keeps its + // reservation and a failure returns it. + let actual = match (&record.outcome.usage, record.succeeded()) { + (Some(usage), _) => i64::try_from(usage.total_tokens() - usage.cache_read_tokens) + .unwrap_or(i64::MAX), + (None, true) => tokens.reserved, + (None, false) => 0, + }; + let delta = actual - tokens.reserved; + (delta != 0).then(|| { + ( + tokens + .prefixes + .into_iter() + .map(|prefix| (prefix, delta)) + .collect::>(), + tokens.shared, + ) + }) + }); + tokio::spawn(async move { + if let Some(slot) = slot { + store.release(&slot.prefix, &slot.id, slot.shared).await; + } + if let Some((deltas, shared)) = charge { + store.charge(deltas, shared).await; + } + }); + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn body(value: Value) -> Map { + match value { + Value::Object(object) => object, + _ => unreachable!(), + } + } + + #[rstest] + #[case::utc_epoch(0, "1970-01-01 00:00:00 UTC")] + #[case::utc_later(1_792_238_465, "2026-10-17 12:01:05 UTC")] + fn prints_reset_times_like_the_proxy(#[case] seconds: u64, #[case] expected: &str) { + assert_eq!(utc(seconds), expected); + } + + #[rstest] + // 40 characters: 10 input tokens; the small limit floors output at 100_000 / 4 = 1024. + #[case::uncapped(json!({"messages": [{"role": "user", "content": "a".repeat(40)}]}), 100_000, 10 + 1024, None)] + // A small limit floors output at 400 / 4 = 100 and caps the request there. + #[case::small_limit(json!({"messages": [{"role": "user", "content": "a".repeat(40)}]}), 400, 110, Some(100))] + #[case::capped(json!({"max_tokens": 50, "messages": [{"role": "user", "content": [{"type": "text", "text": "a".repeat(8)}]}]}), 400, 52, None)] + #[case::candidates(json!({"max_completion_tokens": 5, "n": 3, "messages": [{"role": "user", "content": "abcd"}]}), 100_000, 16, None)] + #[case::no_text(json!({"messages": []}), 100_000, 1, None)] + fn estimates_tokens_like_the_proxy( + #[case] request: Value, + #[case] min_tpm: i64, + #[case] tokens: i64, + #[case] implicit_max_tokens: Option, + ) { + assert_eq!( + estimate(&body(request), None, "m", min_tpm), + Estimate { + tokens, + implicit_max_tokens + } + ); + } + + #[rstest] + fn configured_output_estimates_come_from_the_key_then_the_team() { + let metadata = |value: Value| Arc::new(body(value)); + let key = KeyFacts { + metadata: metadata(json!({"default_estimated_output_tokens_per_model": {"m": 7}})), + team_metadata: metadata(json!({"default_estimated_output_tokens": 9})), + ..KeyFacts::default() + }; + let request = body(json!({"messages": [{"role": "user", "content": "abcd"}]})); + assert_eq!(estimate(&request, Some(&key), "m", 100_000).tokens, 8); + assert_eq!(estimate(&request, Some(&key), "other", 100_000).tokens, 10); + } + + #[rstest] + fn local_windows_count_reserve_and_reset() { + let mut local = Local::default(); + let prefixes = vec!["{api_key:h}".to_owned()]; + assert_eq!(local.requests(&prefixes, 100, 60), [1]); + assert_eq!(local.requests(&prefixes, 159, 60), [2]); + let scopes = vec![("{api_key:h}".to_owned(), 10)]; + assert_eq!( + local.reserve(&scopes, 6, 159, 60), + Reserved::Counts(vec![6]) + ); + assert_eq!( + local.reserve(&scopes, 6, 159, 60), + Reserved::Over { + index: 0, + current: 6 + } + ); + local.charge(&[("{api_key:h}".to_owned(), -4)]); + assert_eq!( + local.reserve(&scopes, 8, 159, 60), + Reserved::Counts(vec![10]) + ); + // The window started at 100 is over at 160: requests and tokens start again. + assert_eq!(local.requests(&prefixes, 160, 60), [1]); + assert_eq!( + local.reserve(&scopes, 10, 160, 60), + Reserved::Counts(vec![10]) + ); + } + + #[rstest] + fn local_slots_fill_and_release() { + let mut local = Local::default(); + assert_eq!(local.acquire("{api_key:h}", 2, "a"), Ok(1)); + assert_eq!(local.acquire("{api_key:h}", 2, "b"), Ok(2)); + assert_eq!(local.acquire("{api_key:h}", 2, "c"), Err(2)); + local.release("{api_key:h}", "a"); + assert_eq!(local.acquire("{api_key:h}", 2, "c"), Ok(2)); + } +} diff --git a/litellm-rust/crates/gateway/src/spend.rs b/litellm-rust/crates/gateway/src/spend.rs index fc373307bec..a2f0d077819 100644 --- a/litellm-rust/crates/gateway/src/spend.rs +++ b/litellm-rust/crates/gateway/src/spend.rs @@ -648,7 +648,7 @@ fn usage_values(usage: &Usage) -> (Value, Value) { } /// Days since 1970-01-01 as a UTC calendar date. -fn civil(days: i64) -> (i64, u32, u32) { +pub(crate) fn civil(days: i64) -> (i64, u32, u32) { let days = days + 719_468; let era = days.div_euclid(146_097); let day_of_era = days.rem_euclid(146_097); diff --git a/litellm-rust/crates/gateway/tests/rate_limits.rs b/litellm-rust/crates/gateway/tests/rate_limits.rs new file mode 100644 index 00000000000..2e7101bd89c --- /dev/null +++ b/litellm-rust/crates/gateway/tests/rate_limits.rs @@ -0,0 +1,515 @@ +//! The rate limits hook on real requests: refusals and headers as the proxy sends them, tokens +//! reserved and settled to usage, in-flight slots, and counters shared through Redis. Set +//! `LITELLM_TEST_REDIS_URL` to also run the shared cases against a live Redis. + +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{HeaderMap, Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway::{ + coordination::{Coordination, RedisSettings}, + rate_limits::{RateLimitSettings, RateLimiter}, +}; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLimits, KeyLookup, KeyModels, KeyProfile, KeyStatus, Limits, + LookupFuture, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +/// 3 input and 4 output tokens: 7 count against a token limit. +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; + +#[derive(Clone, Default)] +struct Upstream { + calls: Arc, + bodies: Arc>>, +} + +async fn upstream(delay: Duration) -> (String, Upstream) { + let upstream = Upstream::default(); + let seen = upstream.clone(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(move |request: Request| { + let seen = seen.clone(); + async move { + seen.calls.fetch_add(1, Ordering::SeqCst); + let body = to_bytes(request.into_body(), 1 << 20).await.unwrap(); + seen.bodies + .lock() + .unwrap() + .push(serde_json::from_slice(&body).unwrap_or(Value::Null)); + tokio::time::sleep(delay).await; + ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response() + } + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + (format!("http://{address}"), upstream) +} + +struct Keys(HashMap); + +impl KeyLookup for Keys { + fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> { + Box::pin(async move { Ok(self.0.get(hash).cloned()) }) + } +} + +fn hash(token: &str) -> KeyHash { + KeyHash::from_token(&SecretValue::new(token)) +} + +fn key(team_id: Option<&str>, limits: KeyLimits) -> KeyStatus { + KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: team_id.map(str::to_owned), + profile: Arc::new(KeyProfile { + user_id: Some("user-1".into()), + limits, + ..KeyProfile::default() + }), + } +} + +fn rpm(limit: i64) -> Limits { + Limits { + rpm: Some(limit), + tpm: None, + } +} + +fn tpm(limit: i64) -> Limits { + Limits { + rpm: None, + tpm: Some(limit), + } +} + +struct Gateway { + app: Router, + upstream: Upstream, +} + +struct Setup { + redis: Option>, + database: Option, + delay: Duration, +} + +impl Default for Setup { + fn default() -> Self { + Self { + redis: None, + database: None, + delay: Duration::ZERO, + } + } +} + +async fn gateway(keys: Vec<(&str, KeyStatus)>, setup: Setup) -> Gateway { + let (address, upstream) = upstream(setup.delay).await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {address} + api_key: sk-upstream +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let limiter = RateLimiter::new(RateLimitSettings::default(), setup.redis, setup.database); + litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(limiter)], Vec::new()) + .await + .unwrap(); + let keys: Arc = Arc::new(Keys( + keys.into_iter() + .map(|(token, status)| (hash(token), status)) + .collect(), + )); + let app = litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys)); + Gateway { app, upstream } +} + +struct Reply { + status: u16, + headers: HeaderMap, + body: Value, +} + +async fn send(app: &Router, path: &str, body: Value) -> Reply { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer sk-key") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + Reply { + status, + headers, + body: serde_json::from_slice(&body).unwrap_or(Value::Null), + } +} + +fn message(extra: Value) -> Value { + let mut body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + for (field, value) in extra.as_object().unwrap() { + body[field] = value.clone(); + } + body +} + +async fn messages(app: &Router) -> Reply { + send(app, "/v1/messages", message(json!({}))).await +} + +async fn chat(app: &Router) -> Reply { + send(app, "/v1/chat/completions", message(json!({}))).await +} + +/// Settling runs after the response, off the request path. +async fn settled() { + tokio::time::sleep(Duration::from_millis(100)).await; +} + +fn header<'a>(reply: &'a Reply, name: &str) -> Option<&'a str> { + reply.headers.get(name).map(|value| value.to_str().unwrap()) +} + +/// The message up to the reset time, which moves with the clock. +fn refusal_message(reply: &Reply, path: &[&str]) -> String { + let mut value = &reply.body; + for part in path { + value = &value[*part]; + } + let message = value.as_str().unwrap(); + let (message, reset) = message.split_once(" Limit resets at: ").unwrap(); + assert!(reset.ends_with(" UTC"), "{reset}"); + message.to_owned() +} + +#[rstest] +#[tokio::test] +async fn requests_over_the_key_limit_are_refused_like_the_proxy() { + let limits = KeyLimits { + key: rpm(2), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let first = messages(&gateway.app).await; + assert_eq!(first.status, 200); + assert_eq!( + header(&first, "x-ratelimit-api_key-remaining-requests"), + Some("1") + ); + assert_eq!( + header(&first, "x-ratelimit-api_key-limit-requests"), + Some("2") + ); + assert_eq!(header(&first, "x-litellm-key-rpm-limit"), Some("2")); + assert_eq!(chat(&gateway.app).await.status, 200); + + let expected = format!( + "Rate limit exceeded for api_key: {}. Limit type: requests. Current limit: 2, Remaining: 0.", + hash("sk-key").as_str() + ); + // Chat Completions: the proxy's error body and its retry headers. + let refused = chat(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!(refusal_message(&refused, &["error", "message"]), expected); + assert_eq!(refused.body["error"]["type"], "throttling_error"); + assert_eq!(refused.body["error"]["param"], Value::Null); + assert_eq!(refused.body["error"]["code"], "429"); + assert_eq!(header(&refused, "retry-after"), Some("60")); + assert_eq!(header(&refused, "rate_limit_type"), Some("requests")); + assert!(header(&refused, "reset_at").unwrap().ends_with(" UTC")); + assert_eq!(header(&refused, "x-litellm-key-rpm-limit"), Some("2")); + // Messages: Anthropic's envelope, without the retry headers. + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!(refused.body["type"], "error"); + assert_eq!(refused.body["error"]["type"], "rate_limit_error"); + assert_eq!(refusal_message(&refused, &["error", "message"]), expected); + assert_eq!(header(&refused, "retry-after"), None); + assert_eq!(gateway.upstream.calls.load(Ordering::SeqCst), 2); +} + +#[rstest] +#[case::user(KeyLimits { user: rpm(1), ..KeyLimits::default() }, "user: user-1")] +#[case::team(KeyLimits { team: rpm(1), ..KeyLimits::default() }, "team: team-1")] +#[case::team_member(KeyLimits { team_member: rpm(1), ..KeyLimits::default() }, "team_member: team-1:user-1")] +#[tokio::test] +async fn owners_are_limited_under_their_own_names( + #[case] limits: KeyLimits, + #[case] descriptor: &str, +) { + let gateway = gateway( + vec![("sk-key", key(Some("team-1"), limits))], + Setup::default(), + ) + .await; + assert_eq!(messages(&gateway.app).await.status, 200); + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + format!( + "Rate limit exceeded for {descriptor}. Limit type: requests. Current limit: 1, Remaining: 0." + ) + ); +} + +#[rstest] +#[tokio::test] +async fn tokens_are_reserved_and_settled_to_usage() { + // Each request reserves 1 input + 8 output tokens and settles to 7. + let limits = KeyLimits { + key: tpm(20), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let first = messages(&gateway.app).await; + assert_eq!(first.status, 200); + assert_eq!( + header(&first, "x-ratelimit-api_key-remaining-tokens"), + Some("11") + ); + assert_eq!(header(&first, "x-litellm-key-tpm-limit"), Some("20")); + settled().await; + assert_eq!(messages(&gateway.app).await.status, 200); + settled().await; + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + format!( + "Rate limit exceeded for api_key: {}. Limit type: tokens. Current limit: 20, Remaining: 6.", + hash("sk-key").as_str() + ) + ); +} + +#[rstest] +#[tokio::test] +async fn small_token_limits_cap_uncapped_requests() { + let limits = KeyLimits { + key: tpm(400), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let body = json!({"model": "claude", "messages": [{"role": "user", "content": "hi"}]}); + assert_eq!(send(&gateway.app, "/v1/messages", body).await.status, 200); + assert_eq!( + gateway.upstream.bodies.lock().unwrap()[0]["max_tokens"], + 100 + ); +} + +#[rstest] +#[tokio::test] +async fn in_flight_requests_are_capped_per_key() { + let limits = KeyLimits { + max_parallel_requests: Some(1), + ..KeyLimits::default() + }; + let gateway = gateway( + vec![("sk-key", key(None, limits))], + Setup { + delay: Duration::from_millis(300), + ..Setup::default() + }, + ) + .await; + let (first, second) = tokio::join!(messages(&gateway.app), async { + tokio::time::sleep(Duration::from_millis(100)).await; + messages(&gateway.app).await + }); + assert_eq!(first.status, 200); + assert_eq!( + header( + &first, + "x-ratelimit-api_key-remaining-max_parallel_requests" + ), + Some("0") + ); + assert_eq!(second.status, 429); + assert_eq!( + refusal_message(&second, &["error", "message"]), + format!( + "Rate limit exceeded for api_key: {}. Limit type: max_parallel_requests. Current limit: 1, Remaining: 0.", + hash("sk-key").as_str() + ) + ); + settled().await; + assert_eq!(messages(&gateway.app).await.status, 200); +} + +fn redis_url() -> Option { + std::env::var("LITELLM_TEST_REDIS_URL").ok() +} + +async fn redis(url: &str, namespace: &str) -> Arc { + let coordination = + Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap(); + let router: Arc = Arc::new(NoRouter( + litellm_gateway_inference::ModelRouter::from_config( + &Config::from_yaml("model_list: []").unwrap(), + ) + .unwrap(), + )); + coordination.start(router, None, tokio_util::sync::CancellationToken::new()); + for _ in 0..100 { + if coordination.connected() { + return coordination; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("Redis did not connect"); +} + +struct NoRouter(litellm_gateway_inference::ModelRouter); + +impl litellm_gateway::coordination::RouterHandle for NoRouter { + fn router(&self) -> &litellm_gateway_inference::ModelRouter { + &self.0 + } +} + +#[rstest] +#[tokio::test] +async fn replicas_share_the_proxy_counters() { + let Some(url) = redis_url() else { return }; + let namespace = format!("rate-limits-{}", uuid::Uuid::new_v4()); + let limits = KeyLimits { + key: Limits { + rpm: Some(3), + tpm: Some(1_000), + }, + max_parallel_requests: Some(5), + ..KeyLimits::default() + }; + let mut replicas = Vec::new(); + for _ in 0..2 { + let setup = Setup { + redis: Some(redis(&url, &namespace).await), + ..Setup::default() + }; + replicas.push(gateway(vec![("sk-key", key(None, limits))], setup).await); + } + assert_eq!(messages(&replicas[0].app).await.status, 200); + assert_eq!(messages(&replicas[1].app).await.status, 200); + let third = messages(&replicas[0].app).await; + assert_eq!( + header(&third, "x-ratelimit-api_key-remaining-requests"), + Some("0") + ); + settled().await; + assert_eq!(messages(&replicas[1].app).await.status, 429); + + let mut connection = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let prefix = format!("{namespace}:{{api_key:{}}}", hash("sk-key").as_str()); + let get = async |connection: &mut redis::aio::MultiplexedConnection, suffix: &str| { + redis::cmd("GET") + .arg(format!("{prefix}:{suffix}")) + .query_async::>(connection) + .await + .unwrap() + }; + // The refused request still counted; three settled requests used 7 tokens each. + assert_eq!(get(&mut connection, "requests").await.as_deref(), Some("4")); + assert_eq!(get(&mut connection, "tokens").await.as_deref(), Some("21")); + assert!(get(&mut connection, "window").await.is_some()); + let in_flight: i64 = redis::cmd("ZCARD") + .arg(format!("{prefix}:max_parallel_requests")) + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(in_flight, 0); +} + +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows; the lookups under test use the checked macro" +)] +async fn end_users(pool: &litellm_db::PgPool) { + sqlx::raw_sql( + r#" +INSERT INTO "LiteLLM_BudgetTable" (budget_id, rpm_limit, created_by, updated_by) + VALUES ('one-a-minute', 1, 'x', 'x'); +INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id) + VALUES ('limited', 0, 'one-a-minute'), ('free', 0, NULL); +"#, + ) + .execute(pool) + .await + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn end_users_are_held_to_their_budget_limits() { + let postgres = litellm_db_testing::MigratedPostgres::start().await.unwrap(); + end_users(postgres.pool()).await; + let gateway = gateway( + vec![("sk-key", key(None, KeyLimits::default()))], + Setup { + database: Some(postgres.pool().clone()), + ..Setup::default() + }, + ) + .await; + let as_end_user = |user: &str| message(json!({"metadata": {"user_id": user}})); + for user in ["limited", "free", "free"] { + assert_eq!( + send(&gateway.app, "/v1/messages", as_end_user(user)) + .await + .status, + 200 + ); + } + let refused = send(&gateway.app, "/v1/messages", as_end_user("limited")).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + "Rate limit exceeded for end_user: limited. Limit type: requests. Current limit: 1, Remaining: 0." + ); +} diff --git a/litellm-rust/docs/customer-proxy/rate-limits.svg b/litellm-rust/docs/customer-proxy/rate-limits.svg new file mode 100644 index 00000000000..beff70d8482 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/rate-limits.svg @@ -0,0 +1,55 @@ + + + + + + + Rate limits hook: the proxy's limiter on the proxy's shared counters + + + Scopes, in the proxy's order + api_key: key rpm/tpm (else budget row's), + max_parallel_requests + user, team, team_member (team:user) + end_user: its budget row, or the default + read with the key; end users cached 60 s + no limits set: no Redis call at all + + + Three Redis scripts, in the proxy's phases + 1. requests: +1 on every {scope:id}:requests in one call, fixed window from the first hit; + over when count > limit; the count is never given back + 2. in flight (key only): ZSET of slots on Redis TIME, pruned at 3600 s, full at limit + 3. tokens: reserve chars/4 + output x candidates on every scope, all or none; + output = max_tokens, else the configured estimate, else max(input, min(1024, tpm/4)); + a limit under 4096 caps an uncapped request at that floor upstream + + + + Refusal: 429, first scope over + "Rate limit exceeded for api_key: <hash>. Limit type: requests. + Current limit: 10, Remaining: 0. Limit resets at: ... UTC" + chat: throttling_error + retry-after, rate_limit_type, reset_at + messages: Anthropic rate_limit_error envelope, no retry headers + a refused token reservation gives its slot back + + + + Admitted: request runs + complete responses carry where each limit stands: + x-ratelimit-{scope}-remaining-{type}, -limit-{type} + every response: x-litellm-key-rpm-limit / -tpm-limit + streams get no x-ratelimit-* headers, as in the proxy + + + + Settle (off the request path) + tokens: INCRBYFLOAT (actual - reserved) on each reserved scope, actual = all tokens but cache reads; a failure returns the reservation + in flight: ZREM the request's slot when the response ends, fails, or the client goes away + + + + Same keys, scripts and windows as the Python proxy, under its namespace + ns:{api_key:<hash>}:window | requests | tokens | max_parallel_requests: gateway replicas and Python proxies enforce one limit. + No Redis, or Redis unreachable: the same logic on this replica's counters (fail_closed_rate_limit_enforcement: 503 instead) +