tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-16 18:16:34 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
f5b11b72a6dcc8f4e7a16f00a61bdd124cc171d2
litellm/tests/unified_google_tests/test_google_ai_studio.py
T
Sameer KankuteGitHubCursormateo-berriClaude
 95e3d136e1 test(google): add google-genai SDK proxy integration tests (#29781) 
* test(google): add google-genai SDK proxy integration tests for Gemini and Vertex

Pin google-genai in the CI dependency group and exercise streaming/non-streaming
generate_content through the LiteLLM proxy in the existing unified_google_tests suite.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(test): address Greptile review for google-genai proxy SDK tests

Restore GOOGLE_APPLICATION_CREDENTIALS after the module proxy fixture tears down,
initialize temp-file tracking on the proxy SDK base class, and skip litellm reload
for proxy_genai_sdk tests so the module-scoped proxy server stays consistent.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(test): only load Vertex credentials when keys exist for proxy SDK tests

Avoid writing empty GOOGLE_APPLICATION_CREDENTIALS temp files so Vertex tests
skip cleanly without credentials, use a session-scoped proxy fixture, and clean up
per-test credential temp files.

Co-authored-by: Cursor <cursoragent@cursor.com>

* chore(test): scope google-genai pin to unified_google_tests only

Remove google-genai from the ci dependency group and pin it in
tests/unified_google_tests/requirements.txt for local test installs.

Co-authored-by: Cursor <cursoragent@cursor.com>

* test(google): tie litellm reload skip to proxy fixture dependency

Replace the name-based reload guard with a check on whether the test
requests the google_genai_proxy_url fixture, so the skip stays correct
if the proxy SDK tests are renamed.

* fix(test): stop DatabaseURLSettings tests leaking DATABASE_URL into os.environ

The autouse env scrubber relied on monkeypatch.delenv, but apply_to_env
writes DATABASE_URL straight into os.environ, which monkeypatch never
tracks and therefore never undoes. The synthesized writer.example.com URL
leaked past the last test in this module and into proxy-infra tests that
read DATABASE_URL to decide whether to hit a real database, e.g.
test_deprecated_key_grace_period_cache_hit_path, turning an intended skip
into a ConnectError. Snapshot and restore the managed vars directly so the
original environment is reinstated regardless of how it was mutated.

* test(google): drop redundant per-test vertex credential setup

The session-scoped google_genai_proxy_url fixture already configures
GOOGLE_APPLICATION_CREDENTIALS before the proxy starts, and
_require_proxy_sdk skips when credentials are missing, so the per-test
_setup_vertex_credentials_if_needed helper and its temp-file tracking
never did any work. Remove it to keep the ABC self-contained.

* test(google): declare model_config contract on proxy SDK ABC

_skip_reason_if_credentials_missing reads self.model_config to pick the
provider, but that property was only declared on the sibling
BaseGoogleGenAITest. Make the dependency explicit by adding model_config
as an abstract property on BaseGoogleGenAIProxySDKTest so the ABC is
self-contained and a standalone subclass fails fast instead of hitting an
AttributeError.

* test(google): narrow streaming error catch to Exception

Catching BaseException in the streaming assertion swallowed
KeyboardInterrupt and SystemExit, turning a Ctrl-C into a test failure
message instead of letting pytest interrupt cleanly. Only genuine runtime
errors should be recorded as stream failures, so catch Exception.

* test(google): initialize proxy on the same loop that serves it

The proxy was initialized via asyncio.run() on the main thread, which
creates and tears down a throwaway event loop, while requests were served
on a separate loop in the worker thread. Any asyncio primitive bound to
the init loop would be unusable once serving started. Run initialize()
on the worker thread's loop right before server.serve() so setup and
request handling share a single event loop.

* test(google): drop redundant google-genai requirements pin

google-genai>=1.37.0,<2.0 is already declared in the proxy-runtime extra,
which the google_generate_content_endpoint_testing CI job installs via
uv sync --all-extras. The standalone tests/unified_google_tests/requirements.txt
duplicated that pin with a narrower ==1.37.0 specifier and was never
installed by CI, so it added a second source of truth without changing
what gets installed. Drop it and rely on the proxy-runtime extra.

* chore: revert incidental uv.lock exclude-newer bump

The google-genai ci pin was added and then dropped (it is already
provided by the proxy-runtime group), but each uv lock recomputed the
relative exclude-newer span, leaving only a timestamp bump in uv.lock.
Restore it to the base value so this test-only PR carries no lockfile
change.

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-06-05 21:05:32 +00:00

482 lines
19 KiB
Python
Raw Blame History

from base_google_genai_proxy_sdk_test import BaseGoogleGenAIProxySDKTest
from base_google_test import BaseGoogleGenAITest
import sys
import os
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import pytest
import litellm
import unittest.mock
import json
class TestGoogleGenAIStudio(BaseGoogleGenAITest, BaseGoogleGenAIProxySDKTest):
"""Test Google GenAI Studio"""
@property
def model_config(self):
return {
"model": "gemini/gemini-2.5-flash-lite",
}
@property
def proxy_model_name(self) -> str:
return "gemini-2.5-flash-lite"
@pytest.mark.asyncio
async def test_mock_stream_generate_content_with_tools():
"""Test streaming function call response parsing and validation"""
from litellm.types.google_genai.main import ToolConfigDict
litellm._turn_on_debug()
contents = [
{
"role": "user",
"parts": [
{
"text": "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning"
}
],
}
]
# Mock streaming response chunks that represent a function call response
mock_response_chunk = {
"candidates": [
{
"content": {
"parts": [
{
"functionCall": {
"name": "schedule_meeting",
"args": {
"attendees": ["Bob", "Alice"],
"date": "2025-03-27",
"time": "10:00",
"topic": "Q3 planning",
},
}
}
],
"role": "model",
},
"finishReason": "STOP",
"index": 0,
}
],
"usageMetadata": {
"promptTokenCount": 15,
"candidatesTokenCount": 5,
"totalTokenCount": 20,
},
}
# Convert to bytes as expected by the streaming iterator
raw_chunks = [
f"data: {json.dumps(mock_response_chunk)}\n\n".encode(),
b"data: [DONE]\n\n",
]
# Mock the HTTP handler
with unittest.mock.patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=unittest.mock.AsyncMock,
) as mock_post:
# Create mock response object
mock_response = unittest.mock.MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
# Mock the aiter_bytes method to return our chunks as bytes
async def mock_aiter_bytes():
for chunk in raw_chunks:
yield chunk
mock_response.aiter_bytes = mock_aiter_bytes
mock_post.return_value = mock_response
print(
"\n--- Testing async agenerate_content_stream with function call parsing ---"
)
response = await litellm.google_genai.agenerate_content_stream(
model="gemini/gemini-2.5-flash-lite",
contents=contents,
tools=[
{
"functionDeclarations": [
{
"name": "schedule_meeting",
"description": "Schedules a meeting with specified attendees at a given time and date.",
"parameters": {
"type": "object",
"properties": {
"attendees": {
"type": "array",
"items": {"type": "string"},
"description": "List of people attending the meeting.",
},
"date": {
"type": "string",
"description": "Date of the meeting (e.g., '2024-07-29')",
},
"time": {
"type": "string",
"description": "Time of the meeting (e.g., '15:00')",
},
"topic": {
"type": "string",
"description": "The subject or topic of the meeting.",
},
},
"required": ["attendees", "date", "time", "topic"],
},
}
]
}
],
)
# Collect all chunks and parse function calls
chunks = []
function_calls = []
chunk_count = 0
async for chunk in response:
chunk_count += 1
print(f"Received chunk {chunk_count}: {chunk}")
chunks.append(chunk)
# Stop after a reasonable number of chunks to prevent infinite loop
if chunk_count > 10:
break
# Parse function calls from byte chunks
if isinstance(chunk, bytes):
try:
# Decode bytes to string
chunk_str = chunk.decode("utf-8")
print(f"Decoded chunk: {chunk_str}")
# Extract JSON from Server-Sent Events format (data: {...})
if chunk_str.startswith("data: ") and not chunk_str.startswith(
"data: [DONE]"
):
json_str = chunk_str[6:].strip() # Remove 'data: ' prefix
try:
parsed_json = json.loads(json_str)
print(f"Parsed JSON: {parsed_json}")
# Parse function calls from the JSON
if "candidates" in parsed_json:
for candidate in parsed_json["candidates"]:
if (
"content" in candidate
and "parts" in candidate["content"]
):
for part in candidate["content"]["parts"]:
if "functionCall" in part:
function_calls.append(
{
"name": part["functionCall"][
"name"
],
"args": part["functionCall"][
"args"
],
}
)
print(
f"Found function call: {part['functionCall']}"
)
except json.JSONDecodeError as e:
print(f"Failed to parse JSON: {e}")
except UnicodeDecodeError as e:
print(f"Failed to decode bytes: {e}")
# Handle dict responses (in case some chunks are already parsed)
elif isinstance(chunk, dict):
# Direct dict response
if "candidates" in chunk:
for candidate in chunk["candidates"]:
if "content" in candidate and "parts" in candidate["content"]:
for part in candidate["content"]["parts"]:
if "functionCall" in part:
function_calls.append(
{
"name": part["functionCall"]["name"],
"args": part["functionCall"]["args"],
}
)
# Handle object responses with attributes
elif hasattr(chunk, "candidates") and chunk.candidates:
for candidate in chunk.candidates:
if hasattr(candidate, "content") and candidate.content:
if (
hasattr(candidate.content, "parts")
and candidate.content.parts
):
for part in candidate.content.parts:
if (
hasattr(part, "function_call")
and part.function_call
):
function_calls.append(
{
"name": part.function_call.name,
"args": part.function_call.args,
}
)
# Assertions
print(f"\nFunction calls found: {function_calls}")
print(f"Total chunks received: {chunk_count}")
# Assert we found at least one function call
assert (
len(function_calls) > 0
), "Expected at least one function call in the streaming response"
# Check the first function call
function_call = function_calls[0]
# Assert function name
assert (
function_call["name"] == "schedule_meeting"
), f"Expected function name 'schedule_meeting', got '{function_call['name']}'"
# Assert function arguments
args = function_call["args"]
assert "attendees" in args, "Expected 'attendees' in function call arguments"
assert "date" in args, "Expected 'date' in function call arguments"
assert "time" in args, "Expected 'time' in function call arguments"
assert "topic" in args, "Expected 'topic' in function call arguments"
# Assert specific argument values
assert args["attendees"] == [
"Bob",
"Alice",
], f"Expected attendees ['Bob', 'Alice'], got {args['attendees']}"
assert (
args["date"] == "2025-03-27"
), f"Expected date '2025-03-27', got {args['date']}"
assert args["time"] == "10:00", f"Expected time '10:00', got {args['time']}"
assert (
args["topic"] == "Q3 planning"
), f"Expected topic 'Q3 planning', got {args['topic']}"
print("✅ All function call assertions passed!")
@pytest.mark.asyncio
async def test_validate_post_request_parameters():
"""
Test that the correct parameters are sent in the POST request to Google GenAI API
Params validated
1. model
2. contents
3. tools
"""
from litellm.types.google_genai.main import ToolConfigDict
contents = [
{
"role": "user",
"parts": [
{
"text": "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning"
}
],
}
]
tools = [
{
"functionDeclarations": [
{
"name": "schedule_meeting",
"description": "Schedules a meeting with specified attendees at a given time and date.",
"parameters": {
"type": "object",
"properties": {
"attendees": {
"type": "array",
"items": {"type": "string"},
"description": "List of people attending the meeting.",
},
"date": {
"type": "string",
"description": "Date of the meeting (e.g., '2024-07-29')",
},
"time": {
"type": "string",
"description": "Time of the meeting (e.g., '15:00')",
},
"topic": {
"type": "string",
"description": "The subject or topic of the meeting.",
},
},
"required": ["attendees", "date", "time", "topic"],
},
}
]
}
]
# Mock response for the HTTP request
raw_chunks = [b"data: [DONE]\n\n"]
# Mock the HTTP handler to capture the request
with unittest.mock.patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=unittest.mock.AsyncMock,
) as mock_post:
# Create mock response object
mock_response = unittest.mock.MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
# Mock the aiter_bytes method
async def mock_aiter_bytes():
for chunk in raw_chunks:
yield chunk
mock_response.aiter_bytes = mock_aiter_bytes
mock_post.return_value = mock_response
print("\n--- Testing POST request parameters validation ---")
# Make the API call
response = await litellm.google_genai.agenerate_content_stream(
model="gemini/gemini-2.5-flash-lite", contents=contents, tools=tools
)
# Consume the response to ensure the request is made
async for chunk in response:
pass
# Validate that the HTTP post was called
assert mock_post.called, "Expected HTTP POST to be called"
# Get the call arguments
call_args, call_kwargs = mock_post.call_args
print(f"POST call args: {call_args}")
print(f"POST call kwargs: {call_kwargs}")
# Validate URL contains the correct endpoint
if call_args:
url = call_args[0] if len(call_args) > 0 else call_kwargs.get("url")
assert url is not None, "Expected URL to be provided"
assert (
"generativelanguage.googleapis.com" in url
), f"Expected Google API URL, got: {url}"
assert (
"streamGenerateContent" in url
), f"Expected streamGenerateContent endpoint, got: {url}"
print(f"✅ URL validation passed: {url}")
# Get the request data/json from the call
request_data = None
if "data" in call_kwargs:
# If data is passed as bytes, decode it
if isinstance(call_kwargs["data"], bytes):
request_data = json.loads(call_kwargs["data"].decode("utf-8"))
else:
request_data = call_kwargs["data"]
elif "json" in call_kwargs:
request_data = call_kwargs["json"]
assert request_data is not None, "Expected request data to be provided"
print(f"Request data: {json.dumps(request_data, indent=2)}")
# Validate model field
assert "model" in request_data, "Expected 'model' field in request data"
# Model might be transformed, but should contain gemini-2.5-flash-lite
model_value = request_data["model"]
assert (
"gemini-2.5-flash-lite" in model_value
), f"Expected model to contain 'gemini-2.5-flash-lite', got: {model_value}"
print(f"✅ Model validation passed: {model_value}")
# Validate contents field
assert "contents" in request_data, "Expected 'contents' field in request data"
request_contents = request_data["contents"]
assert isinstance(request_contents, list), "Expected contents to be a list"
assert len(request_contents) > 0, "Expected at least one content item"
# Check the first content item
first_content = request_contents[0]
assert "role" in first_content, "Expected 'role' in content item"
assert (
first_content["role"] == "user"
), f"Expected role 'user', got: {first_content['role']}"
assert "parts" in first_content, "Expected 'parts' in content item"
assert isinstance(first_content["parts"], list), "Expected parts to be a list"
assert len(first_content["parts"]) > 0, "Expected at least one part"
# Check the text content
first_part = first_content["parts"][0]
assert "text" in first_part, "Expected 'text' in part"
expected_text = "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning"
assert (
first_part["text"] == expected_text
), f"Expected text '{expected_text}', got: {first_part['text']}"
print(f"✅ Contents validation passed")
# Validate tools field
assert "tools" in request_data, "Expected 'tools' field in request data"
request_tools = request_data["tools"]
assert isinstance(request_tools, list), "Expected tools to be a list"
assert len(request_tools) > 0, "Expected at least one tool"
# Check the first tool
first_tool = request_tools[0]
assert (
"functionDeclarations" in first_tool
), "Expected 'functionDeclarations' in tool"
function_declarations = first_tool["functionDeclarations"]
assert isinstance(
function_declarations, list
), "Expected functionDeclarations to be a list"
assert (
len(function_declarations) > 0
), "Expected at least one function declaration"
# Check the function declaration
func_decl = function_declarations[0]
assert "name" in func_decl, "Expected 'name' in function declaration"
assert (
func_decl["name"] == "schedule_meeting"
), f"Expected function name 'schedule_meeting', got: {func_decl['name']}"
assert (
"description" in func_decl
), "Expected 'description' in function declaration"
assert (
"parameters" in func_decl
), "Expected 'parameters' in function declaration"
# Check function parameters
params = func_decl["parameters"]
assert "type" in params, "Expected 'type' in parameters"
assert (
params["type"] == "object"
), f"Expected parameters type 'object', got: {params['type']}"
assert "properties" in params, "Expected 'properties' in parameters"
assert "required" in params, "Expected 'required' in parameters"
# Check required fields
required_fields = params["required"]
expected_required = ["attendees", "date", "time", "topic"]
assert set(required_fields) == set(
expected_required
), f"Expected required fields {expected_required}, got: {required_fields}"
print(f"✅ Tools validation passed")
print("✅ All POST request parameter validations passed!")
Reference in New Issue View Git Blame Copy Permalink