mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
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:
parent
fb75a132e5
commit
edfc8744db
7 changed files with 117 additions and 18 deletions
|
|
@ -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 ()
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue