fix(proxy): hold one candidate for a count too large to price and read system parts like content parts

litellm/google_genai/adapters/transformation.py

@@ -9,7 +9,6 @@ from typing_extensions import ReadOnly, Required, TypedDict
 
 from litellm import verbose_logger
 from litellm.exceptions import BadRequestError
-from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
 from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params
 from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types, normalize_tool_schema
@@ -250,16 +249,8 @@ def _genai_tool_call(part: object) -> ChatCompletionAssistantToolCall | None:
     )
 
 
-def _system_part_text(part: object) -> str | None:
-    fields: Final = as_str_mapping(part)
-    text: Final = fields.get("text") if fields is not None else None
-    return text if isinstance(text, str) else None
-
-
 def _system_texts(parts: object) -> tuple[str, ...]:
-    return tuple(
-        text for text in map(_system_part_text, parts if isinstance(parts, Sequence) else ()) if text is not None
-    )
+    return tuple(part["text"] for part in _genai_message_parts(_genai_content_parts(parts)) if part["type"] == "text")
 
 
 def _genai_content_messages(content: object) -> tuple[AllMessageValues, ...]:

litellm/proxy/spend_tracking/budget_reservation.py

@@ -1449,8 +1449,8 @@ def _max_cost_for_cost_info(
                 tier_rate(tier, "output_cost_per_token"),
                 tier_rate(tier, "output_cost_per_reasoning_token"),
             )
-            return (estimated_input_tokens * tier_rate(tier, "input_cost_per_token")) + (
-                output_tokens * output_multiplier * output_rate
+            return (estimated_input_tokens * tier_rate(tier, "input_cost_per_token")) + _output_cost(
+                output_tokens, output_multiplier, output_rate
             )
 
     input_cost_per_token: Final = _to_float(model_info.get("input_cost_per_token"))
@@ -1467,7 +1467,7 @@ def _max_cost_for_cost_info(
     # under-reserving reasoning-heavy requests.
     output_rate = max(output_cost_per_token or 0.0, output_cost_per_reasoning_token or 0.0)
     if output_cost_per_token is not None or output_cost_per_reasoning_token is not None:
-        cost += output_tokens * output_multiplier * output_rate
+        cost += _output_cost(output_tokens, output_multiplier, output_rate)
     elif output_tokens > 0:
         return None
 
@@ -1706,6 +1706,13 @@ def _get_output_multiplier(request_body: Mapping[str, object], route: str) -> in
     return max((1, *(count for count in map(_to_int, candidates) if count is not None)))
 
 
+def _output_cost(output_tokens: int, output_multiplier: int, output_rate: float) -> float:
+    try:
+        return output_tokens * output_multiplier * output_rate
+    except OverflowError:
+        return output_tokens * output_rate
+
+
 def _is_input_only_route(route: str) -> bool:
     return any(
         route_part in route
@@ -1731,7 +1738,7 @@ def _to_int(value: object) -> int | None:
         return None
     try:
         return int(value)
-    except (TypeError, ValueError):
+    except (TypeError, ValueError, OverflowError):
         return None
 
 

litellm/proxy/spend_tracking/input_tokens.py

@@ -266,7 +266,9 @@ def _gemini_content_is_countable(content: Mapping[str, object]) -> bool:
 
 
 def _gemini_parts_are_countable(parts: object, fields: frozenset[str]) -> bool:
-    return all(_gemini_part_is_countable(part, fields) for part in iter_items(parts))
+    single_part: Final = as_str_mapping(parts)
+    items: Final = (single_part,) if single_part is not None else iter_items(parts)
+    return all(_gemini_part_is_countable(part, fields) for part in items)
 
 
 def _gemini_part_is_countable(part: object, fields: frozenset[str]) -> bool:

tests/unit/google_genai/test_google_genai_adapter.py

@@ -215,14 +215,31 @@ def test_every_system_instruction_text_part_reaches_the_system_message(
 @pytest.mark.parametrize(
     ("system_parts", "expected_messages"),
     [
-        (
-            [{"text": "Answer tersely"}, {"text": None}, {"text": 3}, "text", {"text": "Use British spelling"}],
+        pytest.param(
+            [
+                {"text": "Answer tersely"},
+                {"text": None},
+                {"text": 3},
+                {"text": ["x"]},
+                {"text": "Use British spelling"},
+            ],
             [{"role": "system", "content": "Answer tersely\nUse British spelling"}, {"role": "user", "content": "hi"}],
+            id="malformed_parts_skipped_like_contents",
         ),
-        (None, [{"role": "user", "content": "hi"}]),
+        pytest.param(
+            {"text": "Answer tersely"},
+            [{"role": "system", "content": "Answer tersely"}, {"role": "user", "content": "hi"}],
+            id="single_object_parts",
+        ),
+        pytest.param(
+            ["Answer tersely", {"text": "Use British spelling"}],
+            [{"role": "system", "content": "Answer tersely\nUse British spelling"}, {"role": "user", "content": "hi"}],
+            id="bare_string_part_read_like_contents",
+        ),
+        pytest.param(None, [{"role": "user", "content": "hi"}], id="no_parts"),
     ],
 )
-def test_a_system_instruction_part_without_string_text_is_skipped(
+def test_system_instruction_parts_follow_the_contents_part_rules(
     system_parts: object, expected_messages: list[dict[str, str]]
 ) -> None:
     request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion(

tests/unit/proxy/spend_tracking/test_budget_reservation.py

@@ -254,6 +254,7 @@ def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> fl
     (
         {"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}, GEMINI_AUDIO_PART]}]},
         {"contents": [{"role": "user", "parts": [{"fileData": {"mimeType": "video/mp4", "fileUri": "gs://b/v.mp4"}}]}]},
+        {"contents": [{"role": "user", "parts": {"fileData": {"mimeType": "video/mp4", "fileUri": "gs://b/v.mp4"}}}]},
         {"contents": [{"role": "user", "parts": [{"inline_data": {"mime_type": "image/png", "data": "iVBORw0K"}}]}]},
         {"contents": [{"role": "model", "parts": [{"executableCode": {"language": "PYTHON", "code": "print(1)"}}]}]},
         {"contents": [GEMINI_TEXT_CONTENT, {"role": "function", "parts": [{"text": GEMINI_PROMPT * 500}]}]},
@@ -266,6 +267,7 @@ def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> fl
             ]
         },
         {"systemInstruction": {"parts": [GEMINI_AUDIO_PART]}, "contents": [GEMINI_TEXT_CONTENT]},
+        {"systemInstruction": {"parts": GEMINI_AUDIO_PART}, "contents": [GEMINI_TEXT_CONTENT]},
         {"tools": [{"function_declarations": [{"name": "lookup"}]}], "contents": [GEMINI_TEXT_CONTENT]},
         {"tools": [{"googleSearch": {}}], "contents": [GEMINI_TEXT_CONTENT]},
         {"cachedContent": "cachedContents/abc", "contents": [GEMINI_TEXT_CONTENT]},
@@ -335,6 +337,26 @@ def test_a_gemini_output_cap_of_zero_or_below_holds_the_default_output(cap: int)
     )
 
 
+UNPRICEABLE_COUNTS: Final = (10**400, "9" * 400, 1e308, math.inf)
+
+
+@pytest.mark.parametrize("count", UNPRICEABLE_COUNTS)
+def test_a_candidate_count_too_large_to_price_holds_one_candidate(count: object) -> None:
+    one_candidate_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]})
+    assert one_candidate_hold is not None
+    assert (
+        _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": count}})
+        == one_candidate_hold
+    )
+
+
+@pytest.mark.parametrize("count", UNPRICEABLE_COUNTS)
+def test_an_n_too_large_to_price_holds_one_choice(count: object) -> None:
+    one_choice_hold: Final = _chat_hold(GEMINI_CHAT_BODY)
+    assert one_choice_hold is not None
+    assert _chat_hold({**GEMINI_CHAT_BODY, "n": count}) == one_choice_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

tests/unit/proxy/spend_tracking/test_input_tokens.py

@@ -118,3 +118,24 @@ def test_a_stray_contents_field_never_replaces_the_request_input_count(input_fie
     count: Final = count_input_tokens_for_model(request_body=body, model=CL100K_MODEL)
     assert count is not None and count > 1000
     assert count_input_tokens_for_model(request_body={**body, "contents": []}, model=CL100K_MODEL) == count
+
+
+GEMINI_TEXT: Final = "count every one of these words " * 200
+GEMINI_GREETING: Final = {"role": "user", "parts": [{"text": "hi"}]}
+
+
+def _gemini_count(body: dict[str, object], model: str) -> int | None:
+    return count_input_tokens_for_model(request_body=body, model=model, gemini_route=True)
+
+
+@pytest.mark.parametrize("parts", ([GEMINI_TEXT], {"text": GEMINI_TEXT}), ids=("bare_string", "single_object"))
+@pytest.mark.parametrize("model", (CL100K_MODEL, "claude-sonnet-5-5"))
+def test_gemini_parts_count_the_same_in_every_shape_the_adapter_reads(model: str, parts: object) -> None:
+    listed_contents: Final = _gemini_count({"contents": [{"role": "user", "parts": [{"text": GEMINI_TEXT}]}]}, model)
+    listed_system: Final = _gemini_count(
+        {"systemInstruction": {"parts": [{"text": GEMINI_TEXT}]}, "contents": [GEMINI_GREETING]}, model
+    )
+    assert listed_contents is not None and listed_contents > 1000
+    assert listed_system is not None and listed_system > 1000
+    assert _gemini_count({"contents": [{"role": "user", "parts": parts}]}, model) == listed_contents
+    assert _gemini_count({"systemInstruction": {"parts": parts}, "contents": [GEMINI_GREETING]}, model) == listed_system