ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #22077 from Harshit28j/litellm_gemini_trace_id_missingv2

Browse Source 
Litellm gemini trace id missingv2
This commit is contained in:
Harshit Jain 2026-02-25 15:09:42 +05:30 committed by GitHub
commit 66c49dbb9c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 163 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/my-website/docs/generateContent.md
View File

@ -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
---

53
litellm/proxy/google_endpoints/endpoints.py
View File

@ -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(

124
tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py
View File

@ -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