diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index d46dae47f06..40ab377dd5e 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/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, ...]: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 566fbc886ad..352a166fae0 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/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 diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 30de3e6f596..dbf90aadfa0 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/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: diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 1f581a05031..d622a1b85b9 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/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( diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 12d509d5198..a719dbca091 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/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 diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index 72c9ff4ae55..6800dd8f91b 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/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