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

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)))
 

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

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: