Merge pull request #47446 from BerriAI/litellm_backport_google_count_shapes_stable_1_104_x

fix(count_tokens): restore Google request compatibility on stable/1.104.x

litellm/proxy/google_endpoints/endpoints.py

@@ -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 *
@@ -31,6 +31,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)
@@ -51,7 +56,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(
@@ -111,7 +116,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_google_gemini_proxy_request.py

@@ -9,11 +9,12 @@ The request payload is correctly processed and forwarded to the httpx client.
 
 import json
 import unittest.mock
-from typing import Optional
+from typing import Final, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
+from pydantic import JsonValue, ValidationError
 
 # Add the parent directory to the system path
 
@@ -493,3 +494,39 @@ async def test_gemini_proxy_config_with_custom_api_base():
 if __name__ == "__main__":
     # Run the tests
     pytest.main([__file__, "-v"])
+
+
+@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:
+    from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
+
+    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:
+    from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
+
+    with pytest.raises(ValidationError):
+        _google_count_tokens_request(
+            "counted",
+            {"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}},
+        )