diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 5fbe780379e..5dbcbcba8e1 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/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, diff --git a/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py index 6f8f90efc73..c465b6167e6 100644 --- a/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py +++ b/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}}, + )