mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
feat(gateway): enforce key, user, team, member and end-user rate limits on shared counters (#47050)
* feat(gateway): enforce key, user, team, member and end-user rate limits on shared counters
Adds a built-in rate limits hook with the proxy's v3 limiter semantics:
requests per window, tokens per window and in-flight requests per key,
on the proxy's own {scope:id}:window|requests|tokens|max_parallel_requests
Redis keys and with its scripts, so replicas and Python proxies sharing
the Redis enforce one limit. Tokens are reserved from the proxy's
estimate and settled to usage when the response ends. Refusals carry the
proxy's message, body and headers on each endpoint, and complete
responses carry its x-ratelimit-* headers. Without Redis the counters are
kept per replica. The key store now reads key, user, team and member
rpm/tpm limits and the key's max_parallel_requests, and end users' budget
limits are read with their budgets; keys with plain rate limits are no
longer refused.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* test(gateway-auth): expect the narrowed unsupported rate limit message
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
---------
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Nate Armstrong
GitHub
parent
b83353906b
commit
0919823385
21 changed files with 2293 additions and 152 deletions
|
|
@ -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"
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
|
|
@ -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<f64>,
|
||||
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<f64>,
|
||||
/// 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<EndUser, Error> {
|
||||
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,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -19,6 +19,9 @@ pub struct VerificationToken {
|
|||
pub key_alias: Option<String>,
|
||||
pub key_name: Option<String>,
|
||||
pub metadata: Value,
|
||||
/// The key's `rpm_limit` and `tpm_limit`, else its budget table's.
|
||||
pub limits: RateLimits,
|
||||
pub max_parallel_requests: Option<i64>,
|
||||
pub restrictions: Restrictions,
|
||||
pub team_id: Option<String>,
|
||||
/// `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<f64>,
|
||||
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<i64>,
|
||||
pub tpm: Option<i64>,
|
||||
}
|
||||
|
||||
/// 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<f64>,
|
||||
/// The team's default member budget, `metadata.team_member_budget_id`.
|
||||
pub default_max_budget: Option<f64>,
|
||||
/// `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<String>,
|
||||
|
|
@ -61,6 +74,7 @@ pub struct User {
|
|||
pub spend: f64,
|
||||
pub max_budget: Option<f64>,
|
||||
pub models: Vec<String>,
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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<Map<String, Value>>,
|
||||
pub team_metadata: Arc<Map<String, Value>>,
|
||||
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<i64>,
|
||||
pub tpm: Option<i64>,
|
||||
}
|
||||
|
||||
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<i64>,
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -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<Map<String, Value>>,
|
||||
pub team_metadata: Arc<Map<String, Value>>,
|
||||
pub budgets: KeyBudgets,
|
||||
pub limits: KeyLimits,
|
||||
}
|
||||
|
||||
/// What the gateway knows about a request before it is sent, shared by every hook and observer.
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -121,6 +121,7 @@ fn key_facts(hooks: &Hooks, identity: &AuthenticatedRequest) -> Option<Arc<KeyFa
|
|||
metadata: profile.metadata.clone(),
|
||||
team_metadata: profile.team_metadata.clone(),
|
||||
budgets: profile.budgets.clone(),
|
||||
limits: profile.limits,
|
||||
}))
|
||||
}
|
||||
(AuthenticationMethod::MasterKey, _) => Some(Arc::new(KeyFacts {
|
||||
|
|
|
|||
|
|
@ -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<String>);
|
||||
|
||||
/// End users' budgets, read from the proxy's database and reused for a minute.
|
||||
struct EndUsers {
|
||||
pool: PgPool,
|
||||
cache: Mutex<HashMap<EndUserKey, (EndUser, Instant)>>,
|
||||
}
|
||||
|
||||
impl EndUsers {
|
||||
async fn get(&self, id: &str, default_budget_id: Option<&str>) -> Option<EndUser> {
|
||||
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();
|
||||
|
|
|
|||
70
litellm-rust/crates/gateway/src/end_users.rs
Normal file
70
litellm-rust/crates/gateway/src/end_users.rs
Normal file
|
|
@ -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<String>, Option<String>);
|
||||
|
||||
pub(crate) struct EndUsers {
|
||||
pool: PgPool,
|
||||
cache: Mutex<HashMap<EndUserKey, (EndUser, Instant)>>,
|
||||
}
|
||||
|
||||
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<EndUser> {
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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<Refusal> {
|
||||
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<Restriction> {
|
||||
|
|
@ -195,6 +221,14 @@ fn metadata_restriction(metadata: &Value) -> Option<Restriction> {
|
|||
.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))
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -120,6 +120,13 @@ async fn main() -> Result<(), Box<dyn Error>> {
|
|||
coordination.clone(),
|
||||
database.clone(),
|
||||
)));
|
||||
hooks.push(Arc::new(
|
||||
litellm_gateway::rate_limits::RateLimiter::for_gateway(
|
||||
&config,
|
||||
coordination.clone(),
|
||||
database.clone(),
|
||||
),
|
||||
));
|
||||
}
|
||||
let mut observers: Vec<Arc<dyn litellm_gateway_hooks::Observer>> = Vec::new();
|
||||
let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config));
|
||||
|
|
|
|||
1193
litellm-rust/crates/gateway/src/rate_limits.rs
Normal file
1193
litellm-rust/crates/gateway/src/rate_limits.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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);
|
||||
|
|
|
|||
515
litellm-rust/crates/gateway/tests/rate_limits.rs
Normal file
515
litellm-rust/crates/gateway/tests/rate_limits.rs
Normal file
|
|
@ -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<AtomicUsize>,
|
||||
bodies: Arc<Mutex<Vec<Value>>>,
|
||||
}
|
||||
|
||||
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<Body>| {
|
||||
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<KeyHash, KeyStatus>);
|
||||
|
||||
impl KeyLookup for Keys {
|
||||
fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> {
|
||||
Box::pin(async move { Ok(self.0.get(hash).cloned()) })
|
||||
}
|
||||
}
|
||||
|
||||
fn 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<Arc<Coordination>>,
|
||||
database: Option<litellm_db::PgPool>,
|
||||
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<dyn KeyLookup> = 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<String> {
|
||||
std::env::var("LITELLM_TEST_REDIS_URL").ok()
|
||||
}
|
||||
|
||||
async fn redis(url: &str, namespace: &str) -> Arc<Coordination> {
|
||||
let coordination =
|
||||
Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap();
|
||||
let router: Arc<dyn litellm_gateway::coordination::RouterHandle> = 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::<Option<String>>(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."
|
||||
);
|
||||
}
|
||||
55
litellm-rust/docs/customer-proxy/rate-limits.svg
Normal file
55
litellm-rust/docs/customer-proxy/rate-limits.svg
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 620" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
|
||||
<defs>
|
||||
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
|
||||
<marker id="r" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#dc2626"/></marker>
|
||||
</defs>
|
||||
<rect width="1000" height="620" fill="#ffffff"/>
|
||||
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">Rate limits hook: the proxy's limiter on the proxy's shared counters</text>
|
||||
|
||||
<rect x="24" y="60" width="300" height="150" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
|
||||
<text x="38" y="84" font-weight="600" fill="#1e3a8a">Scopes, in the proxy's order</text>
|
||||
<text x="38" y="106" fill="#1e293b" font-size="13">api_key: key rpm/tpm (else budget row's),</text>
|
||||
<text x="38" y="124" fill="#1e293b" font-size="13"> max_parallel_requests</text>
|
||||
<text x="38" y="142" fill="#1e293b" font-size="13">user, team, team_member (team:user)</text>
|
||||
<text x="38" y="160" fill="#1e293b" font-size="13">end_user: its budget row, or the default</text>
|
||||
<text x="38" y="182" fill="#64748b" font-size="13">read with the key; end users cached 60 s</text>
|
||||
<text x="38" y="200" fill="#64748b" font-size="13">no limits set: no Redis call at all</text>
|
||||
|
||||
<rect x="360" y="60" width="616" height="150" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
|
||||
<text x="374" y="84" font-weight="600" fill="#92400e">Three Redis scripts, in the proxy's phases</text>
|
||||
<text x="374" y="106" fill="#1e293b" font-size="13">1. requests: +1 on every {scope:id}:requests in one call, fixed window from the first hit;</text>
|
||||
<text x="374" y="124" fill="#1e293b" font-size="13"> over when count > limit; the count is never given back</text>
|
||||
<text x="374" y="142" fill="#1e293b" font-size="13">2. in flight (key only): ZSET of slots on Redis TIME, pruned at 3600 s, full at limit</text>
|
||||
<text x="374" y="160" fill="#1e293b" font-size="13">3. tokens: reserve chars/4 + output x candidates on every scope, all or none;</text>
|
||||
<text x="374" y="178" fill="#1e293b" font-size="13"> output = max_tokens, else the configured estimate, else max(input, min(1024, tpm/4));</text>
|
||||
<text x="374" y="196" fill="#1e293b" font-size="13"> a limit under 4096 caps an uncapped request at that floor upstream</text>
|
||||
<line x1="324" y1="135" x2="356" y2="135" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<rect x="24" y="246" width="460" height="130" rx="8" fill="#fef2f2" stroke="#dc2626"/>
|
||||
<text x="38" y="270" font-weight="600" fill="#991b1b">Refusal: 429, first scope over</text>
|
||||
<text x="38" y="292" fill="#1e293b" font-size="13">"Rate limit exceeded for api_key: <hash>. Limit type: requests.</text>
|
||||
<text x="38" y="310" fill="#1e293b" font-size="13"> Current limit: 10, Remaining: 0. Limit resets at: ... UTC"</text>
|
||||
<text x="38" y="332" fill="#1e293b" font-size="13">chat: throttling_error + retry-after, rate_limit_type, reset_at</text>
|
||||
<text x="38" y="350" fill="#1e293b" font-size="13">messages: Anthropic rate_limit_error envelope, no retry headers</text>
|
||||
<text x="38" y="368" fill="#64748b" font-size="13">a refused token reservation gives its slot back</text>
|
||||
<path d="M500 210 L500 230 L254 230 L254 242" fill="none" stroke="#dc2626" stroke-width="1.5" marker-end="url(#r)"/>
|
||||
|
||||
<rect x="516" y="246" width="460" height="130" rx="8" fill="#f8fafc" stroke="#64748b"/>
|
||||
<text x="530" y="270" font-weight="600" fill="#0f172a">Admitted: request runs</text>
|
||||
<text x="530" y="292" fill="#334155" font-size="13">complete responses carry where each limit stands:</text>
|
||||
<text x="530" y="310" fill="#334155" font-size="13">x-ratelimit-{scope}-remaining-{type}, -limit-{type}</text>
|
||||
<text x="530" y="328" fill="#334155" font-size="13">every response: x-litellm-key-rpm-limit / -tpm-limit</text>
|
||||
<text x="530" y="350" fill="#64748b" font-size="13">streams get no x-ratelimit-* headers, as in the proxy</text>
|
||||
<path d="M760 210 L760 242" fill="none" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<rect x="24" y="412" width="952" height="80" rx="8" fill="#f5f3ff" stroke="#8b5cf6"/>
|
||||
<text x="38" y="436" font-weight="600" fill="#5b21b6">Settle (off the request path)</text>
|
||||
<text x="38" y="458" fill="#1e293b" font-size="13">tokens: INCRBYFLOAT (actual - reserved) on each reserved scope, actual = all tokens but cache reads; a failure returns the reservation</text>
|
||||
<text x="38" y="476" fill="#1e293b" font-size="13">in flight: ZREM the request's slot when the response ends, fails, or the client goes away</text>
|
||||
<line x1="746" y1="376" x2="746" y2="408" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<rect x="24" y="522" width="952" height="76" rx="8" fill="#ecfdf5" stroke="#10b981"/>
|
||||
<text x="38" y="546" font-weight="600" fill="#065f46">Same keys, scripts and windows as the Python proxy, under its namespace</text>
|
||||
<text x="38" y="568" fill="#1e293b" font-size="13">ns:{api_key:<hash>}:window | requests | tokens | max_parallel_requests: gateway replicas and Python proxies enforce one limit.</text>
|
||||
<text x="38" y="586" fill="#64748b" font-size="13">No Redis, or Redis unreachable: the same logic on this replica's counters (fail_closed_rate_limit_enforcement: 503 instead)</text>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 5.3 KiB |
Loading…
Add table
Reference in a new issue