From 897b47201189de4e96e8681aefa458ccdfefec16 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Thu, 5 Feb 2026 05:25:26 +0530 Subject: [PATCH 1/2] fix: Metadata / Trace ID Missing in S3 Streaming Callbacks --- litellm/proxy/google_endpoints/endpoints.py | 53 ++++++++++++++++----- 1 file changed, 41 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 5c41b371ca..9768d93e92 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,6 +1,10 @@ +from datetime import datetime + 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 @@ -17,7 +21,8 @@ router = APIRouter( dependencies=[Depends(user_api_key_auth)], ) @router.post( - "/models/{model_name:path}:generateContent", dependencies=[Depends(user_api_key_auth)] + "/models/{model_name:path}:generateContent", + dependencies=[Depends(user_api_key_auth)], ) async def google_generate_content( request: Request, @@ -36,12 +41,12 @@ async def google_generate_content( data = await _read_request_body(request=request) 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, @@ -51,7 +56,19 @@ async def google_generate_content( 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") @@ -103,6 +120,18 @@ async def google_stream_generate_content( 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") @@ -247,11 +276,11 @@ async def create_interaction( ) data = await _read_request_body(request=request) - + # Default to gemini provider for interactions if "custom_llm_provider" not in data: data["custom_llm_provider"] = "gemini" - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -301,7 +330,7 @@ async def get_interaction( ): """ Get an interaction by ID. - + Per OpenAPI spec: GET /{api_version}/interactions/{interaction_id} """ from litellm.proxy.proxy_server import ( @@ -319,7 +348,7 @@ async def get_interaction( ) data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"} - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -369,7 +398,7 @@ async def delete_interaction( ): """ Delete an interaction by ID. - + Per OpenAPI spec: DELETE /{api_version}/interactions/{interaction_id} """ from litellm.proxy.proxy_server import ( @@ -387,7 +416,7 @@ async def delete_interaction( ) data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"} - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -437,7 +466,7 @@ async def cancel_interaction( ): """ Cancel an interaction by ID. - + Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}:cancel """ from litellm.proxy.proxy_server import ( @@ -455,7 +484,7 @@ async def cancel_interaction( ) data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"} - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( From 23e84eb7897c54fa46f6f46129dcfac94d83e18a Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Thu, 5 Feb 2026 05:25:26 +0530 Subject: [PATCH 2/2] fix: Metadata / Trace ID Missing in S3 Streaming Callbacks --- docs/my-website/docs/generateContent.md | 1 + .../test_google_api_endpoints.py | 124 +++++++++++++++++- 2 files changed, 122 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/generateContent.md b/docs/my-website/docs/generateContent.md index 4453e5ce06..bf8e1b6c03 100644 --- a/docs/my-website/docs/generateContent.md +++ b/docs/my-website/docs/generateContent.md @@ -15,6 +15,7 @@ Use LiteLLM to call Google AI's generateContent endpoints for text generation, m | Streaming | ✅ | | | Fallbacks | ✅ | between supported models | | Loadbalancing | ✅ | between supported models | +| Metadata Tracking | ✅ | passes trace ID, metadata to observability callbacks (e.g. S3, Langfuse) | ## Usage --- 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 2f2eaa905b..205d724c2b 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 @@ -2,10 +2,9 @@ """ Test to verify the Google GenAI proxy API endpoints """ -import asyncio import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest @@ -13,7 +12,6 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -import litellm def test_google_generate_content_endpoint(): @@ -401,3 +399,123 @@ def test_google_generate_content_with_image_config(): assert "contents" in called_data assert len(called_data["contents"]) == 1 assert called_data["contents"][0]["role"] == "user" + + +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)""" + 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 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( + "/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, + }, + ) + + 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 + + +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""" + 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}") + + app = FastAPI() + app.include_router(google_router) + client = TestClient(app) + + 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: + 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, + }, + ) + + 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