fix(vertex_ai): skip malformed system parts and read Gemini output fields only on Gemini bodies

A systemInstruction part whose text is not a string is skipped, as main did, instead of failing the request. The budget hold reads generationConfig output and candidate fields, and counts contents, only when the body has no messages, prompt, input, query or documents field. An output cap of 0 or below counts as no cap. Test recorders keep tuples instead of mutating lists.
This commit is contained in:
mateo-berri 2026-10-29 01:18:47 -07:00
parent fb75a132e5
commit edfc8744db
7 changed files with 117 additions and 18 deletions

View file

@ -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 ()
)

View file

@ -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:

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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",

View file

@ -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