diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 0363a6d3d1c..594b97138ef 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -629,13 +629,8 @@ class GoogleGenAIAdapter: """Transform Google GenAI contents to OpenAI messages format""" system_parts: Final = system_instruction.get("parts", ()) if system_instruction else () system_texts: Final = tuple(part["text"] for part in system_parts if "text" in part) - system_content: Final = ( - system_texts[0] - if len(system_texts) == 1 - else [ChatCompletionTextObject(type="text", text=text) for text in system_texts] - ) system_messages: Final = ( - (ChatCompletionSystemMessage(role="system", content=system_content),) if system_texts else () + (ChatCompletionSystemMessage(role="system", content="\n".join(system_texts)),) if system_texts else () ) return [*system_messages, *chain.from_iterable(_genai_content_messages(content) for content in contents)] diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 5cada93bd33..d20e2371f09 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -7,6 +7,7 @@ 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 @@ -1651,6 +1652,7 @@ 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: @@ -1658,11 +1660,18 @@ 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"), - _nested_field(request_body, "generationConfig", "maxOutputTokens"), + *_gemini_config_values(request_body, "maxOutputTokens", "max_output_tokens"), ) return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) @@ -1671,7 +1680,7 @@ def _get_output_multiplier(request_body: Mapping[str, object]) -> int: candidates: Final = ( request_body.get("n"), request_body.get("best_of"), - _nested_field(request_body, "generationConfig", "candidateCount"), + *_gemini_config_values(request_body, "candidateCount", "candidate_count"), ) 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 a7e9981ac62..a2559f22892 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -22,6 +22,7 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.responses.litellm_completion_transformation.compaction_items import ( billable_input_since_newest_compaction, holds_compaction_summary, @@ -39,6 +40,16 @@ from litellm.rust_bridge.token_counter import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 _GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction") +_GEMINI_CACHED_CONTENT_FIELDS: Final = ("cachedContent", "cached_content") +_GEMINI_FUNCTION_RESPONSE_FIELDS: Final = ("functionResponse", "function_response") +_GEMINI_TEXT_PART_FIELDS: Final = frozenset({"text", "thought", "thoughtSignature", "thought_signature"}) +_GEMINI_COUNTED_PART_FIELDS: Final = _GEMINI_TEXT_PART_FIELDS | { + "functionCall", + "function_call", + *_GEMINI_FUNCTION_RESPONSE_FIELDS, +} +_GEMINI_COUNTED_ROLES: Final = frozenset({"user", "model"}) +_GEMINI_COUNTED_TOOL_FIELDS: Final = frozenset({"functionDeclarations"}) _GeminiContents: TypeAlias = list[dict[str, object]] | dict[str, object] _GEMINI_CONTENTS: Final = TypeAdapter[_GeminiContents](_GeminiContents) _INPUT_SIZE_FIELDS: Final = ( @@ -182,16 +193,60 @@ def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) - system_instruction: Final = next( (request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None ) + contents: Final = _GEMINI_CONTENTS.validate_python(request_body["contents"]) + tools: Final = request_body.get("tools") + if not _gemini_body_is_countable(request_body, contents, system_instruction, tools): + return None translate: Final = GoogleGenAIAdapter().translate_generate_content_to_completion # pyright: ignore[reportUnknownMemberType] # adapter's **kwargs is untyped - chat_request: Final = translate( - model=model, - contents=_GEMINI_CONTENTS.validate_python(request_body["contents"]), - systemInstruction=system_instruction, - tools=request_body.get("tools"), - ) + chat_request: Final = translate(model=model, contents=contents, systemInstruction=system_instruction, tools=tools) return count_input_tokens_for_model(request_body=chat_request, model=model) +def _gemini_body_is_countable( + request_body: Mapping[str, object], + contents: _GeminiContents, + system_instruction: object, + tools: object, +) -> bool: + content_items: Final = contents if isinstance(contents, list) else (contents,) + system: Final = as_str_mapping(system_instruction) + return ( + not any(field in request_body for field in _GEMINI_CACHED_CONTENT_FIELDS) + and all(_gemini_content_is_countable(content) for content in content_items) + and ( + system_instruction is None + or (system is not None and _gemini_parts_are_countable(system.get("parts", ()), _GEMINI_TEXT_PART_FIELDS)) + ) + and all(_gemini_tool_is_countable(tool) for tool in iter_items(tools or ())) + ) + + +def _gemini_content_is_countable(content: Mapping[str, object]) -> bool: + return content.get("role", "user") in _GEMINI_COUNTED_ROLES and _gemini_parts_are_countable( + content.get("parts", ()), _GEMINI_COUNTED_PART_FIELDS + ) + + +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)) + + +def _gemini_part_is_countable(part: object, fields: frozenset[str]) -> bool: + if isinstance(part, str): + return True + mapping: Final = as_str_mapping(part) + return ( + mapping is not None + and mapping.keys() <= fields + and not any("parts" in (as_str_mapping(mapping.get(field)) or ()) for field in _GEMINI_FUNCTION_RESPONSE_FIELDS) + ) + + +def _gemini_tool_is_countable(tool: object) -> bool: + mapping: Final = as_str_mapping(tool) + return mapping is not None and mapping.keys() <= _GEMINI_COUNTED_TOOL_FIELDS + + def _count_text_tokens(model: str, text: object) -> int: if text is None: return 0 diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 7d1e8e2caa4..e1ffaadb46d 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -186,14 +186,11 @@ def test_native_function_call_preserves_arguments_and_matching_response(args: ob ("system_texts", "expected_content"), [ (("Answer tersely",), "Answer tersely"), - ( - ("Answer tersely", "Use British spelling"), - [{"type": "text", "text": "Answer tersely"}, {"type": "text", "text": "Use British spelling"}], - ), + (("Answer tersely", "Use British spelling"), "Answer tersely\nUse British spelling"), ], ) def test_every_system_instruction_text_part_reaches_the_system_message( - system_texts: tuple[str, ...], expected_content: str | list[dict[str, str]] + system_texts: tuple[str, ...], expected_content: str ) -> None: request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion( model="gemini/gemini-3.8-flash", diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 78431304198..2b654d9085d 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -193,6 +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}), ({}, {}), ), ) @@ -217,7 +218,7 @@ def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window request_body={ "model": "gemini-3.8-flash", "messages": [ - {"role": "system", "content": [{"type": "text", "text": text} for text in system_texts]}, + {"role": "system", "content": "\n".join(system_texts)}, {"role": "user", "content": GEMINI_PROMPT}, ], **chat_output_limit, @@ -230,6 +231,66 @@ def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window assert native_cost == chat_cost +GEMINI_ROUTE: Final = "/v1beta/models/gemini-3.8-flash:generateContent" +GEMINI_TEXT_CONTENT: Final = {"role": "user", "parts": [{"text": GEMINI_PROMPT}]} +GEMINI_AUDIO_PART: Final = {"inlineData": {"mimeType": "audio/wav", "data": "UklGRg" * 40_000}} + + +def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> float | None: + return estimate_request_max_cost( + request_body=body, + route=GEMINI_ROUTE, + llm_router=None, + input_token_counts={} if input_tokens is None else {"gemini-3.8-flash": input_tokens}, + ) + + +@pytest.mark.parametrize( + "body", + ( + {"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": [{"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}]}]}, + { + "contents": [ + { + "role": "user", + "parts": [{"functionResponse": {"name": "listen", "response": {}, "parts": [GEMINI_AUDIO_PART]}}], + } + ] + }, + {"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]}, + ), +) +def test_gemini_body_the_prompt_count_cannot_see_whole_keeps_the_context_window_hold(body: dict[str, object]) -> None: + context_window: Final = litellm.get_model_info("gemini-3.8-flash")["max_input_tokens"] + assert context_window is not None + assert _gemini_hold(body) == _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}, input_tokens=context_window) + + +def test_gemini_tool_turns_and_thought_signatures_count_as_prompt() -> None: + context_window: Final = litellm.get_model_info("gemini-3.8-flash")["max_input_tokens"] + assert context_window is not None + tool_turns: Final = { + "tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "object"}}]}], + "contents": [ + GEMINI_TEXT_CONTENT, + {"role": "model", "parts": [{"functionCall": {"name": "lookup", "args": {}}, "thoughtSignature": "c2ln"}]}, + {"role": "user", "parts": [{"functionResponse": {"name": "lookup", "response": {"answer": "hi"}}}]}, + ], + } + text_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) + tool_turns_hold: Final = _gemini_hold(tool_turns) + context_window_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}, input_tokens=context_window) + assert text_hold is not None and tool_turns_hold is not None and context_window_hold is not None + assert text_hold < tool_turns_hold < context_window_hold + + def _tiered_deployment(input_cost_per_token: float) -> Deployment: return Deployment( model_name="tiered-group",