fix(proxy): read only camelCase Gemini output fields in the budget hold

/v1beta forwards only generationConfig.maxOutputTokens to non-Gemini models, so a snake_case cap the model never receives must not shrink the hold. Snake_case bodies now get the default output hold.
This commit is contained in:
mateo-berri 2026-10-28 19:07:38 -07:00
parent b22e2ca8b6
commit d793d2917f
2 changed files with 3 additions and 12 deletions

View file

@ -7,7 +7,6 @@ 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
@ -1652,7 +1651,6 @@ 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:
@ -1660,18 +1658,11 @@ 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_config_values(request_body: Mapping[str, object], *fields: str) -> tuple[object, ...]:
return tuple(
_nested_field(request_body, config_field, field)
for config_field, field in product(_GEMINI_CONFIG_FIELDS, fields)
)
def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None:
candidates: Final = (
*(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS),
_nested_field(request_body, "inferenceConfig", "maxTokens"),
*_gemini_config_values(request_body, "maxOutputTokens", "max_output_tokens"),
_nested_field(request_body, "generationConfig", "maxOutputTokens"),
)
return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None)
@ -1680,7 +1671,7 @@ def _get_output_multiplier(request_body: Mapping[str, object]) -> int:
candidates: Final = (
request_body.get("n"),
request_body.get("best_of"),
*_gemini_config_values(request_body, "candidateCount", "candidate_count"),
_nested_field(request_body, "generationConfig", "candidateCount"),
)
return max((1, *(count for count in map(_to_int, candidates) if count is not None)))

View file

@ -193,7 +193,7 @@ 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}}, {"max_tokens": 5, "n": 3}),
({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {}),
({}, {}),
),
)