From 84a4505d4248c0eddab99800bec8823a9e4c4519 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 04:13:55 -0700 Subject: [PATCH] fix(proxy): take the budget hold's request shape from the route On generateContent routes the hold counts only contents and reads its output cap and multiplier from generationConfig, so stray fields of other APIs can no longer hide the prompt. A max_tokens of 0 holds no output again, while negative values and a Gemini maxOutputTokens of 0 or below hold the default output --- .../spend_tracking/budget_reservation.py | 33 +++++----- litellm/proxy/spend_tracking/input_tokens.py | 31 +++++---- .../spend_tracking/test_budget_reservation.py | 66 +++++++++++++++++-- 3 files changed, 95 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 55e59c0f8b0..566fbc886ad 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -44,7 +44,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.spend_tracking.input_tokens import ( count_input_tokens, count_input_tokens_for_model, - is_gemini_contents_body, tokenizer_models, ) from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope @@ -1441,7 +1440,7 @@ def _max_cost_for_cost_info( if estimated_input_tokens is None or output_tokens is None: return None - output_multiplier: Final = _get_output_multiplier(request_body=request_body) + output_multiplier: Final = _get_output_multiplier(request_body=request_body, route=route) tiered_pricing: Final = model_info.get("tiered_pricing") if isinstance(tiered_pricing, list) and tiered_pricing: tier: Final = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=estimated_input_tokens) @@ -1611,6 +1610,7 @@ async def count_request_input_tokens( request_body=request_body, raw_body=raw_body, models=tuple(dict.fromkeys(chain.from_iterable(counted_as.values()))), + gemini_route=RouteChecks.is_generate_content_route(route), ) return MappingProxyType( { @@ -1631,7 +1631,9 @@ def _estimate_input_tokens( counted: Final = ( input_tokens if input_tokens is not None - else count_input_tokens_for_model(request_body=request_body, model=model) + else count_input_tokens_for_model( + request_body=request_body, model=model, gemini_route=RouteChecks.is_generate_content_route(route) + ) ) if counted is not None: return counted @@ -1654,7 +1656,7 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 - requested: Final = _requested_output_tokens(request_body) + requested: Final = _requested_output_tokens(request_body, route) # Clamp at min(requested-or-default, model_max-or-default). Two purposes: # (1) Without an explicit cap we still need a finite reservation so the @@ -1678,33 +1680,28 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field: def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object, ...]: - if not is_gemini_contents_body(request_body): - return () 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: - gemini_output_cap: Final = ( - _nested_field(request_body, "generationConfig", "maxOutputTokens") - if is_gemini_contents_body(request_body) - else None - ) +def _requested_output_tokens(request_body: Mapping[str, object], route: str) -> int | None: + if RouteChecks.is_generate_content_route(route): + gemini_cap: Final = _to_int(_nested_field(request_body, "generationConfig", "maxOutputTokens")) + return gemini_cap if gemini_cap is not None and gemini_cap > 0 else None candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), _nested_field(request_body, "inferenceConfig", "maxTokens"), - gemini_output_cap, ) - return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens > 0), None) + return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens >= 0), None) -def _get_output_multiplier(request_body: Mapping[str, object]) -> int: +def _get_output_multiplier(request_body: Mapping[str, object], route: str) -> int: candidates: Final = ( - request_body.get("n"), - request_body.get("best_of"), - *_gemini_candidate_counts(request_body), + _gemini_candidate_counts(request_body) + if RouteChecks.is_generate_content_route(route) + else (request_body.get("n"), request_body.get("best_of")) ) return max((1, *(count for count in map(_to_int, candidates) if count is not None))) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 7f33cda7592..30de3e6f596 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -44,7 +44,6 @@ if TYPE_CHECKING: TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 -_NON_GEMINI_INPUT_FIELDS: Final = ("messages", "prompt", "input", "query", "documents") _GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction") _GEMINI_CACHED_CONTENT_FIELDS: Final = ("cachedContent", "cached_content") _GEMINI_FUNCTION_RESPONSE_FIELDS: Final = ("functionResponse", "function_response") @@ -103,6 +102,7 @@ async def count_input_tokens( request_body: dict[str, object], raw_body: bytes | None, models: Sequence[str], + gemini_route: bool = False, ) -> Mapping[str, int]: """Input-token count per model, sharing one native count across models that select the same tokenizer.""" @@ -110,13 +110,14 @@ async def count_input_tokens( (model, rust_tokenizer(model)) for model in models ) groups: Final[tuple[RustTokenizer | None, ...]] = tuple(dict.fromkeys(tokenizer for _, tokenizer in tokenizers)) - counted_raw_body: Final = None if holds_compaction_summary(request_body.get("input")) else raw_body + counted_raw_body: Final = None if gemini_route or holds_compaction_summary(request_body.get("input")) else raw_body group_counts: Final = [ await _count_group( request_body=request_body, raw_body=counted_raw_body, tokenizer=tokenizer, models=tuple(model for model, selected in tokenizers if selected == tokenizer), + gemini_route=gemini_route, ) for tokenizer in groups ] @@ -130,14 +131,16 @@ async def _count_group( raw_body: bytes | None, tokenizer: RustTokenizer | None, models: tuple[str, ...], + gemini_route: bool, ) -> Mapping[str, int]: async def python() -> Mapping[str, int]: if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: - return _count_input_tokens_for_models(request_body=request_body, models=models) + return _count_input_tokens_for_models(request_body=request_body, models=models, gemini_route=gemini_route) return await asyncio.to_thread( _count_input_tokens_for_models, request_body=request_body, models=models, + gemini_route=gemini_route, ) if tokenizer is None or raw_body is None: @@ -172,18 +175,30 @@ async def _native_counts( def _count_input_tokens_for_models( request_body: dict[str, object], models: Sequence[str], + gemini_route: bool, ) -> Mapping[str, int]: return MappingProxyType( { model: tokens for model in models - if (tokens := count_input_tokens_for_model(request_body=request_body, model=model)) is not None + if ( + tokens := count_input_tokens_for_model( + request_body=request_body, model=model, gemini_route=gemini_route + ) + ) + is not None } ) -def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: +def count_input_tokens_for_model(request_body: dict, model: str, gemini_route: bool = False) -> int | None: try: + if gemini_route: + return ( + _count_gemini_input_tokens(request_body=request_body, model=model) + if "contents" in request_body + else None + ) if "messages" in request_body: try: return litellm.token_counter( @@ -194,8 +209,6 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: ) except ValueError: return _count_text_tokens(model=model, text=request_body.get("messages")) - if is_gemini_contents_body(request_body): - return _count_gemini_input_tokens(request_body=request_body, model=model) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -214,10 +227,6 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: return None -def is_gemini_contents_body(request_body: Mapping[str, object]) -> bool: - return "contents" in request_body and not any(field in request_body for field in _NON_GEMINI_INPUT_FIELDS) - - def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) -> int | None: system_instruction: Final = next( (request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index f5e6fa9af48..12d509d5198 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import math from datetime import datetime, timedelta, timezone from typing import Final @@ -25,6 +26,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, + estimate_request_input_cost, estimate_request_max_cost, get_budget_window_start, release_unbound_budget_reservation, @@ -324,15 +326,67 @@ def test_gemini_output_fields_on_a_chat_body_leave_its_hold_alone(gemini_output_ @pytest.mark.parametrize("cap", (0, -1000)) -def test_an_output_cap_of_zero_or_below_holds_the_default_output(cap: int) -> None: - uncapped_gemini_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) - uncapped_chat_hold: Final = _chat_hold(GEMINI_CHAT_BODY) - assert uncapped_gemini_hold is not None and uncapped_chat_hold is not None +def test_a_gemini_output_cap_of_zero_or_below_holds_the_default_output(cap: int) -> None: + uncapped_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) + assert uncapped_hold is not None assert ( _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": cap}}) - == uncapped_gemini_hold + == uncapped_hold ) - assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": cap}) == uncapped_chat_hold + + +def test_a_negative_max_tokens_holds_the_default_output() -> None: + uncapped_hold: Final = _chat_hold(GEMINI_CHAT_BODY) + assert uncapped_hold is not None + assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": -1000}) == uncapped_hold + + +def test_a_text_completion_with_max_tokens_zero_holds_only_its_prompt() -> None: + body: Final = {"model": "gemini-3.8-flash", "prompt": GEMINI_PROMPT * 50, "max_tokens": 0} + prompt_cost: Final = estimate_request_input_cost(request_body=body, route="/v1/completions", llm_router=None) + assert prompt_cost is not None and prompt_cost > 0 + assert estimate_request_max_cost(request_body=body, route="/v1/completions", llm_router=None) == prompt_cost + + +GEMINI_LONG_CONTENTS: Final = {"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT * 500}]}]} + + +@pytest.mark.parametrize( + "stray_fields", + ( + {"prompt": ""}, + {"messages": [], "n": 5}, + {"input": "", "max_tokens": 1}, + {"query": "", "documents": [], "max_completion_tokens": 1}, + {"inferenceConfig": {"maxTokens": 1}}, + ), +) +def test_fields_of_other_apis_on_a_gemini_route_leave_its_hold_alone(stray_fields: dict[str, object]) -> None: + body: Final = {**GEMINI_LONG_CONTENTS, "generationConfig": {"maxOutputTokens": 4000, "candidateCount": 2}} + gemini_hold: Final = _gemini_hold(body) + assert gemini_hold is not None + assert _gemini_hold({**body, **stray_fields}) == gemini_hold + + +@pytest.mark.parametrize("stray_input", ({"prompt": ""}, {"messages": []}, {"input": ""})) +@pytest.mark.parametrize("model", ("gpt-4", "claude-sonnet-5-5")) +@pytest.mark.asyncio +async def test_a_stray_input_field_on_a_gemini_route_never_hides_the_counted_contents( + model: str, stray_input: dict[str, object] +) -> None: + async def counted(body: dict[str, object]) -> dict[str, int]: + return dict( + await count_request_input_tokens( + request_body=body, + route=f"/v1beta/models/{model}:generateContent", + llm_router=None, + raw_body=json.dumps(body).encode(), + ) + ) + + contents_counts: Final = await counted(GEMINI_LONG_CONTENTS) + assert contents_counts[model] > 1000 + assert await counted({**GEMINI_LONG_CONTENTS, **stray_input}) == contents_counts def _tiered_deployment(input_cost_per_token: float) -> Deployment: