diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql new file mode 100644 index 00000000000..9306edd7d75 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql @@ -0,0 +1,6 @@ +-- AlterTable +-- Cumulative per-team token budget, mirroring the existing dollar budget +-- (max_budget / spend). max_tokens_budget is the cap (null = untracked); +-- tokens_used is the accumulator reset on the shared budget_reset_at tick. +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "max_tokens_budget" BIGINT; +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "tokens_used" BIGINT NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f3ebee52dcd..3bef9af3b9e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -197,6 +197,8 @@ model LiteLLM_TeamTable { tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? + max_tokens_budget BigInt? // cumulative token cap over budget_duration; null = untracked (mirrors max_budget) + tokens_used BigInt @default(0) // cumulative tokens consumed in the current window (mirrors spend) blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/models/team.py b/litellm/models/team.py index 8edf10703b1..79b3a6c0617 100644 --- a/litellm/models/team.py +++ b/litellm/models/team.py @@ -74,6 +74,8 @@ class TeamBase(LiteLLMPydanticObjectBase): tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None + # cumulative token cap over budget_duration (mirrors max_budget); null = untracked + max_tokens_budget: int | None = None budget_duration: str | None = None budget_limits: list[BudgetLimitEntry] | None = None models: list[str] = [] @@ -86,6 +88,7 @@ class TeamBase(LiteLLMPydanticObjectBase): class LiteLLM_TeamTable(TeamBase): team_id: str spend: float | None = None + tokens_used: int | None = None # cumulative tokens consumed this window (mirrors spend) max_parallel_requests: int | None = None budget_duration: str | None = None budget_reset_at: datetime | None = None diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b6c81f17761..17978256526 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1947,6 +1947,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None + max_tokens_budget: int | None = None # cumulative token cap over budget_duration models: list | None = None blocked: bool | None = None budget_duration: str | None = None @@ -5111,6 +5112,9 @@ class DBSpendUpdateTransactions(TypedDict): end_user_list_transactions: dict[str, float] | None key_list_transactions: dict[str, float] | None team_list_transactions: dict[str, float] | None + # cumulative token deltas per team, flushed to LiteLLM_TeamTable.tokens_used + # alongside team_list_transactions (spend). Tokens budgets only apply to teams. + team_token_list_transactions: dict[str, int] | None team_member_list_transactions: dict[str, float] | None org_list_transactions: dict[str, float] | None org_member_list_transactions: ReadOnly[dict[str, float] | None] @@ -5124,6 +5128,7 @@ class SpendUpdateQueueItem(TypedDict, total=False): entity_type: Litellm_EntityType entity_id: str response_cost: float | None + tokens: int | None # total tokens for this request; summed into tokens_used for TEAM updates class ToolDiscoveryQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 2b8535cb113..2524996d42f 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -118,6 +118,9 @@ class SpendUpdateQueue(BaseUpdateQueue): current_cost = _in_memory_map[_key].get("response_cost", 0) or 0 update_cost = update.get("response_cost", 0) or 0 _in_memory_map[_key]["response_cost"] = current_cost + update_cost + current_tokens = _in_memory_map[_key].get("tokens", 0) or 0 + update_tokens = update.get("tokens", 0) or 0 + _in_memory_map[_key]["tokens"] = current_tokens + update_tokens for _key, update in _in_memory_map.items(): aggregated_spend_updates.append(update) @@ -135,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue): end_user_list_transactions={}, key_list_transactions={}, team_list_transactions={}, + team_token_list_transactions={}, team_member_list_transactions={}, org_list_transactions={}, org_member_list_transactions={}, @@ -216,6 +220,15 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict[entity_id] += response_cost or 0 + # Tokens budgets only apply to teams, so sum the token delta into a + # parallel map flushed to LiteLLM_TeamTable.tokens_used beside spend. + if entity_type == Litellm_EntityType.TEAM: + team_tokens = db_spend_update_transactions["team_token_list_transactions"] + if team_tokens is None: + team_tokens = {} + db_spend_update_transactions["team_token_list_transactions"] = team_tokens + team_tokens[entity_id] = (team_tokens.get(entity_id) or 0) + (update.get("tokens") or 0) + return db_spend_update_transactions async def _emit_new_item_added_to_queue_event( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f3ebee52dcd..3bef9af3b9e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -197,6 +197,8 @@ model LiteLLM_TeamTable { tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? + max_tokens_budget BigInt? // cumulative token cap over budget_duration; null = untracked (mirrors max_budget) + tokens_used BigInt @default(0) // cumulative tokens consumed in the current window (mirrors spend) blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/schema.prisma b/schema.prisma index f3ebee52dcd..3bef9af3b9e 100644 --- a/schema.prisma +++ b/schema.prisma @@ -197,6 +197,8 @@ model LiteLLM_TeamTable { tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? + max_tokens_budget BigInt? // cumulative token cap over budget_duration; null = untracked (mirrors max_budget) + tokens_used BigInt @default(0) // cumulative tokens consumed in the current window (mirrors spend) blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at")