mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
Merge pull request #26914 from BerriAI/litellm_googleGenContentHooks
Run pre_call_hook on Google generateContent endpoints
This commit is contained in:
Michael-RZ-Berri
GitHub
commit
e4fb325a3a
3 changed files with 223 additions and 555 deletions
|
|
@ -1,10 +1,6 @@
|
|||
from datetime import datetime
|
||||
from fastapi import APIRouter, Depends, Request, Response
|
||||
from fastapi.responses import ORJSONResponse
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
from fastapi.responses import ORJSONResponse, StreamingResponse
|
||||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
|
@ -30,12 +26,17 @@ async def google_generate_content(
|
|||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
llm_router,
|
||||
proxy_config,
|
||||
proxy_logging_obj,
|
||||
select_data_generator,
|
||||
user_api_base,
|
||||
user_max_tokens,
|
||||
user_model,
|
||||
user_request_timeout,
|
||||
user_temperature,
|
||||
version,
|
||||
)
|
||||
|
||||
|
|
@ -43,48 +44,33 @@ async def google_generate_content(
|
|||
if "model" not in data:
|
||||
data["model"] = model_name
|
||||
|
||||
# Extract generationConfig and pass it as config parameter
|
||||
generation_config = data.pop("generationConfig", None)
|
||||
if generation_config:
|
||||
data["config"] = generation_config
|
||||
|
||||
# Add user authentication metadata for cost tracking
|
||||
data = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_config=proxy_config,
|
||||
general_settings=general_settings,
|
||||
version=version,
|
||||
)
|
||||
|
||||
# Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id
|
||||
data["litellm_call_id"] = request.headers.get(
|
||||
"x-litellm-call-id", str(uuid.uuid4())
|
||||
)
|
||||
logging_obj, data = litellm.utils.function_setup(
|
||||
original_function="agenerate_content",
|
||||
rules_obj=litellm.utils.Rules(),
|
||||
start_time=datetime.now(),
|
||||
**data,
|
||||
)
|
||||
data["litellm_logging_obj"] = logging_obj
|
||||
|
||||
# call router
|
||||
if llm_router is None:
|
||||
raise HTTPException(status_code=500, detail="Router not initialized")
|
||||
response = await llm_router.agenerate_content(**data)
|
||||
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=response,
|
||||
request_data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
logging_obj=logging_obj,
|
||||
version=version,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
fastapi_response.headers.update(success_headers)
|
||||
return response
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
return await processor.base_process_llm_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route_type="agenerate_content",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=llm_router,
|
||||
general_settings=general_settings,
|
||||
proxy_config=proxy_config,
|
||||
select_data_generator=select_data_generator,
|
||||
model=model_name,
|
||||
user_model=user_model,
|
||||
user_temperature=user_temperature,
|
||||
user_request_timeout=user_request_timeout,
|
||||
user_max_tokens=user_max_tokens,
|
||||
user_api_base=user_api_base,
|
||||
version=version,
|
||||
)
|
||||
except Exception as e:
|
||||
raise await processor._handle_llm_api_exception(
|
||||
e=e,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
version=version,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -101,73 +87,52 @@ async def google_stream_generate_content(
|
|||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
llm_router,
|
||||
proxy_config,
|
||||
proxy_logging_obj,
|
||||
select_data_generator,
|
||||
user_api_base,
|
||||
user_max_tokens,
|
||||
user_model,
|
||||
user_request_timeout,
|
||||
user_temperature,
|
||||
version,
|
||||
)
|
||||
|
||||
data = await _read_request_body(request=request)
|
||||
|
||||
if "model" not in data:
|
||||
data["model"] = model_name
|
||||
data["stream"] = True
|
||||
|
||||
data["stream"] = True # enforce streaming for this endpoint
|
||||
|
||||
# Extract generationConfig and pass it as config parameter
|
||||
generation_config = data.pop("generationConfig", None)
|
||||
if generation_config:
|
||||
data["config"] = generation_config
|
||||
|
||||
# Add user authentication metadata for cost tracking
|
||||
data = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_config=proxy_config,
|
||||
general_settings=general_settings,
|
||||
version=version,
|
||||
)
|
||||
|
||||
# Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id
|
||||
data["litellm_call_id"] = request.headers.get(
|
||||
"x-litellm-call-id", str(uuid.uuid4())
|
||||
)
|
||||
logging_obj, data = litellm.utils.function_setup(
|
||||
original_function="agenerate_content_stream",
|
||||
rules_obj=litellm.utils.Rules(),
|
||||
start_time=datetime.now(),
|
||||
**data,
|
||||
)
|
||||
data["litellm_logging_obj"] = logging_obj
|
||||
|
||||
# call router
|
||||
if llm_router is None:
|
||||
raise HTTPException(status_code=500, detail="Router not initialized")
|
||||
response = await llm_router.agenerate_content_stream(**data)
|
||||
|
||||
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=response,
|
||||
request_data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
logging_obj=logging_obj,
|
||||
version=version,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# Check if response is an async iterator (streaming response)
|
||||
if response is not None and hasattr(response, "__aiter__"):
|
||||
return StreamingResponse(
|
||||
content=response,
|
||||
media_type="text/event-stream",
|
||||
headers=success_headers,
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
return await processor.base_process_llm_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route_type="agenerate_content_stream",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=llm_router,
|
||||
general_settings=general_settings,
|
||||
proxy_config=proxy_config,
|
||||
select_data_generator=select_data_generator,
|
||||
model=model_name,
|
||||
user_model=user_model,
|
||||
user_temperature=user_temperature,
|
||||
user_request_timeout=user_request_timeout,
|
||||
user_max_tokens=user_max_tokens,
|
||||
user_api_base=user_api_base,
|
||||
version=version,
|
||||
)
|
||||
except Exception as e:
|
||||
raise await processor._handle_llm_api_exception(
|
||||
e=e,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
version=version,
|
||||
)
|
||||
fastapi_response.headers.update(success_headers)
|
||||
return response
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Test to verify the Google GenAI proxy API endpoints
|
|||
"""
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -13,520 +13,171 @@ sys.path.insert(
|
|||
) # Adds the parent directory to the system path
|
||||
|
||||
|
||||
def test_google_generate_content_endpoint():
|
||||
"""Test that the google_generate_content endpoint correctly routes requests"""
|
||||
# Skip this test if we can't import the required modules due to missing dependencies
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
def _build_test_client():
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def _patch_base_process(return_value=None):
|
||||
"""Patch ProxyBaseLLMRequestProcessing.base_process_llm_request so endpoint
|
||||
tests don't run the full pipeline. Returns the AsyncMock so callers can
|
||||
inspect call args."""
|
||||
if return_value is None:
|
||||
return_value = {"test": "response"}
|
||||
return patch(
|
||||
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request",
|
||||
new_callable=AsyncMock,
|
||||
return_value=return_value,
|
||||
)
|
||||
|
||||
|
||||
def test_google_generate_content_endpoint():
|
||||
"""generateContent routes through ProxyBaseLLMRequestProcessing with the
|
||||
agenerate_content route_type — that pipeline runs pre_call_hook +
|
||||
during_call_hook + post_call_success_hook for every guardrail callback."""
|
||||
try:
|
||||
client = _build_test_client()
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
|
||||
# Create a test client
|
||||
client = TestClient(app)
|
||||
|
||||
# Mock the router's agenerate_content method
|
||||
with patch("litellm.proxy.proxy_server.llm_router") as mock_router:
|
||||
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
|
||||
|
||||
# Send a request to the endpoint
|
||||
with _patch_base_process() as mock_base:
|
||||
response = client.post(
|
||||
"/v1beta/models/test-model:generateContent",
|
||||
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
|
||||
)
|
||||
|
||||
# Verify the response
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"test": "response"}
|
||||
|
||||
# Verify that agenerate_content was called
|
||||
mock_router.agenerate_content.assert_called_once()
|
||||
mock_base.assert_called_once()
|
||||
kwargs = mock_base.call_args.kwargs
|
||||
assert kwargs["route_type"] == "agenerate_content"
|
||||
assert kwargs["model"] == "test-model"
|
||||
|
||||
|
||||
def test_google_stream_generate_content_endpoint():
|
||||
"""Test that the google_stream_generate_content endpoint correctly routes streaming requests"""
|
||||
# Skip this test if we can't import the required modules due to missing dependencies
|
||||
"""streamGenerateContent must route through the same processor with the
|
||||
streaming route_type so the guardrail pipeline runs."""
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
client = _build_test_client()
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
|
||||
# Create a test client
|
||||
client = TestClient(app)
|
||||
|
||||
# Mock the router's agenerate_content_stream method to return a stream
|
||||
async def mock_stream_generator():
|
||||
yield 'data: {"test": "stream_chunk_1"}\n\n'
|
||||
yield 'data: {"test": "stream_chunk_2"}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
with patch("litellm.proxy.proxy_server.llm_router") as mock_router:
|
||||
mock_router.agenerate_content_stream = AsyncMock(
|
||||
return_value=mock_stream_generator()
|
||||
)
|
||||
|
||||
# Send a request to the endpoint
|
||||
with (
|
||||
_patch_base_process() as mock_base,
|
||||
patch(
|
||||
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__",
|
||||
return_value=None,
|
||||
) as mock_init,
|
||||
):
|
||||
response = client.post(
|
||||
"/v1beta/models/test-model:streamGenerateContent",
|
||||
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
|
||||
)
|
||||
|
||||
# Verify the response
|
||||
assert response.status_code == 200
|
||||
mock_base.assert_called_once()
|
||||
kwargs = mock_base.call_args.kwargs
|
||||
assert kwargs["route_type"] == "agenerate_content_stream"
|
||||
assert kwargs["model"] == "test-model"
|
||||
|
||||
# Verify that agenerate_content_stream was called with correct parameters
|
||||
mock_router.agenerate_content_stream.assert_called_once()
|
||||
call_args = mock_router.agenerate_content_stream.call_args
|
||||
assert call_args[1]["stream"] is True
|
||||
assert call_args[1]["model"] == "test-model"
|
||||
assert call_args[1]["contents"] == [
|
||||
# stream=True must be forced into the data the processor receives.
|
||||
init_kwargs = mock_init.call_args.kwargs
|
||||
assert init_kwargs["data"]["stream"] is True
|
||||
assert init_kwargs["data"]["model"] == "test-model"
|
||||
assert init_kwargs["data"]["contents"] == [
|
||||
{"role": "user", "parts": [{"text": "Hello"}]}
|
||||
]
|
||||
|
||||
|
||||
def test_google_generate_content_with_cost_tracking_metadata():
|
||||
"""Test that the google_generate_content endpoint includes user metadata for cost tracking"""
|
||||
def test_google_generate_content_data_flows_through_processor():
|
||||
"""The body the client sends must reach ProxyBaseLLMRequestProcessing
|
||||
intact so the pipeline can apply guardrails to it."""
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
client = _build_test_client()
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
|
||||
# Create a test client
|
||||
client = TestClient(app)
|
||||
|
||||
# Mock all required proxy server dependencies
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
|
||||
patch("litellm.proxy.proxy_server.version", "1.0.0"),
|
||||
_patch_base_process(),
|
||||
patch(
|
||||
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
|
||||
) as mock_add_data,
|
||||
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__",
|
||||
return_value=None,
|
||||
) as mock_init,
|
||||
):
|
||||
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
|
||||
|
||||
# Mock add_litellm_data_to_request to return data with metadata
|
||||
async def mock_add_litellm_data(
|
||||
data, request, user_api_key_dict, proxy_config, general_settings, version
|
||||
):
|
||||
# Simulate adding user metadata
|
||||
data["litellm_metadata"] = {
|
||||
"user_api_key_user_id": "test-user-id",
|
||||
"user_api_key_team_id": "test-team-id",
|
||||
"user_api_key": "hashed-key",
|
||||
}
|
||||
return data
|
||||
|
||||
mock_add_data.side_effect = mock_add_litellm_data
|
||||
|
||||
# Send a request to the endpoint
|
||||
response = client.post(
|
||||
client.post(
|
||||
"/v1beta/models/test-model:generateContent",
|
||||
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
|
||||
headers={"Authorization": "Bearer sk-test-key"},
|
||||
)
|
||||
|
||||
# Verify the response
|
||||
assert response.status_code == 200
|
||||
|
||||
# Verify that add_litellm_data_to_request was called
|
||||
mock_add_data.assert_called_once()
|
||||
|
||||
# Verify that agenerate_content was called with metadata
|
||||
mock_router.agenerate_content.assert_called_once()
|
||||
call_args = mock_router.agenerate_content.call_args
|
||||
called_data = call_args[1]
|
||||
|
||||
# Verify that litellm_metadata exists and contains user information
|
||||
assert "litellm_metadata" in called_data
|
||||
assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id"
|
||||
assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id"
|
||||
|
||||
|
||||
def test_google_stream_generate_content_with_cost_tracking_metadata():
|
||||
"""Test that the google_stream_generate_content endpoint includes user metadata for cost tracking"""
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
|
||||
# Create a test client
|
||||
client = TestClient(app)
|
||||
|
||||
# Mock the router's agenerate_content_stream method to return a stream
|
||||
mock_stream = AsyncMock()
|
||||
mock_stream.__aiter__ = lambda self: mock_stream
|
||||
mock_stream.__anext__.side_effect = StopAsyncIteration
|
||||
|
||||
# Mock all required proxy server dependencies
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
|
||||
patch("litellm.proxy.proxy_server.version", "1.0.0"),
|
||||
patch(
|
||||
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
|
||||
) as mock_add_data,
|
||||
):
|
||||
mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream)
|
||||
|
||||
# Mock add_litellm_data_to_request to return data with metadata
|
||||
async def mock_add_litellm_data(
|
||||
data, request, user_api_key_dict, proxy_config, general_settings, version
|
||||
):
|
||||
# Simulate adding user metadata
|
||||
data["litellm_metadata"] = {
|
||||
"user_api_key_user_id": "test-user-id",
|
||||
"user_api_key_team_id": "test-team-id",
|
||||
"user_api_key": "hashed-key",
|
||||
}
|
||||
return data
|
||||
|
||||
mock_add_data.side_effect = mock_add_litellm_data
|
||||
|
||||
# Send a request to the endpoint
|
||||
response = client.post(
|
||||
"/v1beta/models/test-model:streamGenerateContent",
|
||||
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
|
||||
headers={"Authorization": "Bearer sk-test-key"},
|
||||
)
|
||||
|
||||
# Verify the response
|
||||
assert response.status_code == 200
|
||||
|
||||
# Verify that add_litellm_data_to_request was called
|
||||
mock_add_data.assert_called_once()
|
||||
|
||||
# Verify that agenerate_content_stream was called with metadata
|
||||
mock_router.agenerate_content_stream.assert_called_once()
|
||||
call_args = mock_router.agenerate_content_stream.call_args
|
||||
called_data = call_args[1]
|
||||
|
||||
# Verify that litellm_metadata exists and contains user information
|
||||
assert "litellm_metadata" in called_data
|
||||
assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id"
|
||||
assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id"
|
||||
# Verify stream is set to True
|
||||
assert called_data["stream"] is True
|
||||
|
||||
|
||||
def test_google_generate_content_with_system_instruction():
|
||||
"""
|
||||
Test that systemInstruction is correctly passed through from the endpoint to the router.
|
||||
|
||||
This test verifies the fix for systemInstruction being dropped when forwarding
|
||||
requests to Vertex AI through the Google GenAI endpoint.
|
||||
"""
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
# Create a FastAPI app and include the router
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
|
||||
# Create a test client
|
||||
client = TestClient(app)
|
||||
|
||||
# Mock all required proxy server dependencies
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
|
||||
patch("litellm.proxy.proxy_server.version", "1.0.0"),
|
||||
patch(
|
||||
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
|
||||
) as mock_add_data,
|
||||
):
|
||||
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
|
||||
|
||||
# Mock add_litellm_data_to_request to pass through data unchanged
|
||||
async def mock_add_litellm_data(
|
||||
data, request, user_api_key_dict, proxy_config, general_settings, version
|
||||
):
|
||||
return data
|
||||
|
||||
mock_add_data.side_effect = mock_add_litellm_data
|
||||
|
||||
# Define the systemInstruction to test
|
||||
system_instruction = {"parts": [{"text": "Your name is Doodle."}]}
|
||||
|
||||
# Send a request with systemInstruction
|
||||
response = client.post(
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
json={
|
||||
"systemInstruction": system_instruction,
|
||||
"contents": [
|
||||
{"parts": [{"text": "What is your name?"}], "role": "user"}
|
||||
],
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test-key"},
|
||||
)
|
||||
|
||||
# Verify the response
|
||||
assert response.status_code == 200
|
||||
|
||||
# Verify that agenerate_content was called
|
||||
mock_router.agenerate_content.assert_called_once()
|
||||
call_args = mock_router.agenerate_content.call_args
|
||||
called_data = call_args[1]
|
||||
|
||||
# Verify that systemInstruction is present in the call arguments
|
||||
assert "systemInstruction" in called_data
|
||||
assert called_data["systemInstruction"] == system_instruction
|
||||
assert (
|
||||
called_data["systemInstruction"]["parts"][0]["text"]
|
||||
== "Your name is Doodle."
|
||||
)
|
||||
|
||||
# Verify contents are also present
|
||||
assert "contents" in called_data
|
||||
assert len(called_data["contents"]) == 1
|
||||
assert called_data["contents"][0]["role"] == "user"
|
||||
|
||||
|
||||
def test_google_generate_content_with_image_config():
|
||||
"""
|
||||
Test that imageConfig is correctly passed through from generationConfig to the router.
|
||||
|
||||
This test verifies that imageConfig parameters (aspectRatio, imageSize) are preserved
|
||||
when forwarding requests to Google GenAI through the endpoint.
|
||||
"""
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
# Create a FastAPI app and include the router
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
|
||||
# Create a test client
|
||||
client = TestClient(app)
|
||||
|
||||
# Mock all required proxy server dependencies
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
|
||||
patch("litellm.proxy.proxy_server.version", "1.0.0"),
|
||||
patch(
|
||||
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
|
||||
) as mock_add_data,
|
||||
):
|
||||
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
|
||||
|
||||
# Mock add_litellm_data_to_request to pass through data unchanged
|
||||
async def mock_add_litellm_data(
|
||||
data, request, user_api_key_dict, proxy_config, general_settings, version
|
||||
):
|
||||
return data
|
||||
|
||||
mock_add_data.side_effect = mock_add_litellm_data
|
||||
|
||||
# Send a request with generationConfig containing imageConfig
|
||||
response = client.post(
|
||||
"/v1beta/models/gemini-3-pro-image-preview:generateContent",
|
||||
json={
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{
|
||||
"text": "Create a vibrant infographic about photosynthesis"
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
|
||||
"systemInstruction": {"parts": [{"text": "Your name is Doodle."}]},
|
||||
"generationConfig": {
|
||||
"responseModalities": ["TEXT", "IMAGE"],
|
||||
"imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"},
|
||||
},
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test-key"},
|
||||
)
|
||||
|
||||
# Verify the response
|
||||
assert response.status_code == 200
|
||||
|
||||
# Verify that agenerate_content was called
|
||||
mock_router.agenerate_content.assert_called_once()
|
||||
call_args = mock_router.agenerate_content.call_args
|
||||
called_data = call_args[1]
|
||||
|
||||
# Verify that config is present in the call arguments
|
||||
assert "config" in called_data
|
||||
|
||||
# Verify that imageConfig is preserved in the config
|
||||
assert "imageConfig" in called_data["config"]
|
||||
assert called_data["config"]["imageConfig"]["aspectRatio"] == "9:16"
|
||||
assert called_data["config"]["imageConfig"]["imageSize"] == "4K"
|
||||
|
||||
# Verify that responseModalities is also preserved
|
||||
assert "responseModalities" in called_data["config"]
|
||||
assert called_data["config"]["responseModalities"] == ["TEXT", "IMAGE"]
|
||||
|
||||
# Verify contents are also present
|
||||
assert "contents" in called_data
|
||||
assert len(called_data["contents"]) == 1
|
||||
assert called_data["contents"][0]["role"] == "user"
|
||||
data = mock_init.call_args.kwargs["data"]
|
||||
assert data["model"] == "test-model"
|
||||
assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
|
||||
assert data["systemInstruction"] == {
|
||||
"parts": [{"text": "Your name is Doodle."}]
|
||||
}
|
||||
# generationConfig arrives intact here; the rename to `config` is
|
||||
# done downstream in route_request (see test_route_llm_request).
|
||||
assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"]
|
||||
assert data["generationConfig"]["imageConfig"]["aspectRatio"] == "9:16"
|
||||
|
||||
|
||||
def test_google_generate_content_metadata_and_trace_id_callbacks():
|
||||
"""Test that google_generate_content sets litellm_call_id and logging_obj for callbacks (e.g. S3, Langfuse)"""
|
||||
def test_google_generate_content_forwards_call_id_header():
|
||||
"""The endpoint must forward the x-litellm-call-id header to the processor
|
||||
so the helper can stamp it on the logging object. Trace continuity from
|
||||
client → callbacks (S3, Langfuse, etc.) depends on this header surviving
|
||||
the hop through these endpoints."""
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
client = _build_test_client()
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
# Create a FastAPI app and include the router
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
|
||||
# Create a test client
|
||||
client = TestClient(app)
|
||||
|
||||
# Mock all required proxy server dependencies
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
|
||||
patch("litellm.proxy.proxy_server.version", "1.0.0"),
|
||||
patch(
|
||||
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
|
||||
) as mock_add_data,
|
||||
):
|
||||
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
|
||||
|
||||
# Mock add_litellm_data_to_request to return data with metadata
|
||||
async def mock_add_litellm_data(
|
||||
data, request, user_api_key_dict, proxy_config, general_settings, version
|
||||
):
|
||||
# Simulate adding user metadata
|
||||
data["litellm_metadata"] = {
|
||||
"user_api_key_user_id": "test-user-id",
|
||||
}
|
||||
return data
|
||||
|
||||
mock_add_data.side_effect = mock_add_litellm_data
|
||||
|
||||
# Send a request to the endpoint with x-litellm-call-id header
|
||||
test_call_id = "test-custom-call-id"
|
||||
response = client.post(
|
||||
with _patch_base_process() as mock_base:
|
||||
client.post(
|
||||
"/v1beta/models/test-model:generateContent",
|
||||
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
|
||||
headers={
|
||||
"Authorization": "Bearer sk-test-key",
|
||||
"x-litellm-call-id": test_call_id,
|
||||
},
|
||||
headers={"x-litellm-call-id": "trace-abc-123"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
mock_router.agenerate_content.assert_called_once()
|
||||
call_args = mock_router.agenerate_content.call_args
|
||||
called_data = call_args[1]
|
||||
|
||||
# Verify that the litellm_logging_obj got assigned in the final called_data to router
|
||||
assert "litellm_logging_obj" in called_data
|
||||
assert "litellm_call_id" in called_data
|
||||
assert called_data["litellm_call_id"] == test_call_id
|
||||
forwarded_request = mock_base.call_args.kwargs["request"]
|
||||
assert forwarded_request.headers.get("x-litellm-call-id") == "trace-abc-123"
|
||||
|
||||
|
||||
def test_google_stream_generate_content_metadata_and_trace_id_callbacks():
|
||||
"""Test that google_stream_generate_content sets litellm_call_id and logging_obj for callbacks"""
|
||||
def test_google_count_tokens_unchanged():
|
||||
"""countTokens has its own path and isn't affected by the pipeline change."""
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
client = _build_test_client()
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(google_router)
|
||||
client = TestClient(app)
|
||||
fake_response = MagicMock()
|
||||
fake_response.original_response = {
|
||||
"totalTokens": 7,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
|
||||
}
|
||||
fake_response.total_tokens = 7
|
||||
|
||||
mock_stream = AsyncMock()
|
||||
mock_stream.__aiter__ = lambda self: mock_stream
|
||||
mock_stream.__anext__.side_effect = StopAsyncIteration
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
|
||||
patch("litellm.proxy.proxy_server.version", "1.0.0"),
|
||||
patch(
|
||||
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
|
||||
) as mock_add_data,
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.token_counter",
|
||||
new_callable=AsyncMock,
|
||||
return_value=fake_response,
|
||||
):
|
||||
mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream)
|
||||
|
||||
async def mock_add_litellm_data(
|
||||
data, request, user_api_key_dict, proxy_config, general_settings, version
|
||||
):
|
||||
data["litellm_metadata"] = {
|
||||
"user_api_key_user_id": "test-user-id",
|
||||
}
|
||||
return data
|
||||
|
||||
mock_add_data.side_effect = mock_add_litellm_data
|
||||
|
||||
test_call_id = "test-custom-stream-call-id"
|
||||
response = client.post(
|
||||
"/v1beta/models/test-model:streamGenerateContent",
|
||||
json={"contents": [{"role": "user", "parts": [{"text": "Hello stream"}]}]},
|
||||
headers={
|
||||
"Authorization": "Bearer sk-test-key",
|
||||
"x-litellm-call-id": test_call_id,
|
||||
},
|
||||
"/v1beta/models/test-model:countTokens",
|
||||
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
mock_router.agenerate_content_stream.assert_called_once()
|
||||
call_args = mock_router.agenerate_content_stream.call_args
|
||||
called_data = call_args[1]
|
||||
|
||||
assert "litellm_logging_obj" in called_data
|
||||
assert "litellm_call_id" in called_data
|
||||
assert called_data["litellm_call_id"] == test_call_id
|
||||
body = response.json()
|
||||
assert body["totalTokens"] == 7
|
||||
|
|
|
|||
|
|
@ -239,3 +239,55 @@ async def test_route_request_with_router_settings_override_preserves_existing():
|
|||
assert call_kwargs["num_retries"] == 10
|
||||
# Key/team timeout should be applied since not in request
|
||||
assert call_kwargs["timeout"] == 30
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route_type", ["agenerate_content", "agenerate_content_stream"]
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_route_request_maps_generation_config_for_google_routes(route_type):
|
||||
"""For Google generate_content routes, route_request must rename
|
||||
`generationConfig` (Google's wire format) to `config` (the kwarg the
|
||||
router method expects). Without this mapping the request reaches the
|
||||
LLM with the field under the wrong name and the config is dropped."""
|
||||
data = {
|
||||
"model": "gemini-2.5-flash",
|
||||
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
|
||||
"generationConfig": {
|
||||
"responseModalities": ["TEXT", "IMAGE"],
|
||||
"imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"},
|
||||
},
|
||||
}
|
||||
llm_router = MagicMock()
|
||||
getattr(llm_router, route_type).return_value = "ok"
|
||||
|
||||
await route_request(data, llm_router, None, route_type)
|
||||
|
||||
call_kwargs = getattr(llm_router, route_type).call_args[1]
|
||||
assert "generationConfig" not in call_kwargs
|
||||
assert "config" in call_kwargs
|
||||
assert call_kwargs["config"]["responseModalities"] == ["TEXT", "IMAGE"]
|
||||
assert call_kwargs["config"]["imageConfig"]["aspectRatio"] == "9:16"
|
||||
assert call_kwargs["config"]["imageConfig"]["imageSize"] == "4K"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route_type", ["agenerate_content", "agenerate_content_stream"]
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_route_request_preserves_existing_config_for_google_routes(route_type):
|
||||
"""If the caller already supplies `config`, route_request must not
|
||||
overwrite it with `generationConfig`."""
|
||||
data = {
|
||||
"model": "gemini-2.5-flash",
|
||||
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
|
||||
"config": {"existing": True},
|
||||
"generationConfig": {"shouldNotWin": True},
|
||||
}
|
||||
llm_router = MagicMock()
|
||||
getattr(llm_router, route_type).return_value = "ok"
|
||||
|
||||
await route_request(data, llm_router, None, route_type)
|
||||
|
||||
call_kwargs = getattr(llm_router, route_type).call_args[1]
|
||||
assert call_kwargs["config"] == {"existing": True}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue