mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
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:
parent
f453c42509
commit
fbda977c5e
3 changed files with 17 additions and 20 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue