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..852414ff4b7 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql @@ -0,0 +1,8 @@ +SET LOCAL lock_timeout = '10s'; + +-- 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 83a0e6f9e05..e4e83a77ac7 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 5970ede5c7e..1f40a2c4009 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1958,6 +1958,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 @@ -5179,6 +5180,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, float] | None # mutable-ok: accumulator # writable-ok: queue-reassigned team_member_list_transactions: dict[str, float] | None org_list_transactions: dict[str, float] | None org_member_list_transactions: ReadOnly[dict[str, float] | None] @@ -5192,6 +5196,7 @@ class SpendUpdateQueueItem(TypedDict, total=False): entity_type: Litellm_EntityType entity_id: str response_cost: float | None + tokens: int | None # writable-ok: total request tokens; dedup accumulates in place class ToolDiscoveryQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index d0708fcda2f..3c5c1e20900 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1407,6 +1407,11 @@ async def common_checks( proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, ), + team_max_tokens_check( + team_object=team_object, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + ), _team_multi_budget_check(team_object=team_object), _virtual_key_multi_budget_check(valid_token=valid_token) if valid_token is not None else None, _team_soft_budget_check( @@ -6908,6 +6913,67 @@ async def team_max_budget_check( ) +async def team_max_tokens_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + """Reject a request when the team is over its cumulative token budget. + + Mirrors ``team_max_budget_check``: dollars are a poor proxy for a token + allowance once a team mixes models at different prices, so a team can cap + tokens directly (``max_tokens_budget`` + ``budget_duration``). The usage is + read from the cross-pod ``tokens:team:{id}`` counter (Redis-first), falling + back to the ``tokens_used`` column loaded in auth. No budget set on the team + means the check is skipped, so existing teams are unaffected. + """ + if team_object is None: + return + max_tokens_budget: Final = team_object.max_tokens_budget + if max_tokens_budget is None: + return + + from litellm.proxy.proxy_server import get_current_tokens_used + + tokens_used: Final = await get_current_tokens_used( + counter_key=f"tokens:team:{team_object.team_id}", + fallback_tokens=team_object.tokens_used or 0, + ) + + if tokens_used > max_tokens_budget: + if valid_token: + call_info: Final = CallInfo( + token=valid_token.token, + spend=float(tokens_used), + max_budget=float(max_tokens_budget), + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.TEAM, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) + ) + + raise litellm.BudgetExceededError( + current_cost=float(tokens_used), + max_budget=float(max_tokens_budget), + message=( + f"Token budget has been exceeded! Team={team_object.team_id} " + f"Tokens used: {tokens_used}, Max tokens: {max_tokens_budget}" + ), + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, + ) + + +_team_max_tokens_check: Final = team_max_tokens_check + + _team_max_budget_check: Final = team_max_budget_check diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index ccf2232e379..581b7e511e1 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1400,6 +1400,9 @@ class ResetBudgetJob: team_id = getattr(t.row, "team_id", None) if team_id: await self._invalidate_spend_counter(f"spend:team:{team_id}") + # Drop the token-budget counter too, so the next read + # re-seeds from the freshly-zeroed tokens_used column. + await self._invalidate_spend_counter(f"tokens:team:{team_id}") end_time = time.time() outcome: Final = _ChunkOutcome( diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 166a42adfc0..e892dea5429 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1060,12 +1060,16 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + team_tokens: Final = payload.get( # pyright: ignore[reportUnknownMemberType] # SpendLogsPayload.get overload is partially unknown; total_tokens is a required int + "total_tokens", 0 + ) try: await self._update_team_db( response_cost=response_cost, team_id=team_id, user_id=user_id, prisma_client=prisma_client, + tokens=team_tokens, ) except Exception: verbose_proxy_logger.debug( @@ -1281,6 +1285,7 @@ class DBSpendUpdateWriter: team_id: str | None, user_id: str | None, prisma_client: PrismaClient | None, + tokens: int | None = None, ): try: if team_id is None or prisma_client is None: @@ -1294,6 +1299,7 @@ class DBSpendUpdateWriter: entity_type=Litellm_EntityType.TEAM, entity_id=team_id, response_cost=response_cost, + tokens=tokens, ) ) @@ -2165,6 +2171,7 @@ class DBSpendUpdateWriter: ### UPDATE TEAM TABLE ### team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"] + team_token_list_transactions: Final = db_spend_update_transactions.get("team_token_list_transactions") or {} verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions) if team_list_transactions is not None and len(team_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): @@ -2174,12 +2181,26 @@ class DBSpendUpdateWriter: async with transaction.batch_() as batcher: # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. for team_id, response_cost in sorted(team_list_transactions.items()): + # Increment tokens_used beside spend for token budgets; 0 when + # no token budget is in play, which leaves the column at its default. + token_increment = team_token_list_transactions.get(team_id) or 0 verbose_proxy_logger.debug( - "Updating spend for team id=%s by %s", team_id, response_cost + "Updating spend for team id=%s by %s (tokens +%s)", + team_id, + response_cost, + token_increment, + ) + team_update_data = ( + { + "spend": {"increment": response_cost}, + "tokens_used": {"increment": token_increment}, + } + if token_increment + else {"spend": {"increment": response_cost}} ) batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists where={"team_id": team_id}, - data={"spend": {"increment": response_cost}}, + data=team_update_data, ) break except Exception as e: diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 898aee493df..60f28a18768 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -74,6 +74,7 @@ _SpendTransactionField: TypeAlias = Literal[ "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", @@ -88,6 +89,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( "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", @@ -957,6 +959,9 @@ class RedisUpdateBuffer: end_user_list_transactions=_merged_entity_transactions(list_of_transactions, "end_user_list_transactions"), key_list_transactions=_merged_entity_transactions(list_of_transactions, "key_list_transactions"), team_list_transactions=_merged_entity_transactions(list_of_transactions, "team_list_transactions"), + team_token_list_transactions=_merged_entity_transactions( + list_of_transactions, "team_token_list_transactions" + ), team_member_list_transactions=_merged_entity_transactions( list_of_transactions, "team_member_list_transactions" ), 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 e88b9fc2bfb..f0d2eeeef7e 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -121,6 +121,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) @@ -138,6 +141,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={}, @@ -219,6 +223,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/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6fbd0b8f1b3..555c3e16272 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -382,6 +382,16 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ) + # An unpriced model has no dollar cost and skips the spend pipeline below, but it + # still consumes tokens a team token budget must count, so advance the counter here. + if ( + response_cost is None + and team_id is not None + and ProxyUpdateSpend.disable_spend_updates() is not True + and not _is_unbilled_interaction_response(completion_response) + ): + await _increment_team_token_counter(team_id=team_id, total_tokens=_extract_total_tokens(sl_object)) + if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) verbose_proxy_logger.debug( @@ -735,6 +745,35 @@ class _IncrementSpendCounters(Protocol): ) -> None: ... +def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> int: + """Total tokens for this request, for the team token-budget counter. + + Read from the StandardLoggingPayload, the same source the batched tokens_used + DB write uses, so the cross-pod counter and the column stay in agreement. + """ + if sl_object is None: + return 0 + return sl_object.get( # pyright: ignore[reportUnknownMemberType] # StandardLoggingPayload.get overload is partially unknown; total_tokens is a required int + "total_tokens", 0 + ) + + +async def _increment_team_token_counter(team_id: str, total_tokens: int) -> None: + """Advance the cross-pod ``tokens:team:{id}`` counter so the next request's + auth check sees this request's tokens before the batched ``tokens_used`` + flush lands. Best-effort, a counter hiccup must never fail a request that was + already served. + """ + if not total_tokens: + return + from litellm.proxy.proxy_server import increment_spend_counter_cache + + try: + await increment_spend_counter_cache(counter_key=f"tokens:team:{team_id}", increment=total_tokens) + except Exception: # noqa: BLE001 # best-effort cross-pod token counter; a Redis failure must not break the response or the spend write + verbose_proxy_logger.warning("Failed to increment team token-budget counter for team_id=%s", team_id) + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", increment_spend_counters: _IncrementSpendCounters, @@ -876,6 +915,19 @@ async def _update_database_and_spend_counters_in_batch( finally: budget_reservation["finalized"] = True raise + + # Mirror the spend counter for the team token budget: bump the cross-pod + # tokens:team:{id} counter so the next request's auth check sees this + # request's tokens before the batched DB flush of tokens_used lands. + # Awaited like the spend counter above; best-effort so a counter hiccup + # never fails a request that was already billed. + if team_id is not None: + sl_object: Final = cast( # cast-ok: typed shape of the untyped CustomLogger callback payload + "StandardLoggingPayload | None", + kwargs.get("standard_logging_object"), # pyright: ignore[reportUnknownMemberType] # untyped callback kwargs + ) + await _increment_team_token_counter(team_id=team_id, total_tokens=_extract_total_tokens(sl_object)) + return True diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py index 6a67093fde8..26d48d96845 100644 --- a/litellm/proxy/management_endpoints/scim/scim_transformations.py +++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py @@ -165,7 +165,7 @@ class ScimTransformations: raise HTTPException(status_code=500, detail={"error": "No database connected"}) if isinstance(team, dict): - team = LiteLLM_TeamTable(**team) + team = LiteLLM_TeamTable.model_validate(team) # Get team members with proper display names scim_members: Final[list[SCIMMember]] = [] diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index d1e1c4bd25e..edf8e874570 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -351,7 +351,7 @@ async def add_team_callbacks( # authenticated key holder could overwrite another team's logging # config (and read back the credentials they wrote). if not await get_team_access().allows( - user_api_key_dict, LiteLLM_TeamTable(**_existing_team.model_dump()), TEAM_OR_ORG_ADMIN + user_api_key_dict, LiteLLM_TeamTable.model_validate(_existing_team.model_dump()), TEAM_OR_ORG_ADMIN ): team_access_denied() @@ -534,7 +534,7 @@ async def delete_team_callback( # deregister its callbacks, otherwise any authenticated key holder could # silence another team's observability integration. if not await get_team_access().allows( - user_api_key_dict, LiteLLM_TeamTable(**_existing_team.model_dump()), TEAM_OR_ORG_ADMIN + user_api_key_dict, LiteLLM_TeamTable.model_validate(_existing_team.model_dump()), TEAM_OR_ORG_ADMIN ): team_access_denied() @@ -667,7 +667,7 @@ async def disable_team_logging( # team may disable its logging — otherwise any authenticated key # holder can silence audit logging for any team. if not await get_team_access().allows( - user_api_key_dict, LiteLLM_TeamTable(**_existing_team.model_dump()), TEAM_OR_ORG_ADMIN + user_api_key_dict, LiteLLM_TeamTable.model_validate(_existing_team.model_dump()), TEAM_OR_ORG_ADMIN ): team_access_denied() @@ -808,7 +808,7 @@ async def get_team_callbacks( # (Langfuse / Langsmith / GCS). Only proxy admins / org admins / # team admins of THIS team may read them. if not await get_team_access().allows( - user_api_key_dict, LiteLLM_TeamTable(**_existing_team.model_dump()), TEAM_OR_ORG_ADMIN + user_api_key_dict, LiteLLM_TeamTable.model_validate(_existing_team.model_dump()), TEAM_OR_ORG_ADMIN ): team_access_denied() diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0f40662ad83..134223c7564 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1481,6 +1481,7 @@ async def new_team( - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement. - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement. - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget + - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped. - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed. @@ -1801,9 +1802,8 @@ async def new_team( ) ## ADD TO TEAM TABLE - complete_team_data: Final = LiteLLM_TeamTable( - **data_json, - model_id=_model_id, + complete_team_data: Final = LiteLLM_TeamTable.model_validate( + {**data_json, "model_id": _model_id}, ) # Set Management Endpoint Metadata Fields @@ -2260,6 +2260,7 @@ async def update_team( - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget + - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped. - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dba4b987725..92dd9558762 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3152,6 +3152,28 @@ async def get_current_spend( return current +async def get_current_tokens_used(counter_key: str, fallback_tokens: int) -> int: + """Cross-pod cumulative token usage for a team's token budget. + + Reads the Redis-backed ``tokens:team:{id}`` counter first (authoritative + across pods), the same way :func:`get_current_spend` reads the spend + counter, and falls back to the ``tokens_used`` column loaded in auth when + the counter is cold or Redis is unreachable. The counter read is floored at + that column. An idle counter expires after its TTL and the next post-call + increment re-seeds it from only that request's delta, so a bare counter read + can sit below the recorded usage and let a team keep spending past the cap. + Flooring at the authoritative column closes that gap the same way + ``get_current_spend`` reconciles a stale-low counter against recorded spend, + and also covers a counter that missed an increment when its logging payload + failed to build. + """ + value, _ = await read_spend_counter_cache_value(counter_key=counter_key) + floor: Final = int(fallback_tokens or 0) + if value is not None: + return max(int(value), floor) + return floor + + @with_service_target(SPEND_COUNTERS_TARGET) async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None: """Raise a counter that has fallen below the authoritative DB spend (e.g. diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 83a0e6f9e05..e4e83a77ac7 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/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 95f4f5e81be..02449491ec5 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -56,9 +56,15 @@ class TeamSpendResetWrites: table: BatchTable def queue_spend_reset(self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: + data: Final = dict(_spend_reset_data(budget_reset_at, spend_decrement)) + # Token budgets are team-only, so zero the cumulative tokens_used on the + # same tick that resets spend. A hard reset (vs spend's decrement) is + # acceptable at a budget boundary: a token allowance need not carry the + # handful of tokens that can race the reset the way billing-exact spend does. + data["tokens_used"] = 0 self.table.update( where={"team_id": team_id}, - data=_spend_reset_data(budget_reset_at, spend_decrement), + data=data, ) diff --git a/schema.prisma b/schema.prisma index 83a0e6f9e05..e4e83a77ac7 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") diff --git a/tests/unit/proxy/auth/test_team_token_budget.py b/tests/unit/proxy/auth/test_team_token_budget.py new file mode 100644 index 00000000000..30080a31ee2 --- /dev/null +++ b/tests/unit/proxy/auth/test_team_token_budget.py @@ -0,0 +1,187 @@ +"""Tests for the per-team cumulative token budget check. + +team_max_tokens_check mirrors team_max_budget_check: a team can cap tokens +directly (max_tokens_budget + budget_duration) because a dollar cap drifts once +the team mixes models at different per-token prices. Usage is read from the +cross-pod tokens:team:{id} counter (get_current_tokens_used), falling back to +the tokens_used column loaded in auth. +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth +from litellm.models.team import LiteLLM_TeamTable +from litellm.proxy.auth.auth_checks import team_max_tokens_check +from litellm.proxy.utils import ProxyLogging + + +def _proxy_logging_obj() -> ProxyLogging: + obj = ProxyLogging.__new__(ProxyLogging) + + async def _noop(*args, **kwargs): + return None + + obj.budget_alerts = _noop # type: ignore[attr-defined] + return obj + + +def _patch_tokens(value): + async def _effect(counter_key, fallback_tokens): + return value if value is not None else int(fallback_tokens or 0) + + return patch( + "litellm.proxy.proxy_server.get_current_tokens_used", + new_callable=AsyncMock, + side_effect=_effect, + ) + + +@pytest.mark.asyncio +async def test_no_budget_set_is_skipped(): + """A team with no token budget is never blocked, so existing teams are unaffected.""" + team = LiteLLM_TeamTable(team_id="t1", tokens_used=10**9) + token = UserAPIKeyAuth(token="hashed", team_id="t1") + with _patch_tokens(10**9) as get_tokens: + result = await team_max_tokens_check( + team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj() + ) + assert result is None + assert get_tokens.await_count == 0 + + +@pytest.mark.asyncio +async def test_under_budget_passes(): + team = LiteLLM_TeamTable(team_id="t1", max_tokens_budget=1_000_000, tokens_used=0) + token = UserAPIKeyAuth(token="hashed", team_id="t1") + with _patch_tokens(500_000) as get_tokens: + result = await team_max_tokens_check( + team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj() + ) + assert result is None + assert get_tokens.await_count == 1 + + +@pytest.mark.asyncio +async def test_over_budget_raises(): + team = LiteLLM_TeamTable(team_id="t1", max_tokens_budget=1_000_000, tokens_used=0) + token = UserAPIKeyAuth(token="hashed", team_id="t1") + with _patch_tokens(1_200_000): + with pytest.raises(litellm.BudgetExceededError) as exc: + await team_max_tokens_check(team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj()) + assert exc.value.current_cost == 1_200_000 + assert exc.value.max_budget == 1_000_000 + assert exc.value.entity_type == Litellm_EntityType.TEAM.value + + +@pytest.mark.asyncio +async def test_cold_counter_falls_back_to_db_tokens_used(): + """When the counter is cold (None), usage comes from the tokens_used column + loaded in auth; an over-budget column still blocks.""" + team = LiteLLM_TeamTable(team_id="t1", max_tokens_budget=1_000_000, tokens_used=1_500_000) + token = UserAPIKeyAuth(token="hashed", team_id="t1") + with _patch_tokens(None): + with pytest.raises(litellm.BudgetExceededError): + await team_max_tokens_check(team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj()) + + +@pytest.mark.asyncio +async def test_exactly_at_cap_passes(): + """Enforcement is strictly greater-than, matching the dollar budget check.""" + team = LiteLLM_TeamTable(team_id="t1", max_tokens_budget=1_000_000, tokens_used=0) + token = UserAPIKeyAuth(token="hashed", team_id="t1") + with _patch_tokens(1_000_000) as get_tokens: + result = await team_max_tokens_check( + team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj() + ) + assert result is None + assert get_tokens.await_count == 1 + + +@pytest.mark.asyncio +async def test_get_current_tokens_used_prefers_the_cross_pod_counter(): + """A populated tokens:team counter is authoritative over the DB column.""" + from litellm.proxy.proxy_server import get_current_tokens_used + + async def _counter(counter_key): + return 500_000, None + + with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _counter): + used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=7) + assert used == 500_000 + + +@pytest.mark.asyncio +async def test_get_current_tokens_used_falls_back_to_db_when_counter_cold(): + """A cold counter (None) falls back to the tokens_used column loaded in auth.""" + from litellm.proxy.proxy_server import get_current_tokens_used + + async def _cold(counter_key): + return None, None + + with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _cold): + used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=1_234) + assert used == 1_234 + + +def test_extract_total_tokens_reads_the_standard_logging_payload(): + """The counter delta is the StandardLoggingPayload total, the DB write's source.""" + from litellm.proxy.hooks.proxy_track_cost_callback import _extract_total_tokens + + assert _extract_total_tokens({"total_tokens": 4242}) == 4242 + + +def test_extract_total_tokens_is_zero_without_a_payload(): + """No payload leaves the counter untouched rather than raising.""" + from litellm.proxy.hooks.proxy_track_cost_callback import _extract_total_tokens + + assert _extract_total_tokens(None) == 0 + + +@pytest.mark.asyncio +async def test_get_current_tokens_used_floors_at_recorded_db_usage(): + """A counter that reseeded low after a TTL expiry, or missed an increment when its + logging payload failed, must not read below the tokens_used column loaded in auth, + otherwise a team slips past the cap once the counter sits under recorded usage.""" + from litellm.proxy.proxy_server import get_current_tokens_used + + async def _stale_low(counter_key): + return 200, None + + with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _stale_low): + used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=1_100) + assert used == 1_100 + + +@pytest.mark.asyncio +async def test_unpriced_model_advances_the_team_token_counter(): # test-quality-ok: no in-process counter state + """An unpriced model returns response_cost=None and skips the spend pipeline, but its + tokens must still count against the team, otherwise it bypasses max_tokens_budget. The + cross-pod counter is Redis-backed, so the issued increment is the only in-process signal.""" + from datetime import datetime + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + + kwargs = { + "model": "custom-unpriced", + "call_type": "acompletion", + "litellm_params": {"metadata": {"user_api_key_team_id": "t1"}}, + "standard_logging_object": { + "response_cost": None, + "response_cost_failure_debug_info": "missing custom price", + "request_tags": None, + "total_tokens": 321, + }, + "stream": False, + } + + with patch("litellm.proxy.proxy_server.increment_spend_counter_cache", new_callable=AsyncMock) as inc: + await ProxyDBLogger()._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + ) + inc.assert_awaited_once_with(counter_key="tokens:team:t1", increment=321) diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 70f7774e691..7e9dd11ba72 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -431,7 +431,9 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): assert write["where"] == {"team_id": "tid-1"} assert write["data"]["spend"] == {"decrement": 500.0} assert write["data"]["budget_reset_at"] > now - assert set(write["data"].keys()) == {"spend", "budget_reset_at"} + # Team resets also hard-zero the token-budget accumulator for the new window. + assert write["data"]["tokens_used"] == 0 + assert set(write["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"} def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): @@ -663,7 +665,11 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}" assert writes[0]["where"] == where assert writes[0]["data"]["spend"] == {"decrement": decrement} - assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"} + # Teams carry the extra token-budget accumulator, zeroed for the new window. + expected_keys = {"spend", "budget_reset_at"} + if table_name == "team": + expected_keys = expected_keys | {"tokens_used"} + assert set(writes[0]["data"].keys()) == expected_keys # The budget tier's cascade rides the same batch machinery. assert _batch_writes(mock_prisma_client, "enduser") == [ @@ -4216,9 +4222,13 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c assert len(user_writes) == 1 assert user_writes[0]["where"] == {"user_id": "user2"} assert len(team_writes) == 2 - for c in key_writes + user_writes + team_writes: + for c in key_writes + user_writes: assert set(c["data"].keys()) == {"spend", "budget_reset_at"} assert c["data"]["spend"] == {"decrement": pre_reset_spend[next(iter(c["where"].values()))]} + # Team writes additionally zero the token-budget accumulator for the new window. + for c in team_writes: + assert set(c["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"} + assert c["data"]["spend"] == {"decrement": pre_reset_spend[next(iter(c["where"].values()))]} @pytest.mark.asyncio @@ -4274,7 +4284,7 @@ async def test_reset_budget_teams_partial_failure(): team_writes = [c for c in batch_calls if c["table"] == "team"] assert len(team_writes) == 1 assert team_writes[0]["where"] == {"team_id": "team2"} - assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"} + assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"} assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend} failure_hook_calls = ( diff --git a/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py index 9a4b02ba3ba..ce53d695b8e 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py @@ -360,3 +360,44 @@ async def test_queue_size_reduction_with_large_volume(monkeypatch, spend_queue): ) assert aggregated["user_list_transactions"]["user1"] == 200 * 0.5 assert aggregated["key_list_transactions"]["key1"] == 300 * 1.0 + + +@pytest.mark.asyncio +async def test_team_token_aggregation(spend_queue): + """Per-team token deltas sum into team_token_list_transactions, parallel to + the spend sum, so LiteLLM_TeamTable.tokens_used is flushed on the same batch + cadence as spend. Non-team updates never create token transactions.""" + updates: list[SpendUpdateQueueItem] = [ + {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team1", "response_cost": 0.5, "tokens": 100}, + {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team1", "response_cost": 0.25, "tokens": 50}, + {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team2", "response_cost": 1.0, "tokens": 400}, + # a team update without a tokens field counts as 0 tokens, spend still applies + {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team2", "response_cost": 0.2}, + # a user update carries no team tokens + {"entity_type": Litellm_EntityType.USER, "entity_id": "user1", "response_cost": 0.1, "tokens": 9}, + ] + for update in updates: + await spend_queue.update_queue.put(update) + + aggregated = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + + # spend is unchanged by the token plumbing + assert aggregated["team_list_transactions"]["team1"] == 0.75 + assert aggregated["team_list_transactions"]["team2"] == 1.2 + # tokens sum per team + assert aggregated["team_token_list_transactions"]["team1"] == 150 + assert aggregated["team_token_list_transactions"]["team2"] == 400 + # no team token entry leaked from the user update + assert "user1" not in aggregated["team_token_list_transactions"] + + +@pytest.mark.asyncio +async def test_team_token_aggregation_absent_when_no_tokens(spend_queue): + """A plain team spend update (no tokens) leaves tokens_used untouched: the + token map has the team at 0, which the commit treats as no increment.""" + await spend_queue.update_queue.put( + {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team1", "response_cost": 0.5} + ) + aggregated = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert aggregated["team_list_transactions"]["team1"] == 0.5 + assert aggregated["team_token_list_transactions"].get("team1", 0) == 0 diff --git a/tests/unit/proxy/memory/test_memory_endpoints.py b/tests/unit/proxy/memory/test_memory_endpoints.py index dff0e80fa77..f0a0fa9a259 100644 --- a/tests/unit/proxy/memory/test_memory_endpoints.py +++ b/tests/unit/proxy/memory/test_memory_endpoints.py @@ -197,6 +197,7 @@ def _make_team(team_id: str, *, admin_user_ids: List[str]) -> Any: created_at=now, updated_at=now, spend=0.0, + tokens_used=0, model_spend="{}", model_max_budget="{}", admins=[], diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 468811966de..bac5cafe0f7 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -17010,6 +17010,7 @@ export interface paths { * - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement. * - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement. * - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget + * - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped. * - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. * - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) * - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed. @@ -17241,6 +17242,7 @@ export interface paths { * - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit * - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit * - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget + * - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped. * - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. * - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) * - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed. @@ -34132,6 +34134,8 @@ export interface components { max_budget?: number | null; /** Max Parallel Requests */ max_parallel_requests?: number | null; + /** Max Tokens Budget */ + max_tokens_budget?: number | null; /** * Members * @default [] @@ -34190,6 +34194,8 @@ export interface components { team_id: string; /** Team Member Permissions */ team_member_permissions?: string[] | null; + /** Tokens Used */ + tokens_used?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -35667,6 +35673,8 @@ export interface components { max_budget?: number | null; /** Max Parallel Requests */ max_parallel_requests?: number | null; + /** Max Tokens Budget */ + max_tokens_budget?: number | null; /** * Members * @default [] @@ -35725,6 +35733,8 @@ export interface components { team_id: string; /** Team Member Permissions */ team_member_permissions?: string[] | null; + /** Tokens Used */ + tokens_used?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -39016,6 +39026,8 @@ export interface components { guardrails?: string[] | null; /** Max Budget */ max_budget?: number | null; + /** Max Tokens Budget */ + max_tokens_budget?: number | null; /** Mcp Rpm Limit */ mcp_rpm_limit?: { [key: string]: number; @@ -40306,6 +40318,8 @@ export interface components { guardrails?: string[] | null; /** Max Budget */ max_budget?: number | null; + /** Max Tokens Budget */ + max_tokens_budget?: number | null; /** Mcp Rpm Limit */ mcp_rpm_limit?: { [key: string]: number; @@ -45973,6 +45987,8 @@ export interface components { max_budget?: number | null; /** Max Parallel Requests */ max_parallel_requests?: number | null; + /** Max Tokens Budget */ + max_tokens_budget?: number | null; /** * Members * @default [] @@ -46031,6 +46047,8 @@ export interface components { team_id: string; /** Team Member Permissions */ team_member_permissions?: string[] | null; + /** Tokens Used */ + tokens_used?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -46112,6 +46130,8 @@ export interface components { max_budget?: number | null; /** Max Parallel Requests */ max_parallel_requests?: number | null; + /** Max Tokens Budget */ + max_tokens_budget?: number | null; /** * Members * @default [] @@ -46175,6 +46195,8 @@ export interface components { team_id: string; /** Team Member Permissions */ team_member_permissions?: string[] | null; + /** Tokens Used */ + tokens_used?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -48452,6 +48474,8 @@ export interface components { guardrails?: string[] | null; /** Max Budget */ max_budget?: number | null; + /** Max Tokens Budget */ + max_tokens_budget?: number | null; /** Mcp Rpm Limit */ mcp_rpm_limit?: { [key: string]: number;