Merge fbf50f2c07 into e16d79d635

litellm/proxy/google_endpoints/endpoints.py

@@ -1,11 +1,11 @@
 from collections.abc import Mapping
 from itertools import chain
 from types import MappingProxyType
-from typing import Final
+from typing import Final, cast
 
 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 (cast(object, 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,

tests/unit/proxy/google_endpoints/test_endpoints.py

@@ -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}},
+        )