From cb4561e27bb59e92c0a5ce6e63301ec89c440a2c Mon Sep 17 00:00:00 2001 From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com> Date: Sun, 11 Oct 2026 22:05:22 -0500 Subject: [PATCH] fix(proxy): satisfy type-discipline ratchet for token-budget code Keep the custom LIT rule counts at their merge-base ceiling: - _extract_total_tokens: single-assignment Final locals (no rebinding) and Mapping[str, object] for the read-only kwargs param - team_update_data: build the payload functionally under a Final binding - max_tokens_budget / reset data: declare Final - DBSpendUpdateTransactions.team_token_list_transactions and SpendUpdateQueueItem.tokens: mark the in-place accumulators mutable-ok / writable-ok (they are reassigned by the queue, mirroring their spend siblings) --- litellm/proxy/_types.py | 4 +-- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/db/db_spend_update_writer.py | 11 +++++--- .../proxy/hooks/proxy_track_cost_callback.py | 27 +++++++++---------- litellm/repositories/unit_of_work.py | 2 +- 5 files changed, 24 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 483d9485ec5..95c3244be6c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5114,7 +5114,7 @@ class DBSpendUpdateTransactions(TypedDict): 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 + 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] @@ -5128,7 +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 + 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 7cc6b685a2a..cfeb13f3b3f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6537,7 +6537,7 @@ async def team_max_tokens_check( 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) + max_tokens_budget: Final = getattr(team_object, "max_tokens_budget", None) if team_object is None or max_tokens_budget is None: return diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 7e5137b8ade..657456f37a9 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -2187,9 +2187,14 @@ class DBSpendUpdateWriter: response_cost, token_increment, ) - team_update_data: dict = {"spend": {"increment": response_cost}} - if token_increment: - team_update_data["tokens_used"] = {"increment": token_increment} + team_update_data: Final = ( + { + "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=team_update_data, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1352719aaa6..3c3a05353e3 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -733,27 +733,24 @@ class _IncrementSpendCounters(Protocol): ) -> None: ... -def _extract_total_tokens(kwargs: dict, completion_response: object) -> int: +def _extract_total_tokens(kwargs: Mapping[str, object], 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 + sl_object: Final = kwargs.get("standard_logging_object") or {} + sl_total: Final = sl_object.get("total_tokens") + if sl_total: + return int(sl_total) + usage: Final = ( + completion_response.get("usage") + if isinstance(completion_response, dict) + else getattr(completion_response, "usage", None) + ) + usage_total: Final = usage.get("total_tokens") if isinstance(usage, dict) else getattr(usage, "total_tokens", None) + return int(usage_total) if usage_total else 0 async def _update_database_and_spend_counters( diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 37f2cfbeaef..3f8e44022b2 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -56,7 +56,7 @@ class TeamSpendResetWrites: table: BatchTable def queue_spend_reset(self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: - data = dict(_spend_reset_data(budget_reset_at, spend_decrement)) + 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