fix(count_tokens): normalize Google system parts and null tools

This commit is contained in:
mateo-berri 2026-10-19 16:52:43 -07:00
parent 497f31d43a
commit a9571f3104
2 changed files with 41 additions and 4 deletions

View file

@ -5,7 +5,7 @@ from typing import Final
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse, ORJSONResponse
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator
from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types
from litellm.proxy._types import *
@ -34,6 +34,11 @@ class _GoogleCountTokensSystemInstruction(BaseModel):
model_config = ConfigDict(frozen=True)
parts: tuple[_GoogleCountTokensTextPart, ...]
@field_validator("parts", mode="before")
@classmethod
def normalize_parts(cls, parts: object) -> object:
return (parts,) if isinstance(parts, Mapping) else parts
class _GoogleCountTokensFunctionDeclaration(BaseModel):
model_config = ConfigDict(frozen=True, populate_by_name=True)
@ -54,7 +59,7 @@ class _GoogleCountTokensRequest(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True)
contents: tuple[Mapping[str, JsonValue], ...] = ()
system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction")
tools: tuple[_GoogleCountTokensTool, ...] = ()
tools: tuple[_GoogleCountTokensTool, ...] | None = ()
router: Final = APIRouter(
@ -114,7 +119,7 @@ def _google_count_tokens_request(model_name: str, data: object) -> TokenCountReq
if validated_data.system_instruction is not None
else None
)
declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools)
declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools or ())
tools: Final = [
{
"name": declaration.name,

View file

@ -6,7 +6,7 @@ from typing import Final
import pytest
from dotenv import load_dotenv
from pydantic import JsonValue
from pydantic import JsonValue, ValidationError
from starlette.requests import Request
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens
@ -104,3 +104,35 @@ async def test_proxy_gemini_to_openai_like_model_token_counting():
)
assert response.get("totalTokens") > 0
@pytest.mark.parametrize("system_field", ("systemInstruction", "system_instruction"))
@pytest.mark.parametrize("parts", ({"text": "Follow these rules"}, [{"text": "Follow these rules"}]))
@pytest.mark.parametrize("tools", (None, []))
def test_google_count_tokens_accepts_equivalent_system_parts_and_no_tools(
system_field: str, parts: JsonValue, tools: JsonValue
) -> None:
contents: Final = [{"role": "user", "parts": [{"text": "Count this request"}]}]
request: Final = _google_count_tokens_request(
"counted", {"contents": contents, system_field: {"parts": parts}, "tools": tools}
)
assert request.messages == [{"role": "user", "content": "Count this request"}]
assert request.system == [{"type": "text", "text": "Follow these rules"}]
assert request.tools is None
assert request.contents == contents
@pytest.mark.parametrize(
"parts",
(
{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}},
[{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}],
),
)
def test_google_count_tokens_rejects_media_only_system_parts(parts: JsonValue) -> None:
with pytest.raises(ValidationError):
_google_count_tokens_request(
"counted",
{"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}},
)