diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56773789059..c0c1d12e9ee 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1404,6 +1404,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( @@ -6518,6 +6523,66 @@ 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. + """ + max_tokens_budget = getattr(team_object, "max_tokens_budget", None) + if team_object is None or 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/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index da8a7fb0f15..5b9e878113c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1066,6 +1066,7 @@ class DBSpendUpdateWriter: team_id=team_id, user_id=user_id, prisma_client=prisma_client, + tokens=int(payload.get("total_tokens") or 0), ) except Exception: verbose_proxy_logger.debug( @@ -1281,6 +1282,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 +1296,7 @@ class DBSpendUpdateWriter: entity_type=Litellm_EntityType.TEAM, entity_id=team_id, response_cost=response_cost, + tokens=tokens, ) ) @@ -2165,6 +2168,9 @@ 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 +2180,21 @@ 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: dict = {"spend": {"increment": response_cost}} + if token_increment: + team_update_data["tokens_used"] = {"increment": token_increment} 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/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0938d3f35c3..d670855820c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -733,6 +733,29 @@ class _IncrementSpendCounters(Protocol): ) -> None: ... +def _extract_total_tokens(kwargs: dict, completion_response: object) -> int: + """Total tokens for this request, for the team token-budget counter. + + Prefers the StandardLoggingPayload (always populated for a billed call), + falling back to the response usage. Returns 0 when unknown, which leaves the + token counter untouched. + """ + sl_object = kwargs.get("standard_logging_object") or {} + total = sl_object.get("total_tokens") + if total: + return int(total) + usage = getattr(completion_response, "usage", None) + if usage is None and isinstance(completion_response, dict): + usage = completion_response.get("usage") + if usage is not None: + total = getattr(usage, "total_tokens", None) + if total is None and isinstance(usage, dict): + total = usage.get("total_tokens") + if total: + return int(total) + return 0 + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", increment_spend_counters: _IncrementSpendCounters, @@ -874,6 +897,27 @@ 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: + total_tokens: Final = _extract_total_tokens(kwargs=kwargs, completion_response=completion_response) + if total_tokens: + 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: + verbose_proxy_logger.warning( + "Failed to increment team token-budget counter for team_id=%s", team_id + ) + return True diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 089647c6297..24fba4a5635 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3020,6 +3020,24 @@ 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 is incremented + post-call in the cost callback and invalidated on the budget-reset tick, so + a read after a reset re-seeds from the freshly-zeroed column. Token budgets + tolerate the spend path's fallback-to-DB behavior for the same reason spend + does: the column is the authoritative value the reset job maintains. + """ + value, _ = await read_spend_counter_cache_value(counter_key=counter_key) + if value is not None: + return int(value) + return int(fallback_tokens or 0) + + @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.