From 3317b2fda9cb1e5fdd495bea272a7853633877ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:14:15 -0700 Subject: [PATCH] fix(proxy): count every Gemini system part and candidateCount in the budget hold A Gemini generateContent body with several systemInstruction text parts lost every part after the first when LiteLLM translated it to a chat request. That dropped the later system text for non-Gemini models behind /v1beta/models/:generateContent, and made the budget hold count only the first part. The adapter now keeps every text part, as a plain string when there is one part and as text blocks when there are several. The output hold now also multiplies by generationConfig.candidateCount, the same way chat bodies use n, so a request asking for several candidates holds budget for all of them. --- .../google_genai/adapters/transformation.py | 10 ++++--- .../spend_tracking/budget_reservation.py | 26 +++++++++++-------- .../google_genai/test_google_genai_adapter.py | 21 +++++++++++++++ .../spend_tracking/test_budget_reservation.py | 18 +++++++++---- 4 files changed, 56 insertions(+), 19 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index e1085dd7857..0363a6d3d1c 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -628,10 +628,14 @@ class GoogleGenAIAdapter: ) -> list[AllMessageValues]: """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_parts[0]["text"]),) - if system_parts and "text" in system_parts[0] - else () + (ChatCompletionSystemMessage(role="system", content=system_content),) 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 cfc865a84d1..5cada93bd33 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -16,6 +16,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.proxy._types import ( @@ -1652,24 +1653,27 @@ def _estimate_output_tokens( _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") +def _nested_field(request_body: Mapping[str, object], config_field: str, field: str) -> object: + config: Final = as_str_mapping(request_body.get(config_field)) + return config.get(field) if config is not None else None + + def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: - inference_config: Final = request_body.get("inferenceConfig") - generation_config: Final = request_body.get("generationConfig") candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), - inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, - generation_config.get("maxOutputTokens") if isinstance(generation_config, Mapping) else None, + _nested_field(request_body, "inferenceConfig", "maxTokens"), + _nested_field(request_body, "generationConfig", "maxOutputTokens"), ) return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) -def _get_output_multiplier(request_body: dict) -> int: - output_multiplier = 1 - for key in ("n", "best_of"): - value = _to_int(request_body.get(key)) - if value is not None: - output_multiplier = max(output_multiplier, value) - return output_multiplier +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"), + ) + return max((1, *(count for count in map(_to_int, candidates) if count is not None))) def _is_input_only_route(route: 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 0292e58c029..7d1e8e2caa4 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -182,6 +182,27 @@ def test_native_function_call_preserves_arguments_and_matching_response(args: ob ] +@pytest.mark.parametrize( + ("system_texts", "expected_content"), + [ + (("Answer tersely",), "Answer tersely"), + ( + ("Answer tersely", "Use British spelling"), + [{"type": "text", "text": "Answer tersely"}, {"type": "text", "text": "Use British spelling"}], + ), + ], +) +def test_every_system_instruction_text_part_reaches_the_system_message( + system_texts: tuple[str, ...], expected_content: str | list[dict[str, str]] +) -> None: + request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion( + model="gemini/gemini-3.8-flash", + contents=[{"role": "user", "parts": [{"text": "hi"}]}], + systemInstruction={"parts": [{"text": text} for text in system_texts]}, + ) + assert request["messages"][0] == {"role": "system", "content": expected_content} + + def test_adapter_import(): """Test that the adapter can be imported successfully""" from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index f1ae7bcb5cd..78431304198 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -177,7 +177,7 @@ def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): GEMINI_PROMPT: Final = "Say hello in five words" -GEMINI_SYSTEM: Final = "Answer tersely" +GEMINI_SYSTEM_PARTS: Final = ("Answer tersely", "Use British spelling") @pytest.mark.parametrize( @@ -190,14 +190,22 @@ GEMINI_SYSTEM: Final = "Answer tersely" ) @pytest.mark.parametrize( ("native_output_config", "chat_output_limit"), - (({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), ({}, {})), + ( + ({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), + ({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}), + ({}, {}), + ), ) +@pytest.mark.parametrize("system_texts", (GEMINI_SYSTEM_PARTS[:1], GEMINI_SYSTEM_PARTS)) def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window( - route: str, native_output_config: dict[str, object], chat_output_limit: dict[str, object] + route: str, + native_output_config: dict[str, object], + chat_output_limit: dict[str, object], + system_texts: tuple[str, ...], ) -> None: native_cost: Final = estimate_request_max_cost( request_body={ - "systemInstruction": {"parts": [{"text": GEMINI_SYSTEM}]}, + "systemInstruction": {"parts": [{"text": text} for text in system_texts]}, "contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}]}], **native_output_config, }, @@ -209,7 +217,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": GEMINI_SYSTEM}, + {"role": "system", "content": [{"type": "text", "text": text} for text in system_texts]}, {"role": "user", "content": GEMINI_PROMPT}, ], **chat_output_limit,