fix(proxy): keep token-budget code within the basedpyright ceiling

Type the token lookups instead of reading untyped callback data:
- _extract_total_tokens takes a StandardLoggingPayload and returns its typed
  total_tokens, the same source the batched tokens_used DB write uses
- read max_tokens_budget off the typed LiteLLM_TeamTable, not getattr
- read total_tokens off the typed SpendLogsPayload, not .get
- drop Final on the loop-assigned team_update_data (illegal, and the gate
  exempts loop vars from the Final rule anyway)
The one remaining untyped boundary (the CustomLogger kwargs dict) is cast once
with a reason.
This commit is contained in:
Srivatsa03 2026-10-11 22:31:07 -05:00
parent f453c42509
commit fbda977c5e
3 changed files with 17 additions and 20 deletions

View file

@ -6537,8 +6537,10 @@ 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: Final = getattr(team_object, "max_tokens_budget", None)
if team_object is None or max_tokens_budget is None:
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

View file

@ -1066,7 +1066,7 @@ class DBSpendUpdateWriter:
team_id=team_id,
user_id=user_id,
prisma_client=prisma_client,
tokens=int(payload.get("total_tokens") or 0),
tokens=payload["total_tokens"],
)
except Exception:
verbose_proxy_logger.debug(
@ -2187,7 +2187,7 @@ class DBSpendUpdateWriter:
response_cost,
token_increment,
)
team_update_data: Final = (
team_update_data = (
{
"spend": {"increment": response_cost},
"tokens_used": {"increment": token_increment},

View file

@ -733,24 +733,15 @@ class _IncrementSpendCounters(Protocol):
) -> None: ...
def _extract_total_tokens(kwargs: Mapping[str, object], completion_response: object) -> int:
def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> 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.
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.
"""
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
if sl_object is None:
return 0
return sl_object["total_tokens"]
async def _update_database_and_spend_counters(
@ -901,7 +892,11 @@ async def _update_database_and_spend_counters_in_batch(
# 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)
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
)
total_tokens: Final = _extract_total_tokens(sl_object)
if total_tokens:
from litellm.proxy.proxy_server import increment_spend_counter_cache