mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
feat(gateway): enforce key, team, user, member and end-user budgets on shared spend counters (#47044)
Adds a built-in budgets hook that runs the proxy's budget checks, in its order and with its messages, on the proxy's Redis spend:* counters. Each request reads and seeds the counters, checks them and reserves its maximum cost in one Redis script, then settles the reservation to its real cost off the request path, so replicas and Python proxies sharing the Redis enforce one budget. Without Redis the counters are kept per replica. The key store now reads user, team member and end-user budgets, and auth no longer checks budgets against stored spend. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Nate Armstrong
GitHub
parent
2d2ca95481
commit
b83353906b
22 changed files with 2159 additions and 208 deletions
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -3820,6 +3820,7 @@ dependencies = [
|
|||
"futures-util",
|
||||
"http 1.4.2",
|
||||
"litellm-cost",
|
||||
"litellm-gateway-auth",
|
||||
"litellm-model-catalog",
|
||||
"memchr",
|
||||
"rstest",
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1,60 @@
|
|||
{
|
||||
"db_name": "PostgreSQL",
|
||||
"query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n (SELECT d.max_budget FROM \"LiteLLM_BudgetTable\" AS d\n WHERE d.budget_id = $2) AS \"default_max_budget?\"\n FROM (SELECT 1) AS one\n LEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\n LEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = e.budget_id",
|
||||
"describe": {
|
||||
"columns": [
|
||||
{
|
||||
"ordinal": 0,
|
||||
"name": "spend?",
|
||||
"type_info": "Float8",
|
||||
"origin": {
|
||||
"Table": {
|
||||
"table": "\"LiteLLM_EndUserTable\"",
|
||||
"name": "spend"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 1,
|
||||
"name": "blocked?",
|
||||
"type_info": "Bool",
|
||||
"origin": {
|
||||
"Table": {
|
||||
"table": "\"LiteLLM_EndUserTable\"",
|
||||
"name": "blocked"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 2,
|
||||
"name": "max_budget?",
|
||||
"type_info": "Float8",
|
||||
"origin": {
|
||||
"Table": {
|
||||
"table": "\"LiteLLM_BudgetTable\"",
|
||||
"name": "max_budget"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"ordinal": 3,
|
||||
"name": "default_max_budget?",
|
||||
"type_info": "Float8",
|
||||
"origin": "Expression"
|
||||
}
|
||||
],
|
||||
"parameters": {
|
||||
"Left": [
|
||||
"Text",
|
||||
"Text"
|
||||
]
|
||||
},
|
||||
"nullable": [
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
null
|
||||
]
|
||||
},
|
||||
"hash": "a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe"
|
||||
}
|
||||
40
litellm-rust/crates/db/src/end_user.rs
Normal file
40
litellm-rust/crates/db/src/end_user.rs
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
use sqlx::PgPool;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
/// An end user's spend and budget, as the proxy's end-user budget checks read them.
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct EndUser {
|
||||
/// `LiteLLM_EndUserTable.spend`, when the end user has a row.
|
||||
pub spend: Option<f64>,
|
||||
pub blocked: bool,
|
||||
/// The row's budget, else the default budget's.
|
||||
pub max_budget: Option<f64>,
|
||||
}
|
||||
|
||||
/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own.
|
||||
pub async fn find_end_user(
|
||||
pool: &PgPool,
|
||||
end_user_id: &str,
|
||||
default_budget_id: Option<&str>,
|
||||
) -> Result<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?"
|
||||
FROM (SELECT 1) AS one
|
||||
LEFT JOIN "LiteLLM_EndUserTable" AS e ON e.user_id = $1
|
||||
LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = e.budget_id"#,
|
||||
end_user_id,
|
||||
default_budget_id,
|
||||
)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
Ok(EndUser {
|
||||
spend: row.spend,
|
||||
blocked: row.blocked.unwrap_or(false),
|
||||
max_budget: row.max_budget.or(row.default_max_budget),
|
||||
})
|
||||
}
|
||||
|
|
@ -1,9 +1,11 @@
|
|||
mod end_user;
|
||||
mod error;
|
||||
mod failure;
|
||||
mod pool;
|
||||
mod spend;
|
||||
mod verification_token;
|
||||
|
||||
pub use end_user::{EndUser, find_end_user};
|
||||
pub use error::Error;
|
||||
pub use failure::DbFailure;
|
||||
pub use pool::{PoolSettings, connect_lazy};
|
||||
|
|
|
|||
|
|
@ -41,9 +41,15 @@ pub struct Team {
|
|||
pub restrictions: Restrictions,
|
||||
}
|
||||
|
||||
/// What a team membership's budget table sets beyond its dollar budget.
|
||||
#[derive(Clone, Debug, Default, Eq, PartialEq)]
|
||||
/// The key's team membership: its spend and budget, and what its budget table sets beyond them.
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct TeamMembership {
|
||||
/// The membership row's spend, when the key's user has a row for its team.
|
||||
pub spend: Option<f64>,
|
||||
/// The membership's own budget, with any temporary increase still in effect.
|
||||
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`.
|
||||
pub rate_limits: bool,
|
||||
/// The member's own model scope; empty inherits the team's.
|
||||
|
|
@ -52,6 +58,8 @@ pub struct TeamMembership {
|
|||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct User {
|
||||
pub spend: f64,
|
||||
pub max_budget: Option<f64>,
|
||||
pub models: Vec<String>,
|
||||
pub restrictions: Restrictions,
|
||||
}
|
||||
|
|
@ -85,7 +93,7 @@ pub async fn find_verification_token(
|
|||
v.blocked,
|
||||
v.models,
|
||||
v.spend,
|
||||
v.max_budget,
|
||||
COALESCE(v.max_budget, kb.max_budget) AS max_budget,
|
||||
v.allowed_routes,
|
||||
v.key_alias,
|
||||
v.key_name,
|
||||
|
|
@ -131,9 +139,16 @@ pub async fn find_verification_token(
|
|||
'$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))
|
||||
AS "member_rate_limits!",
|
||||
mb.allowed_models AS "member_allowed_models?",
|
||||
tm.spend AS "member_spend?",
|
||||
mb.max_budget AS "member_max_budget?",
|
||||
CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC')
|
||||
THEN mb.temp_budget_increase END AS "member_budget_increase?",
|
||||
tdb.max_budget AS "member_default_budget?",
|
||||
v.user_id,
|
||||
u.user_id AS "user_row_id?",
|
||||
u.models AS "user_models?",
|
||||
u.spend AS "user_spend?",
|
||||
u.max_budget AS "user_max_budget?",
|
||||
(u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL
|
||||
OR u.max_parallel_requests IS NOT NULL
|
||||
OR COALESCE(jsonb_path_exists(u.model_max_budget,
|
||||
|
|
@ -147,6 +162,8 @@ pub async fn find_verification_token(
|
|||
LEFT JOIN "LiteLLM_TeamMembership" AS tm
|
||||
ON v.team_id = tm.team_id AND v.user_id = tm.user_id
|
||||
LEFT JOIN "LiteLLM_BudgetTable" AS mb ON tm.budget_id = mb.budget_id
|
||||
LEFT JOIN "LiteLLM_BudgetTable" AS tdb
|
||||
ON tdb.budget_id = t.metadata ->> 'team_member_budget_id'
|
||||
LEFT JOIN "LiteLLM_UserTable" AS u ON v.user_id = u.user_id
|
||||
WHERE v.token = $1"#,
|
||||
token
|
||||
|
|
@ -192,11 +209,18 @@ pub async fn find_verification_token(
|
|||
},
|
||||
}),
|
||||
team_membership: TeamMembership {
|
||||
spend: row.member_spend,
|
||||
max_budget: row
|
||||
.member_max_budget
|
||||
.map(|budget| budget + row.member_budget_increase.unwrap_or(0.0)),
|
||||
default_max_budget: row.member_default_budget,
|
||||
rate_limits: row.member_rate_limits,
|
||||
allowed_models: row.member_allowed_models.unwrap_or_default(),
|
||||
},
|
||||
user_id: row.user_id,
|
||||
user: row.user_row_id.map(|_| User {
|
||||
spend: row.user_spend.unwrap_or(0.0),
|
||||
max_budget: row.user_max_budget,
|
||||
models: row.user_models.unwrap_or_default(),
|
||||
restrictions: Restrictions {
|
||||
rate_limits: row.user_rate_limits,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use litellm_db::{
|
||||
Restrictions, Team, TeamMembership, User, VerificationToken, find_verification_token,
|
||||
EndUser, Restrictions, Team, TeamMembership, User, VerificationToken, find_end_user,
|
||||
find_verification_token,
|
||||
};
|
||||
use litellm_db_testing::MigratedPostgres;
|
||||
use rstest::rstest;
|
||||
|
|
@ -21,9 +22,12 @@ INSERT INTO "LiteLLM_BudgetTable" (budget_id, max_budget, rpm_limit, model_max_b
|
|||
VALUES ('b-1', 5.0, NULL, NULL, ARRAY['gpt-a'], 'x', 'x'),
|
||||
('b-2', NULL, 3, NULL, ARRAY[]::text[], 'x', 'x'),
|
||||
('b-3', NULL, NULL, '{"gpt-a": {"tpm_limit": 7}}', ARRAY[]::text[], 'x', 'x');
|
||||
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, budget_id)
|
||||
VALUES ('user-2', 'team-2', 'b-1'),
|
||||
('user-1', 'team-1', 'b-2');
|
||||
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, budget_id, spend)
|
||||
VALUES ('user-2', 'team-2', 'b-1', 0.75),
|
||||
('user-1', 'team-1', 'b-2', 0);
|
||||
INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id, blocked)
|
||||
VALUES ('end-1', 1.5, 'b-1', false),
|
||||
('end-2', 0.25, NULL, true);
|
||||
INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, created_by, updated_by)
|
||||
VALUES ('org-1', 'org', 'b-1', 'x', 'x');
|
||||
INSERT INTO "LiteLLM_ProjectTable" (project_id, created_by, updated_by) VALUES ('project-1', 'x', 'x');
|
||||
|
|
@ -58,19 +62,20 @@ INSERT INTO "LiteLLM_VerificationToken"
|
|||
'{"gpt-a": {"budget_limit": 1.0, "rpm_limit": 2}}');
|
||||
"#;
|
||||
|
||||
#[expect(
|
||||
clippy::disallowed_methods,
|
||||
reason = "test fixture rows; the lookups under test use the checked macro"
|
||||
)]
|
||||
async fn setup(pool: &sqlx::PgPool) {
|
||||
sqlx::raw_sql(FIXTURES).execute(pool).await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() {
|
||||
let postgres = MigratedPostgres::start().await.unwrap();
|
||||
#[expect(
|
||||
clippy::disallowed_methods,
|
||||
reason = "test fixture rows; the lookup under test uses the checked macro"
|
||||
)]
|
||||
sqlx::raw_sql(FIXTURES)
|
||||
.execute(postgres.pool())
|
||||
.await
|
||||
.unwrap();
|
||||
let pool = postgres.pool();
|
||||
setup(pool).await;
|
||||
let find = |token: &'static str| async move {
|
||||
find_verification_token(pool, token).await.unwrap().unwrap()
|
||||
};
|
||||
|
|
@ -125,11 +130,16 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() {
|
|||
},
|
||||
}),
|
||||
team_membership: TeamMembership {
|
||||
spend: Some(0.0),
|
||||
max_budget: None,
|
||||
default_max_budget: None,
|
||||
rate_limits: true,
|
||||
allowed_models: vec![],
|
||||
},
|
||||
user_id: Some("user-1".into()),
|
||||
user: Some(User {
|
||||
spend: 0.0,
|
||||
max_budget: Some(1.0),
|
||||
models: vec!["gpt-b".into()],
|
||||
restrictions: Restrictions {
|
||||
rate_limits: true,
|
||||
|
|
@ -175,10 +185,16 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() {
|
|||
assert_eq!(
|
||||
member.team_membership,
|
||||
TeamMembership {
|
||||
spend: Some(0.75),
|
||||
max_budget: Some(5.0),
|
||||
default_max_budget: None,
|
||||
rate_limits: false,
|
||||
allowed_models: vec!["gpt-a".into()],
|
||||
}
|
||||
);
|
||||
// A key without its own budget takes its budget row's.
|
||||
assert_eq!(find("restricted").await.max_budget, Some(5.0));
|
||||
assert_eq!(find("budgeted").await.max_budget, Some(4.0));
|
||||
assert_eq!(member.team.unwrap().restrictions, Restrictions::default());
|
||||
assert_eq!(member.user.unwrap().restrictions, Restrictions::default());
|
||||
let orphan = find("orphan").await;
|
||||
|
|
@ -193,3 +209,39 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() {
|
|||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn reads_end_user_budgets_with_their_default() {
|
||||
let postgres = MigratedPostgres::start().await.unwrap();
|
||||
let pool = postgres.pool();
|
||||
setup(pool).await;
|
||||
let find =
|
||||
async |id: &str, default: Option<&str>| find_end_user(pool, id, default).await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
find("end-1", Some("b-2")).await,
|
||||
EndUser {
|
||||
spend: Some(1.5),
|
||||
blocked: false,
|
||||
max_budget: Some(5.0)
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
find("end-2", Some("b-1")).await,
|
||||
EndUser {
|
||||
spend: Some(0.25),
|
||||
blocked: true,
|
||||
max_budget: Some(5.0)
|
||||
}
|
||||
);
|
||||
assert_eq!(find("end-2", None).await.max_budget, None);
|
||||
assert_eq!(
|
||||
find("unknown", Some("b-1")).await,
|
||||
EndUser {
|
||||
spend: None,
|
||||
blocked: false,
|
||||
max_budget: Some(5.0)
|
||||
}
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -297,8 +297,46 @@ pub struct KeyProfile {
|
|||
pub team_alias: Option<String>,
|
||||
pub metadata: Arc<Map<String, Value>>,
|
||||
pub team_metadata: Arc<Map<String, Value>>,
|
||||
pub budgets: KeyBudgets,
|
||||
}
|
||||
|
||||
/// Stored spend against an optional budget, as the key store last read it.
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq)]
|
||||
pub struct Spend {
|
||||
pub spend: f64,
|
||||
pub max_budget: Option<f64>,
|
||||
}
|
||||
|
||||
impl Eq for Spend {}
|
||||
|
||||
/// The spend and budgets of a key and of who it belongs to, read with the key and refreshed
|
||||
/// with it.
|
||||
#[derive(Clone, Debug, Default, Eq, PartialEq)]
|
||||
pub struct KeyBudgets {
|
||||
/// The key's `max_budget`, else its budget table's.
|
||||
pub key: Spend,
|
||||
/// `key_name`, the key's masked form, which budget errors show while it still looks masked.
|
||||
pub key_name: Option<String>,
|
||||
/// The key's team, when its row exists.
|
||||
pub team: Option<Spend>,
|
||||
/// The key's user, when its row exists.
|
||||
pub user: Option<Spend>,
|
||||
/// The key's user within its team, when a membership row exists.
|
||||
pub team_member: Option<MemberSpend>,
|
||||
/// `metadata.end_user_budget_id`: the budget of end users without their own.
|
||||
pub end_user_budget_id: Option<String>,
|
||||
}
|
||||
|
||||
/// A team member's spend, with the membership's own budget and the team's default one.
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq)]
|
||||
pub struct MemberSpend {
|
||||
pub spend: f64,
|
||||
pub max_budget: Option<f64>,
|
||||
pub default_max_budget: Option<f64>,
|
||||
}
|
||||
|
||||
impl Eq for MemberSpend {}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum KeyStatus {
|
||||
Active {
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ bytes.workspace = true
|
|||
futures-util.workspace = true
|
||||
http.workspace = true
|
||||
litellm-cost.workspace = true
|
||||
litellm-gateway-auth.workspace = true
|
||||
litellm-model-catalog.workspace = true
|
||||
memchr = "2"
|
||||
serde.workspace = true
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ use bytes::Bytes;
|
|||
use http::HeaderMap;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use litellm_gateway_auth::keys::KeyBudgets;
|
||||
|
||||
use crate::usage::Usage;
|
||||
|
||||
/// The facts a hook reads or an observer asks for. The gateway builds only what some hook or
|
||||
|
|
@ -101,6 +103,7 @@ pub struct KeyFacts {
|
|||
pub org_id: Option<String>,
|
||||
pub metadata: Arc<Map<String, Value>>,
|
||||
pub team_metadata: Arc<Map<String, Value>>,
|
||||
pub budgets: KeyBudgets,
|
||||
}
|
||||
|
||||
/// What the gateway knows about a request before it is sent, shared by every hook and observer.
|
||||
|
|
|
|||
|
|
@ -12,7 +12,10 @@ use futures_util::future::BoxFuture;
|
|||
use http::{HeaderName, HeaderValue};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::facts::{Annotations, Facts, Record, RequestFacts};
|
||||
use crate::{
|
||||
facts::{Annotations, Facts, Record, RequestFacts},
|
||||
pricing::Prices,
|
||||
};
|
||||
|
||||
/// The parts of a request a hook may change.
|
||||
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
|
||||
|
|
@ -93,6 +96,8 @@ pub struct Refusal {
|
|||
pub status: u16,
|
||||
/// The proxy's error `type`, such as `budget_exceeded`.
|
||||
pub kind: Cow<'static, str>,
|
||||
/// The error's `param`: the string `"None"` for most proxy errors, `null` for some.
|
||||
pub param: Option<&'static str>,
|
||||
pub message: String,
|
||||
pub headers: Vec<(HeaderName, HeaderValue)>,
|
||||
}
|
||||
|
|
@ -106,11 +111,18 @@ impl Refusal {
|
|||
Self {
|
||||
status,
|
||||
kind: kind.into(),
|
||||
param: Some("None"),
|
||||
message: message.into(),
|
||||
headers: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends a `null` `param`, as the proxy's budget errors do.
|
||||
pub fn without_param(mut self) -> Self {
|
||||
self.param = None;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_header(mut self, name: HeaderName, value: HeaderValue) -> Self {
|
||||
self.headers.push((name, value));
|
||||
self
|
||||
|
|
@ -229,6 +241,9 @@ pub trait Hook: Send + Sync + 'static {
|
|||
/// The parts of the request this hook may change.
|
||||
fn writes(&self) -> Targets;
|
||||
|
||||
/// Called once with the deployments' prices before the first request.
|
||||
fn prepare(&self, _prices: &Arc<Prices>) {}
|
||||
|
||||
/// Runs before the request is routed. Returning a [`Settle`] asks to be told how the request
|
||||
/// ended.
|
||||
fn run<'a>(
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ impl Pipeline {
|
|||
let mut reads = observers.reads();
|
||||
for hook in &hooks {
|
||||
reads |= hook.reads();
|
||||
hook.prepare(&prices);
|
||||
}
|
||||
// Cost is computed from usage and the deployment that served the request.
|
||||
if reads.contains(Facts::COST) {
|
||||
|
|
|
|||
|
|
@ -17,12 +17,16 @@ use crate::usage::{ServiceTier, Usage};
|
|||
#[derive(Debug, Default)]
|
||||
pub struct Prices {
|
||||
plans: HashMap<Arc<str>, PricingPlan>,
|
||||
max_output_tokens: HashMap<Arc<str>, u64>,
|
||||
}
|
||||
|
||||
impl Prices {
|
||||
/// Adds a deployment's prices from its cost map fields. A deployment whose fields hold no
|
||||
/// usable prices is left out, and its cost is unknown.
|
||||
pub fn insert(&mut self, id: Arc<str>, fields: &Map<String, Value>) {
|
||||
if let Some(limit) = fields.get("max_output_tokens").and_then(Value::as_u64) {
|
||||
self.max_output_tokens.insert(id.clone(), limit);
|
||||
}
|
||||
let parsed = ParsedRates::parse(fields);
|
||||
if parsed.standard == Rates::EMPTY && parsed.thresholds.is_empty() {
|
||||
return;
|
||||
|
|
@ -43,6 +47,11 @@ impl Prices {
|
|||
self.plans.is_empty()
|
||||
}
|
||||
|
||||
/// The most output tokens the deployment's model can return, when the cost map says.
|
||||
pub fn max_output_tokens(&self, deployment: &str) -> Option<u64> {
|
||||
self.max_output_tokens.get(deployment).copied()
|
||||
}
|
||||
|
||||
/// The cost in dollars of `usage` on a deployment, or `None` when its prices are unknown.
|
||||
pub fn cost(&self, deployment: &str, usage: &Usage) -> Option<f64> {
|
||||
let plan = self.plans.get(deployment)?;
|
||||
|
|
|
|||
|
|
@ -180,7 +180,7 @@ fn refused(refusal: Refusal) -> Response {
|
|||
let body = json!({"error": {
|
||||
"message": refusal.message,
|
||||
"type": refusal.kind,
|
||||
"param": "None",
|
||||
"param": refusal.param,
|
||||
"code": status.as_u16().to_string(),
|
||||
}});
|
||||
let mut response = (status, Json(body)).into_response();
|
||||
|
|
|
|||
|
|
@ -120,6 +120,7 @@ fn key_facts(hooks: &Hooks, identity: &AuthenticatedRequest) -> Option<Arc<KeyFa
|
|||
org_id: None,
|
||||
metadata: profile.metadata.clone(),
|
||||
team_metadata: profile.team_metadata.clone(),
|
||||
budgets: profile.budgets.clone(),
|
||||
}))
|
||||
}
|
||||
(AuthenticationMethod::MasterKey, _) => Some(Arc::new(KeyFacts {
|
||||
|
|
|
|||
1209
litellm-rust/crates/gateway/src/budgets.rs
Normal file
1209
litellm-rust/crates/gateway/src/budgets.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -267,10 +267,16 @@ impl Coordination {
|
|||
}))
|
||||
}
|
||||
|
||||
fn connection(&self) -> Option<ConnectionManager> {
|
||||
/// The shared connection, once Redis has answered.
|
||||
pub(crate) fn connection(&self) -> Option<ConnectionManager> {
|
||||
self.connection.get().cloned()
|
||||
}
|
||||
|
||||
/// `key` under the configured namespace, as the proxy prefixes its keys.
|
||||
pub(crate) fn namespaced(&self, key: &str) -> String {
|
||||
self.settings.key(key)
|
||||
}
|
||||
|
||||
/// Whether the connection shared state is written and read through is up yet.
|
||||
pub fn connected(&self) -> bool {
|
||||
self.connection.get().is_some()
|
||||
|
|
|
|||
|
|
@ -2,7 +2,8 @@ use litellm_db::{Restrictions, VerificationToken};
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_gateway_auth::keys::{
|
||||
AllowedRoutes, KeyModels, KeyProfile, KeyStatus, ModelAccess, Owner, Refusal, Restriction,
|
||||
AllowedRoutes, KeyBudgets, KeyModels, KeyProfile, KeyStatus, MemberSpend, ModelAccess, Owner,
|
||||
Refusal, Restriction, Spend,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -37,6 +38,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K
|
|||
Value::Object(fields) => Arc::new(fields),
|
||||
_ => Arc::default(),
|
||||
};
|
||||
let budgets = budgets(&row);
|
||||
let (team_models, team_alias, team_metadata) = match row.team {
|
||||
Some(team) => (
|
||||
Some(ModelAccess::new(team.models)),
|
||||
|
|
@ -59,10 +61,41 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K
|
|||
team_alias,
|
||||
metadata: object(row.metadata),
|
||||
team_metadata,
|
||||
budgets,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// The spend and budgets the budgets hook checks for this key.
|
||||
fn budgets(row: &VerificationToken) -> KeyBudgets {
|
||||
let membership = &row.team_membership;
|
||||
KeyBudgets {
|
||||
key: Spend {
|
||||
spend: row.spend,
|
||||
max_budget: row.max_budget,
|
||||
},
|
||||
key_name: row.key_name.clone(),
|
||||
team: row.team.as_ref().map(|team| Spend {
|
||||
spend: team.spend,
|
||||
max_budget: team.max_budget,
|
||||
}),
|
||||
user: row.user.as_ref().map(|user| Spend {
|
||||
spend: user.spend,
|
||||
max_budget: user.max_budget,
|
||||
}),
|
||||
team_member: membership.spend.map(|spend| MemberSpend {
|
||||
spend,
|
||||
max_budget: membership.max_budget,
|
||||
default_max_budget: membership.default_max_budget,
|
||||
}),
|
||||
end_user_budget_id: row
|
||||
.metadata
|
||||
.get("end_user_budget_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned),
|
||||
}
|
||||
}
|
||||
|
||||
fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option<Refusal> {
|
||||
let team = match (&row.team_id, &row.team) {
|
||||
(Some(team_id), None) => {
|
||||
|
|
@ -80,29 +113,6 @@ fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option<Ref
|
|||
team_id: team_id.clone(),
|
||||
});
|
||||
}
|
||||
// Budgets pass: the gateway does not record spend for its own traffic, so it refuses only a
|
||||
// key or team whose stored spend has already reached its budget.
|
||||
if let Some(max_budget) = row.max_budget
|
||||
&& max_budget.is_finite()
|
||||
&& row.spend >= max_budget
|
||||
{
|
||||
return Some(budget_exceeded(
|
||||
format!("Key={}", key_label(row)),
|
||||
row.spend,
|
||||
max_budget,
|
||||
));
|
||||
}
|
||||
if let Some((team_id, team)) = team
|
||||
&& let Some(max_budget) = team.max_budget
|
||||
&& max_budget.is_finite()
|
||||
&& team.spend > max_budget
|
||||
{
|
||||
return Some(budget_exceeded(
|
||||
format!("Team={team_id}"),
|
||||
team.spend,
|
||||
max_budget,
|
||||
));
|
||||
}
|
||||
unsupported(row, personal_keys_use_teams)
|
||||
.map(|(restriction, owner)| Refusal::Unsupported { restriction, owner })
|
||||
}
|
||||
|
|
@ -197,57 +207,6 @@ fn is_set(value: &Value) -> bool {
|
|||
}
|
||||
}
|
||||
|
||||
fn budget_exceeded(label: String, spend: f64, max_budget: f64) -> Refusal {
|
||||
Refusal::BudgetExceeded {
|
||||
owner: label,
|
||||
spend: python_float(spend),
|
||||
max_budget: python_float(max_budget),
|
||||
}
|
||||
}
|
||||
|
||||
/// The proxy names a key by its alias, adding `key_name` only while it still looks masked.
|
||||
fn key_label(row: &VerificationToken) -> String {
|
||||
let alias = row
|
||||
.key_alias
|
||||
.as_deref()
|
||||
.filter(|alias| !alias.is_empty())
|
||||
.unwrap_or("key");
|
||||
match row.key_name.as_deref().filter(|name| is_masked(name)) {
|
||||
Some(name) => format!("{alias} ({name})"),
|
||||
None => alias.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The proxy's `_MASKED_KEY_NAME_RE`, `sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?`.
|
||||
fn is_masked(name: &str) -> bool {
|
||||
name.strip_prefix("sk-...").is_some_and(|rest| {
|
||||
rest.is_empty()
|
||||
|| (rest.chars().count() == 4
|
||||
&& rest
|
||||
.chars()
|
||||
.all(|character| !character.is_whitespace() && !character.is_control()))
|
||||
})
|
||||
}
|
||||
|
||||
/// Python's `repr` of a float.
|
||||
fn python_float(value: f64) -> String {
|
||||
if value.is_nan() {
|
||||
return "nan".into();
|
||||
}
|
||||
if value.is_infinite() {
|
||||
return if value > 0.0 { "inf" } else { "-inf" }.into();
|
||||
}
|
||||
let magnitude = value.abs();
|
||||
if magnitude != 0.0 && !(1e-4..1e16).contains(&magnitude) {
|
||||
let scientific = format!("{value:e}");
|
||||
let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0"));
|
||||
let exponent: i32 = exponent.parse().unwrap_or(0);
|
||||
let sign = if exponent < 0 { '-' } else { '+' };
|
||||
return format!("{mantissa}e{sign}{:02}", exponent.abs());
|
||||
}
|
||||
format!("{value:?}")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::SystemTime;
|
||||
|
|
@ -307,7 +266,13 @@ mod tests {
|
|||
},
|
||||
allowed_routes: AllowedRoutes::new(strings(&["llm_api_routes"])),
|
||||
team_id: team.then(|| "team-1".to_owned()),
|
||||
profile: Default::default(),
|
||||
profile: Arc::new(KeyProfile {
|
||||
budgets: KeyBudgets {
|
||||
team: team.then(Spend::default),
|
||||
..KeyBudgets::default()
|
||||
},
|
||||
..KeyProfile::default()
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -318,14 +283,6 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
fn exceeded(owner: &str, spend: &str, max_budget: &str) -> KeyStatus {
|
||||
refused(Refusal::BudgetExceeded {
|
||||
owner: owner.into(),
|
||||
spend: spend.into(),
|
||||
max_budget: max_budget.into(),
|
||||
})
|
||||
}
|
||||
|
||||
fn unsupported(restriction: Restriction, owner: Owner) -> KeyStatus {
|
||||
refused(Refusal::Unsupported { restriction, owner })
|
||||
}
|
||||
|
|
@ -358,57 +315,6 @@ mod tests {
|
|||
0.0,
|
||||
refused(Refusal::TeamBlocked { team_id: "team-1".into() })
|
||||
)]
|
||||
#[case::key_at_budget(
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
Some(10.0),
|
||||
10.0,
|
||||
exceeded("Key=key", "10.0", "10.0")
|
||||
)]
|
||||
#[case::key_under_budget(false, None, None, Some(10.0), 9.5, active(false))]
|
||||
#[case::key_zero_budget(false, None, None, Some(0.0), 0.0, exceeded("Key=key", "0.0", "0.0"))]
|
||||
#[case::key_infinite_budget(false, None, None, Some(f64::INFINITY), 1e20, active(false))]
|
||||
#[case::team_over_budget(
|
||||
false,
|
||||
Some("team-1"),
|
||||
team(false, Some(5.0), 5.25),
|
||||
None,
|
||||
0.0,
|
||||
exceeded("Team=team-1", "5.25", "5.0")
|
||||
)]
|
||||
#[case::team_at_budget(
|
||||
false,
|
||||
Some("team-1"),
|
||||
team(false, Some(5.0), 5.0),
|
||||
None,
|
||||
0.0,
|
||||
active(true)
|
||||
)]
|
||||
#[case::team_over_budget_under_key_budget(
|
||||
false,
|
||||
Some("team-1"),
|
||||
team(false, Some(5.0), 9.0),
|
||||
Some(10.0),
|
||||
2.0,
|
||||
exceeded("Team=team-1", "9.0", "5.0")
|
||||
)]
|
||||
#[case::key_and_team_under_budget(
|
||||
false,
|
||||
Some("team-1"),
|
||||
team(false, Some(5.0), 1.0),
|
||||
Some(10.0),
|
||||
2.0,
|
||||
active(true)
|
||||
)]
|
||||
#[case::key_budget_checked_before_team(
|
||||
false,
|
||||
Some("team-1"),
|
||||
team(false, Some(5.0), 9.0),
|
||||
Some(1.0),
|
||||
2.0,
|
||||
exceeded("Key=key", "2.0", "1.0")
|
||||
)]
|
||||
fn refuses_like_the_proxy(
|
||||
row: VerificationToken,
|
||||
#[case] blocked: bool,
|
||||
|
|
@ -467,6 +373,8 @@ mod tests {
|
|||
team: team_id.and_then(|_| team(false, None, 0.0)),
|
||||
user_id: Some("user-1".into()),
|
||||
user: Some(User {
|
||||
spend: 0.0,
|
||||
max_budget: None,
|
||||
models: strings(models),
|
||||
restrictions,
|
||||
}),
|
||||
|
|
@ -483,6 +391,7 @@ mod tests {
|
|||
team_membership: TeamMembership {
|
||||
rate_limits,
|
||||
allowed_models: strings(allowed_models),
|
||||
..TeamMembership::default()
|
||||
},
|
||||
..on_team(NONE, json!({}), row)
|
||||
}
|
||||
|
|
@ -691,7 +600,7 @@ mod tests {
|
|||
models,
|
||||
allowed_routes,
|
||||
team_id,
|
||||
..
|
||||
profile,
|
||||
} => KeyStatus::Active {
|
||||
expires_at,
|
||||
models,
|
||||
|
|
@ -699,7 +608,7 @@ mod tests {
|
|||
team_id,
|
||||
profile: Arc::new(KeyProfile {
|
||||
user_id: user_id.map(str::to_owned),
|
||||
..KeyProfile::default()
|
||||
..(*profile).clone()
|
||||
}),
|
||||
},
|
||||
other => other,
|
||||
|
|
@ -714,42 +623,4 @@ mod tests {
|
|||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_alias(None, None, "key")]
|
||||
#[case::empty_alias(Some(""), None, "key")]
|
||||
#[case::alias(Some("prod"), None, "prod")]
|
||||
#[case::masked_name(Some("prod"), Some("sk-...abcd"), "prod (sk-...abcd)")]
|
||||
#[case::bare_masked_name(None, Some("sk-..."), "key (sk-...)")]
|
||||
#[case::unmasked_name(Some("prod"), Some("sk-unmasked90"), "prod")]
|
||||
#[case::name_with_space(Some("prod"), Some("sk-...ab d"), "prod")]
|
||||
#[case::name_too_long(Some("prod"), Some("sk-...abcde"), "prod")]
|
||||
fn labels_keys_like_the_proxy(
|
||||
row: VerificationToken,
|
||||
#[case] alias: Option<&str>,
|
||||
#[case] name: Option<&str>,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let row = VerificationToken {
|
||||
key_alias: alias.map(str::to_owned),
|
||||
key_name: name.map(str::to_owned),
|
||||
..row
|
||||
};
|
||||
assert_eq!(key_label(&row), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::whole(10.0, "10.0")]
|
||||
#[case::fraction(0.1, "0.1")]
|
||||
#[case::zero(0.0, "0.0")]
|
||||
#[case::small(0.0001, "0.0001")]
|
||||
#[case::tiny(0.00001, "1e-05")]
|
||||
#[case::large(1e15, "1000000000000000.0")]
|
||||
#[case::huge(1e16, "1e+16")]
|
||||
#[case::huge_fraction(1.5e17, "1.5e+17")]
|
||||
#[case::negative(-2.5, "-2.5")]
|
||||
#[case::infinite(f64::INFINITY, "inf")]
|
||||
fn prints_floats_like_python(#[case] value: f64, #[case] expected: &str) {
|
||||
assert_eq!(python_float(value), expected);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
mod auth;
|
||||
pub mod budgets;
|
||||
pub mod coordination;
|
||||
pub mod drain;
|
||||
mod error;
|
||||
|
|
|
|||
|
|
@ -112,13 +112,22 @@ async fn main() -> Result<(), Box<dyn Error>> {
|
|||
coordination.start(inference.clone(), keys.clone(), shutdown.clone());
|
||||
}
|
||||
// Built-in hooks and observers are added here as the gateway gains them.
|
||||
let mut hooks: Vec<Arc<dyn litellm_gateway_hooks::Hook>> = Vec::new();
|
||||
if database.is_some() {
|
||||
hooks.push(Arc::new(litellm_gateway::budgets::Budgets::for_gateway(
|
||||
&config,
|
||||
&inference,
|
||||
coordination.clone(),
|
||||
database.clone(),
|
||||
)));
|
||||
}
|
||||
let mut observers: Vec<Arc<dyn litellm_gateway_hooks::Observer>> = Vec::new();
|
||||
let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config));
|
||||
let spend_tracking = spend.is_some();
|
||||
if let Some((pool, settings)) = spend {
|
||||
observers.push(litellm_gateway::spend::SpendWriter::start(pool, settings));
|
||||
}
|
||||
litellm_gateway::hooks::install(&inference, &config, Vec::new(), observers).await?;
|
||||
litellm_gateway::hooks::install(&inference, &config, hooks, observers).await?;
|
||||
let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?;
|
||||
let address = listener.local_addr()?;
|
||||
// Without TCP_NODELAY, Nagle holds small SSE chunks until the client's delayed ACK (~40 ms per stream).
|
||||
|
|
|
|||
495
litellm-rust/crates/gateway/tests/budgets.rs
Normal file
495
litellm-rust/crates/gateway/tests/budgets.rs
Normal file
|
|
@ -0,0 +1,495 @@
|
|||
//! The budgets hook on real requests: refusals as the proxy words them, reservations that keep
|
||||
//! concurrent spend under a budget, and counters shared through Redis. Set
|
||||
//! `LITELLM_TEST_REDIS_URL` to also run the shared cases against a live Redis.
|
||||
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, header::CONTENT_TYPE},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use litellm_auth_types::SecretValue;
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway::{
|
||||
budgets::{BudgetSettings, Budgets},
|
||||
coordination::{Coordination, RedisSettings},
|
||||
};
|
||||
use litellm_gateway_auth::keys::{
|
||||
AllowedRoutes, KeyBudgets, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture,
|
||||
MemberSpend, Spend,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpListener;
|
||||
use tower::ServiceExt;
|
||||
|
||||
/// Costs 3 * 0.001 + 4 * 0.002 = 0.011 at the test prices.
|
||||
const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#;
|
||||
const COST: f64 = 0.011;
|
||||
|
||||
async fn upstream() -> (String, Arc<AtomicUsize>) {
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let counted = calls.clone();
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = Router::new().fallback(move |_: Request<Body>| {
|
||||
let counted = counted.clone();
|
||||
async move {
|
||||
counted.fetch_add(1, Ordering::SeqCst);
|
||||
([(CONTENT_TYPE, "application/json")], MESSAGE).into_response()
|
||||
}
|
||||
});
|
||||
tokio::spawn(async move { axum::serve(listener, app).await });
|
||||
(format!("http://{address}"), calls)
|
||||
}
|
||||
|
||||
struct Keys(HashMap<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>, budgets: KeyBudgets) -> KeyStatus {
|
||||
KeyStatus::Active {
|
||||
expires_at: None,
|
||||
models: KeyModels::default(),
|
||||
allowed_routes: AllowedRoutes::default(),
|
||||
team_id: team_id.map(str::to_owned),
|
||||
profile: Arc::new(KeyProfile {
|
||||
alias: Some("prod".into()),
|
||||
user_id: Some("user-1".into()),
|
||||
budgets,
|
||||
..KeyProfile::default()
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn budget(spend: f64, max_budget: f64) -> Spend {
|
||||
Spend {
|
||||
spend,
|
||||
max_budget: Some(max_budget),
|
||||
}
|
||||
}
|
||||
|
||||
struct Gateway {
|
||||
app: Router,
|
||||
calls: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
async fn gateway(
|
||||
keys: Vec<(&str, KeyStatus)>,
|
||||
redis: Option<Arc<Coordination>>,
|
||||
settings: &str,
|
||||
) -> Gateway {
|
||||
gateway_with_database(keys, redis, None, settings).await
|
||||
}
|
||||
|
||||
async fn gateway_with_database(
|
||||
keys: Vec<(&str, KeyStatus)>,
|
||||
redis: Option<Arc<Coordination>>,
|
||||
database: Option<litellm_db::PgPool>,
|
||||
settings: &str,
|
||||
) -> Gateway {
|
||||
let (upstream, calls) = upstream().await;
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list:
|
||||
- model_name: claude
|
||||
litellm_params:
|
||||
model: anthropic/claude-x
|
||||
api_base: {upstream}
|
||||
api_key: sk-upstream
|
||||
input_cost_per_token: 0.001
|
||||
output_cost_per_token: 0.002
|
||||
model_info:
|
||||
id: claude-1
|
||||
- model_name: free
|
||||
litellm_params:
|
||||
model: anthropic/claude-x
|
||||
api_base: {upstream}
|
||||
api_key: sk-upstream
|
||||
input_cost_per_token: 0
|
||||
output_cost_per_token: 0
|
||||
model_info:
|
||||
id: free-1
|
||||
router_settings:
|
||||
num_retries: 0
|
||||
general_settings:
|
||||
master_key: sk-master
|
||||
{settings}"
|
||||
))
|
||||
.unwrap();
|
||||
let inference = litellm_gateway::build_inference(&config).unwrap();
|
||||
let budgets = Budgets::for_gateway(&config, &inference, redis, database);
|
||||
litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(budgets)], Vec::new())
|
||||
.await
|
||||
.unwrap();
|
||||
let keys: Arc<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, calls }
|
||||
}
|
||||
|
||||
async fn send(app: &Router, token: &str, model: &str) -> (u16, Value) {
|
||||
send_as(app, token, model, None).await
|
||||
}
|
||||
|
||||
async fn send_as(app: &Router, token: &str, model: &str, end_user: Option<&str>) -> (u16, Value) {
|
||||
let mut body = json!({
|
||||
"model": model,
|
||||
"max_tokens": 8,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
});
|
||||
if let Some(end_user) = end_user {
|
||||
body["metadata"] = json!({"user_id": end_user});
|
||||
}
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::post("/v1/messages")
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", format!("Bearer {token}"))
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status().as_u16();
|
||||
let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
|
||||
(status, serde_json::from_slice(&body).unwrap_or(Value::Null))
|
||||
}
|
||||
|
||||
/// Settling runs after the response, off the request path.
|
||||
async fn settled() {
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::key_at_budget(
|
||||
None,
|
||||
KeyBudgets { key: budget(10.0, 10.0), ..KeyBudgets::default() },
|
||||
"Budget has been exceeded! Key=prod Current cost: 10.0, Max budget: 10.0"
|
||||
)]
|
||||
#[case::zero_key_budget(
|
||||
None,
|
||||
KeyBudgets { key: budget(0.0, 0.0), ..KeyBudgets::default() },
|
||||
"Budget has been exceeded! Key=prod Current cost: 0.0, Max budget: 0.0"
|
||||
)]
|
||||
#[case::team_over_budget(
|
||||
Some("team-1"),
|
||||
KeyBudgets { team: Some(budget(5.25, 5.0)), ..KeyBudgets::default() },
|
||||
"Budget has been exceeded! Team=team-1 Current cost: 5.25, Max budget: 5.0"
|
||||
)]
|
||||
#[case::personal_user_over_budget(
|
||||
None,
|
||||
KeyBudgets { user: Some(budget(3.0, 3.0)), ..KeyBudgets::default() },
|
||||
"ExceededBudget: User=user-1 over budget. Spend=3.0, Budget=3.0"
|
||||
)]
|
||||
#[case::member_before_key(
|
||||
Some("team-1"),
|
||||
KeyBudgets {
|
||||
key: budget(10.0, 10.0),
|
||||
team_member: Some(MemberSpend { spend: 2.0, max_budget: Some(1.0), default_max_budget: None }),
|
||||
..KeyBudgets::default()
|
||||
},
|
||||
"Budget has been exceeded! TeamMember=user-1:team-1 Current cost: 2.0, Max budget: 1.0"
|
||||
)]
|
||||
#[case::zero_member_budget(
|
||||
Some("team-1"),
|
||||
KeyBudgets {
|
||||
team_member: Some(MemberSpend { spend: 0.0, max_budget: Some(0.0), default_max_budget: None }),
|
||||
..KeyBudgets::default()
|
||||
},
|
||||
"Budget has been exceeded! User=user-1 in Team=team-1 Current cost: 0.0, Max budget: 0.0"
|
||||
)]
|
||||
#[case::team_default_member_budget(
|
||||
Some("team-1"),
|
||||
KeyBudgets {
|
||||
team_member: Some(MemberSpend { spend: 4.0, max_budget: None, default_max_budget: Some(4.0) }),
|
||||
..KeyBudgets::default()
|
||||
},
|
||||
"Budget has been exceeded! TeamMember=user-1:team-1 Current cost: 4.0, Max budget: 4.0"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn over_budget_requests_are_refused_like_the_proxy(
|
||||
#[case] team_id: Option<&str>,
|
||||
#[case] budgets: KeyBudgets,
|
||||
#[case] message: &str,
|
||||
) {
|
||||
let gateway = gateway(vec![("sk-key", key(team_id, budgets))], None, "").await;
|
||||
let (status, body) = send(&gateway.app, "sk-key", "claude").await;
|
||||
assert_eq!(status, 422);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({"error": {"message": message, "type": "budget_exceeded", "param": null, "code": "422"}})
|
||||
);
|
||||
assert_eq!(gateway.calls.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::team_at_budget(
|
||||
Some("team-1"),
|
||||
KeyBudgets { team: Some(budget(5.0, 5.0)), ..KeyBudgets::default() }
|
||||
)]
|
||||
#[case::user_budget_on_team_key(
|
||||
Some("team-1"),
|
||||
KeyBudgets { user: Some(budget(3.0, 3.0)), ..KeyBudgets::default() }
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn budgets_the_proxy_does_not_apply_pass(
|
||||
#[case] team_id: Option<&str>,
|
||||
#[case] budgets: KeyBudgets,
|
||||
) {
|
||||
// A team exactly at its budget passes the proxy's check; only a reservation would refuse it.
|
||||
let gateway = gateway(
|
||||
vec![("sk-key", key(team_id, budgets))],
|
||||
None,
|
||||
" disable_budget_reservation: true\n",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(send(&gateway.app, "sk-key", "claude").await.0, 200);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn free_models_skip_budgets_and_the_status_code_is_configurable() {
|
||||
let budgets = KeyBudgets {
|
||||
key: budget(10.0, 10.0),
|
||||
..KeyBudgets::default()
|
||||
};
|
||||
let gateway = gateway(
|
||||
vec![("sk-key", key(None, budgets))],
|
||||
None,
|
||||
"litellm_settings:\n budget_exceeded_status_code: 429\n",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(send(&gateway.app, "sk-key", "free").await.0, 200);
|
||||
assert_eq!(send(&gateway.app, "sk-key", "claude").await.0, 429);
|
||||
}
|
||||
|
||||
/// Spends a key with a 0.02 budget down: the first request fits, the second's reservation
|
||||
/// shrinks to what is left, and the third finds the budget spent.
|
||||
async fn spend_down(app: &Router) {
|
||||
assert_eq!(send(app, "sk-key", "claude").await.0, 200);
|
||||
settled().await;
|
||||
assert_eq!(send(app, "sk-key", "claude").await.0, 200);
|
||||
settled().await;
|
||||
let (status, body) = send(app, "sk-key", "claude").await;
|
||||
assert_eq!(status, 422, "{body}");
|
||||
let message = body["error"]["message"].as_str().unwrap();
|
||||
assert!(
|
||||
message.starts_with("Budget has been exceeded! Key=prod Current cost: 0.02"),
|
||||
"{message}"
|
||||
);
|
||||
}
|
||||
|
||||
fn budgeted_key() -> KeyStatus {
|
||||
key(
|
||||
None,
|
||||
KeyBudgets {
|
||||
key: budget(0.0, 0.02),
|
||||
..KeyBudgets::default()
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn spend_is_counted_on_the_replica_without_redis() {
|
||||
let gateway = gateway(vec![("sk-key", budgeted_key())], None, "").await;
|
||||
spend_down(&gateway.app).await;
|
||||
assert_eq!(gateway.calls.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn reservations_hold_back_concurrent_requests() {
|
||||
// The first request's reservation takes the whole budget, so the second, sent before the
|
||||
// first settles, finds it spent.
|
||||
let key = key(
|
||||
None,
|
||||
KeyBudgets {
|
||||
key: budget(0.0, 0.0001),
|
||||
..KeyBudgets::default()
|
||||
},
|
||||
);
|
||||
let gateway = gateway(vec![("sk-key", key)], None, "").await;
|
||||
let (first, second) = tokio::join!(
|
||||
send(&gateway.app, "sk-key", "claude"),
|
||||
send(&gateway.app, "sk-key", "claude")
|
||||
);
|
||||
let mut statuses = [first.0, second.0];
|
||||
statuses.sort_unstable();
|
||||
assert_eq!(statuses, [200, 422]);
|
||||
}
|
||||
|
||||
fn redis_url() -> Option<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_spend_through_the_proxy_counters() {
|
||||
let Some(url) = redis_url() else { return };
|
||||
let namespace = format!("budgets-{}", uuid::Uuid::new_v4());
|
||||
let first = gateway(
|
||||
vec![("sk-key", budgeted_key())],
|
||||
Some(redis(&url, &namespace).await),
|
||||
"",
|
||||
)
|
||||
.await;
|
||||
let second = gateway(
|
||||
vec![("sk-key", budgeted_key())],
|
||||
Some(redis(&url, &namespace).await),
|
||||
"",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(send(&first.app, "sk-key", "claude").await.0, 200);
|
||||
settled().await;
|
||||
assert_eq!(send(&second.app, "sk-key", "claude").await.0, 200);
|
||||
settled().await;
|
||||
assert_eq!(send(&first.app, "sk-key", "claude").await.0, 422);
|
||||
|
||||
let mut connection = redis::Client::open(url.as_str())
|
||||
.unwrap()
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.unwrap();
|
||||
let counter = format!("{namespace}:spend:key:{}", hash("sk-key").as_str());
|
||||
let spend: String = redis::cmd("GET")
|
||||
.arg(&counter)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
(spend.parse::<f64>().unwrap() - 2.0 * COST).abs() < 1e-9,
|
||||
"{spend}"
|
||||
);
|
||||
let ttl: i64 = redis::cmd("TTL")
|
||||
.arg(&counter)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!((1..=60).contains(&ttl), "{ttl}");
|
||||
let user: String = redis::cmd("GET")
|
||||
.arg(format!("{namespace}:spend:user:user-1"))
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
(user.parse::<f64>().unwrap() - 2.0 * COST).abs() < 1e-9,
|
||||
"{user}"
|
||||
);
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::disallowed_methods,
|
||||
reason = "test fixture rows; the lookups under test use the checked macro"
|
||||
)]
|
||||
async fn end_users(pool: &litellm_db::PgPool) {
|
||||
sqlx::raw_sql(
|
||||
r#"
|
||||
INSERT INTO "LiteLLM_BudgetTable" (budget_id, max_budget, created_by, updated_by)
|
||||
VALUES ('small', 1.0, 'x', 'x'), ('default', 0.0001, 'x', 'x');
|
||||
INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id)
|
||||
VALUES ('over', 1.5, 'small'), ('at', 1.0, 'small'), ('unbudgeted', 9.0, NULL);
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn end_users_are_held_to_their_budgets() {
|
||||
let postgres = litellm_db_testing::MigratedPostgres::start().await.unwrap();
|
||||
end_users(postgres.pool()).await;
|
||||
let key = key(None, KeyBudgets::default());
|
||||
let gateway = gateway_with_database(
|
||||
vec![("sk-key", key)],
|
||||
None,
|
||||
Some(postgres.pool().clone()),
|
||||
" disable_budget_reservation: true\nlitellm_settings:\n max_end_user_budget_id: default\n",
|
||||
)
|
||||
.await;
|
||||
let (status, body) = send_as(&gateway.app, "sk-key", "claude", Some("over")).await;
|
||||
assert_eq!(status, 422);
|
||||
assert_eq!(
|
||||
body["error"]["message"],
|
||||
"ExceededBudget: End User=over over budget. Spend=1.5, Budget=1.0"
|
||||
);
|
||||
// Strictly over: an end user exactly at budget passes the check.
|
||||
assert_eq!(
|
||||
send_as(&gateway.app, "sk-key", "claude", Some("at"))
|
||||
.await
|
||||
.0,
|
||||
200
|
||||
);
|
||||
// Without a row, only a reservation applies the default budget, and reservations are off.
|
||||
assert_eq!(
|
||||
send_as(&gateway.app, "sk-key", "claude", Some("new"))
|
||||
.await
|
||||
.0,
|
||||
200
|
||||
);
|
||||
// A row without a budget takes the default one.
|
||||
let (status, body) = send_as(&gateway.app, "sk-key", "claude", Some("unbudgeted")).await;
|
||||
assert_eq!(status, 422);
|
||||
assert_eq!(
|
||||
body["error"]["message"],
|
||||
"ExceededBudget: End User=unbudgeted over budget. Spend=9.0, Budget=0.0001"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn settings_follow_the_proxy_defaults() {
|
||||
let settings = BudgetSettings::from_config(&Config::from_yaml("model_list: []").unwrap());
|
||||
assert_eq!(settings.status, 422);
|
||||
assert!(settings.reservations);
|
||||
assert_eq!(settings.ttl, Duration::from_secs(60));
|
||||
}
|
||||
51
litellm-rust/docs/customer-proxy/budgets.svg
Normal file
51
litellm-rust/docs/customer-proxy/budgets.svg
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 600" 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="600" fill="#ffffff"/>
|
||||
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">Budgets hook: proxy checks on the proxy's shared spend counters</text>
|
||||
|
||||
<rect x="24" y="60" width="300" height="140" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
|
||||
<text x="38" y="84" font-weight="600" fill="#1e3a8a">Key store row (cached with the key)</text>
|
||||
<text x="38" y="106" fill="#1e293b" font-size="13">key spend + max_budget (or budget row's)</text>
|
||||
<text x="38" y="124" fill="#1e293b" font-size="13">team, user spend + budgets</text>
|
||||
<text x="38" y="142" fill="#1e293b" font-size="13">team member spend, member budget</text>
|
||||
<text x="38" y="160" fill="#1e293b" font-size="13">(+ active temp increase, or team default)</text>
|
||||
<text x="38" y="182" fill="#64748b" font-size="13">end users: own lookup, cached 60 s</text>
|
||||
|
||||
<rect x="360" y="60" width="616" height="140" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
|
||||
<text x="374" y="84" font-weight="600" fill="#92400e">One Redis script per request (atomic)</text>
|
||||
<text x="374" y="106" fill="#1e293b" font-size="13">1. read spend:key / team / user / end_user / team_member counters; missing or below</text>
|
||||
<text x="374" y="124" fill="#1e293b" font-size="13"> the stored spend: seed or raise to it (the proxy's reseed + set-max repair)</text>
|
||||
<text x="374" y="142" fill="#1e293b" font-size="13">2. checks in the proxy's order: member, key, team (>), user, zero member, end user (>)</text>
|
||||
<text x="374" y="160" fill="#1e293b" font-size="13">3. reserve the max-cost estimate on every counter with a positive budget;</text>
|
||||
<text x="374" y="178" fill="#1e293b" font-size="13"> shrink to what is left, refuse when nothing is (1e-12), as the proxy</text>
|
||||
<line x1="324" y1="130" x2="356" y2="130" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<rect x="24" y="236" width="460" height="110" rx="8" fill="#fef2f2" stroke="#dc2626"/>
|
||||
<text x="38" y="260" font-weight="600" fill="#991b1b">Refusal: 422 (budget_exceeded_status_code)</text>
|
||||
<text x="38" y="282" fill="#1e293b" font-size="13">{"error": {"message", "type": "budget_exceeded",</text>
|
||||
<text x="38" y="300" fill="#1e293b" font-size="13"> "param": null, "code": "422"}} on both endpoints</text>
|
||||
<text x="38" y="322" fill="#1e293b" font-size="13">proxy messages, Python float repr; no provider call</text>
|
||||
<path d="M500 200 L500 220 L254 220 L254 232" fill="none" stroke="#dc2626" stroke-width="1.5" marker-end="url(#r)"/>
|
||||
|
||||
<rect x="516" y="236" width="460" height="110" rx="8" fill="#f8fafc" stroke="#64748b"/>
|
||||
<text x="530" y="260" font-weight="600" fill="#0f172a">Admitted: request runs</text>
|
||||
<text x="530" y="282" fill="#334155" font-size="13">estimate = input tokens x input price +</text>
|
||||
<text x="530" y="300" fill="#334155" font-size="13">min(max_tokens or 16384, model max) x output price</text>
|
||||
<text x="530" y="322" fill="#334155" font-size="13">lease renews reserved counters every TTL/2</text>
|
||||
<path d="M760 200 L760 232" fill="none" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<rect x="24" y="382" width="952" height="96" rx="8" fill="#f5f3ff" stroke="#8b5cf6"/>
|
||||
<text x="38" y="406" font-weight="600" fill="#5b21b6">Settle (off the request path)</text>
|
||||
<text x="38" y="428" fill="#1e293b" font-size="13">success: INCRBYFLOAT (actual - reserved) on reserved counters, actual on the others; failure: release the reservation</text>
|
||||
<text x="38" y="446" fill="#1e293b" font-size="13">a counter that expired mid-request is reseeded and charged the whole cost; EXPIRE ttl (default_redis_ttl, 60 s) refreshed</text>
|
||||
<text x="38" y="464" fill="#1e293b" font-size="13">script failure: the counters are deleted so the next read reseeds them, as the proxy does</text>
|
||||
<line x1="746" y1="346" x2="746" y2="378" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<rect x="24" y="508" width="952" height="72" rx="8" fill="#ecfdf5" stroke="#10b981"/>
|
||||
<text x="38" y="532" font-weight="600" fill="#065f46">Same keys, values and TTLs as the Python proxy, under its namespace</text>
|
||||
<text x="38" y="554" fill="#1e293b" font-size="13">gateway replicas and Python proxies on one Redis enforce one budget; the proxy's reset job deletes counters and they reseed.</text>
|
||||
<text x="38" y="572" fill="#64748b" font-size="13">No Redis, or Redis unreachable: the same logic on this replica's counters (fail_closed_budget_enforcement: 503 instead)</text>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 4.8 KiB |
Loading…
Add table
Reference in a new issue