diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 594b97138ef..f624f69567e 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -9,6 +9,7 @@ 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 @@ -244,6 +245,18 @@ 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 + ) + + def _genai_content_messages(content: object) -> tuple[AllMessageValues, ...]: fields: Final = _validated(_CONFIG_FIELDS, content) if fields is None: @@ -627,8 +640,7 @@ class GoogleGenAIAdapter: system_instruction: _GenAISystemInstruction | None = None, ) -> 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_texts: Final = _system_texts(system_instruction.get("parts") if system_instruction else None) system_messages: Final = ( (ChatCompletionSystemMessage(role="system", content="\n".join(system_texts)),) if system_texts else () ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4efcc06fdce..b1746ca830b 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -41,7 +41,11 @@ from litellm.proxy.common_utils.user_api_key_cache import ( tag_cache_key, team_membership_reservation_cache_key, ) -from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model +from litellm.proxy.spend_tracking.input_tokens import ( + count_input_tokens, + count_input_tokens_for_model, + is_gemini_contents_body, +) from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -1661,6 +1665,8 @@ 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")) @@ -1668,12 +1674,17 @@ def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object 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 + ) 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_output_cap, ) - return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), 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: diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index a2559f22892..452a12784b4 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -39,6 +39,7 @@ from litellm.rust_bridge.token_counter import ( 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") @@ -169,7 +170,7 @@ 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 "contents" in request_body: + 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")) @@ -189,6 +190,10 @@ 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 diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index e1ffaadb46d..ae04719f4a0 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -200,6 +200,27 @@ def test_every_system_instruction_text_part_reaches_the_system_message( assert request["messages"][0] == {"role": "system", "content": expected_content} +@pytest.mark.parametrize( + ("system_parts", "expected_messages"), + [ + ( + [{"text": "Answer tersely"}, {"text": None}, {"text": 3}, "text", {"text": "Use British spelling"}], + [{"role": "system", "content": "Answer tersely\nUse British spelling"}, {"role": "user", "content": "hi"}], + ), + (None, [{"role": "user", "content": "hi"}]), + ], +) +def test_a_system_instruction_part_without_string_text_is_skipped( + system_parts: object, expected_messages: 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": system_parts}, + ) + assert request["messages"] == expected_messages + + 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/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index ee069c01c04..45b262b250d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -8282,10 +8282,10 @@ class TestAzureSpeechRawBodyThroughRealAuth: class _CountingCustomAuth: def __init__(self, only_accepts: str | None = None) -> None: self.only_accepts: Final = only_accepts - self.api_keys_seen: list[str] = [] + self.api_keys_seen: tuple[str, ...] = () async def __call__(self, request: Request, api_key: str) -> UserAPIKeyAuth: - self.api_keys_seen.append(api_key) + self.api_keys_seen = (*self.api_keys_seen, api_key) if self.only_accepts is not None and api_key != self.only_accepts: raise ValueError("unknown caller") return UserAPIKeyAuth(user_id=f"custom-auth-caller-{len(self.api_keys_seen)}") @@ -8294,12 +8294,12 @@ class _CountingCustomAuth: class _ServedCallerRecorder(CustomLogger): def __init__(self) -> None: super().__init__() - self.served_user_ids: list[str | None] = [] + self.served_user_ids: tuple[str | None, ...] = () async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral ) -> dict: - self.served_user_ids.append(user_api_key_dict.user_id) + self.served_user_ids = (*self.served_user_ids, user_api_key_dict.user_id) return data @@ -8378,8 +8378,8 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: ) assert response.status_code == 200, response.text - assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] - assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) + assert recorder.served_user_ids == ("custom-auth-caller-1",) assert vertex.call_count == 1 def test_discovery_search_runs_custom_auth_once_and_is_served_as_its_caller( @@ -8400,8 +8400,8 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: ) assert (response.status_code, response.json()) == (200, {"results": []}) - assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] - assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) + assert recorder.served_user_ids == ("custom-auth-caller-1",) assert discovery.call_count == 1 @pytest.mark.parametrize("scheme", ["bearer", "Basic"]) @@ -8445,8 +8445,8 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: ) assert response.status_code == 200, response.text - assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] - assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) + assert recorder.served_user_ids == ("custom-auth-caller-1",) assert google.call_count == 1 def test_x_litellm_api_key_authenticates_once_and_is_never_forwarded_without_a_vertex_credential( @@ -8465,7 +8465,7 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: ) assert response.status_code == 401, response.text - assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) assert not google.called @pytest.mark.parametrize("header", ["x-goog-api-key", "x-api-key", "API-Key", "Ocp-Apim-Subscription-Key"]) @@ -8482,7 +8482,7 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: assert response.status_code == 401, response.text assert self.VIRTUAL_KEY not in custom_auth.api_keys_seen - assert recorder.served_user_ids == [] + assert recorder.served_user_ids == () assert not google.called diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 01acfbf9a61..42ae046ec72 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -292,6 +292,48 @@ def test_gemini_tool_turns_and_thought_signatures_count_as_prompt() -> None: assert text_hold < tool_turns_hold < context_window_hold +GEMINI_CHAT_BODY: Final = { + "model": "gemini-3.8-flash", + "messages": [{"role": "user", "content": GEMINI_PROMPT}], +} + + +def _chat_hold(body: dict[str, object]) -> float | None: + return estimate_request_max_cost( + request_body=body, + route="/chat/completions", + llm_router=None, + input_token_counts={}, + ) + + +@pytest.mark.parametrize( + "gemini_output_fields", + ( + {"generationConfig": {"maxOutputTokens": 5}}, + {"generationConfig": {"candidateCount": 3}}, + {"generation_config": {"candidate_count": 3}}, + {"contents": [], "generationConfig": {"maxOutputTokens": 5}}, + ), +) +def test_gemini_output_fields_on_a_chat_body_leave_its_hold_alone(gemini_output_fields: dict[str, object]) -> None: + plain_hold: Final = _chat_hold(GEMINI_CHAT_BODY) + assert plain_hold is not None + assert _chat_hold({**GEMINI_CHAT_BODY, **gemini_output_fields}) == plain_hold + + +@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 + assert ( + _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": cap}}) + == uncapped_gemini_hold + ) + assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": cap}) == uncapped_chat_hold + + def _tiered_deployment(input_cost_per_token: float) -> Deployment: return Deployment( model_name="tiered-group", diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index c7f13651471..72c9ff4ae55 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -110,3 +110,11 @@ def test_a_compaction_item_counts_only_its_summary(model: str) -> None: assert count_input_tokens_for_model(request_body=padded_body, model=model) == count_input_tokens_for_model( request_body=plain_body, model=model ) + + +@pytest.mark.parametrize("input_field", ("prompt", "input", "query")) +def test_a_stray_contents_field_never_replaces_the_request_input_count(input_field: str) -> None: + body: Final = {"model": CL100K_MODEL, input_field: "count every one of these words " * 200} + 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