diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 6ada8f58783..967ac9f0ac4 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -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( diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index a35f358f365..434f7953c21 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 96870b6cc77..bfea21e705e 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -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}