scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): count snake_case candidate_count in the Gemini budget hold

Browse source
This commit is contained in:
mateo-berri 2026-10-28 22:17:42 -07:00
parent d793d2917f
commit fb75a132e5
2 changed files with 12 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -7,6 +7,7 @@ import time
from collections.abc import AsyncIterator, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import product
from types import MappingProxyType
from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast
@ -1651,6 +1652,7 @@ def _estimate_output_tokens(
_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens")
_GEMINI_CONFIG_FIELDS: Final = ("generationConfig", "generation_config")
def _nested_field(request_body: Mapping[str, object], config_field: str, field: str) -> object:
@ -1658,6 +1660,13 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field:
return config.get(field) if config is not None else None
def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object, ...]:
return tuple(
_nested_field(request_body, config_field, field)
for config_field, field in product(_GEMINI_CONFIG_FIELDS, ("candidateCount", "candidate_count"))
)
def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None:
candidates: Final = (
*(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS),
@ -1671,7 +1680,7 @@ def _get_output_multiplier(request_body: Mapping[str, object]) -> int:
candidates: Final = (
request_body.get("n"),
request_body.get("best_of"),
_nested_field(request_body, "generationConfig", "candidateCount"),
*_gemini_candidate_counts(request_body),
)
return max((1, *(count for count in map(_to_int, candidates) if count is not None)))

3
tests/unit/proxy/spend_tracking/test_budget_reservation.py
View file

@ -193,7 +193,8 @@ GEMINI_SYSTEM_PARTS: Final = ("Answer tersely", "Use British spelling")
(
({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}),
({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}),
({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {}),
({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {"n": 3}),
({"generationConfig": {"candidate_count": 3}}, {"n": 3}),
({}, {}),
),
)