tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-07 06:29:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] Return response_id == upstream response ID for VertexAI + Google AI studio (Stream+Non stream) (#11456)

Browse Source 
* fix: vertexAI return responseID

* fix: vertexAI return responseID

* test_vertex_ai_response_id

* test: test_vertex_ai_streaming_response_id

* test_vertex_ai_streaming_response_id
This commit is contained in:
Ishaan Jaff
2025-06-05 20:18:55 -07:00
committed by GitHub
parent 23627d6a26
commit f0cb80ec50
4 changed files with 147 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/litellm_core_utils/streaming_handler.py
+9 -11
View File
@@ -85,9 +85,9 @@ class CustomStreamWrapper:
self.system_fingerprint: Optional[str] = None
self.received_finish_reason: Optional[str] = None
self.intermittent_finish_reason: Optional[
str
] = None # finish reasons that show up mid-stream
self.intermittent_finish_reason: Optional[str] = (
None # finish reasons that show up mid-stream
)
self.special_tokens = [
"<|assistant|>",
"<|system|>",
@@ -619,8 +619,6 @@ class CustomStreamWrapper:
model_response = ModelResponseStream(**args)
if self.response_id is not None:
model_response.id = self.response_id
else:
self.response_id = model_response.id # type: ignore
if self.system_fingerprint is not None:
model_response.system_fingerprint = self.system_fingerprint
if hidden_params is not None:
@@ -1305,9 +1303,9 @@ class CustomStreamWrapper:
_json_delta = delta.model_dump()
print_verbose(f"_json_delta: {_json_delta}")
if "role" not in _json_delta or _json_delta["role"] is None:
_json_delta[
"role"
] = "assistant" # mistral's api returns role as None
_json_delta["role"] = (
"assistant" # mistral's api returns role as None
)
if "tool_calls" in _json_delta and isinstance(
_json_delta["tool_calls"], list
):
@@ -1697,9 +1695,9 @@ class CustomStreamWrapper:
chunk = next(self.completion_stream)
if chunk is not None and chunk != b"":
print_verbose(f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}")
processed_chunk: Optional[
ModelResponseStream
] = self.chunk_creator(chunk=chunk)
processed_chunk: Optional[ModelResponseStream] = (
self.chunk_creator(chunk=chunk)
)
print_verbose(
f"PROCESSED CHUNK POST CHUNK CREATOR: {processed_chunk}"
)
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+35 -26
View File
@@ -267,7 +267,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"""
return Tools(googleSearch={})
def _map_function(self, value: List[dict]) -> List[Tools]: # noqa: PLR0915
def _map_function(self, value: List[dict]) -> List[Tools]: # noqa: PLR0915
gtool_func_declarations = []
googleSearch: Optional[dict] = None
googleSearchRetrieval: Optional[dict] = None
@@ -304,9 +304,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return None
for tool in value:
openai_function_object: Optional[
ChatCompletionToolParamFunctionChunk
] = None
openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = (
None
)
if "function" in tool: # tools list
_openai_function_object = ChatCompletionToolParamFunctionChunk( # type: ignore
**tool["function"]
@@ -547,14 +547,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
elif param == "seed":
optional_params["seed"] = value
elif param == "reasoning_effort" and isinstance(value, str):
optional_params[
"thinkingConfig"
] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(value)
optional_params["thinkingConfig"] = (
VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(value)
)
elif param == "thinking":
optional_params[
"thinkingConfig"
] = VertexGeminiConfig._map_thinking_param(
cast(AnthropicThinkingParam, value)
optional_params["thinkingConfig"] = (
VertexGeminiConfig._map_thinking_param(
cast(AnthropicThinkingParam, value)
)
)
elif param == "modalities" and isinstance(value, list):
response_modalities = self.map_response_modalities(value)
@@ -1106,7 +1106,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
model_response.choices.append(choice)
return grounding_metadata, url_context_metadata, safety_ratings, citation_metadata
return (
grounding_metadata,
url_context_metadata,
safety_ratings,
citation_metadata,
)
def transform_response(
self,
@@ -1170,7 +1175,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
)
model_response.choices = []
model_response.id = completion_response.get("responseId", None)
url_context_metadata: List[dict] = []
try:
grounding_metadata, safety_ratings, citation_metadata = [], [], []
if _candidates:
@@ -1190,25 +1196,27 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
## ADD METADATA TO RESPONSE ##
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata)
model_response._hidden_params[
"vertex_ai_grounding_metadata"
] = grounding_metadata
model_response._hidden_params["vertex_ai_grounding_metadata"] = (
grounding_metadata
)
setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata)
model_response._hidden_params[
"vertex_ai_url_context_metadata"
] = url_context_metadata
setattr(
model_response, "vertex_ai_url_context_metadata", url_context_metadata
)
model_response._hidden_params["vertex_ai_url_context_metadata"] = (
url_context_metadata
)
setattr(model_response, "vertex_ai_safety_results", safety_ratings)
model_response._hidden_params[
"vertex_ai_safety_results"
] = safety_ratings # older approach - maintaining to prevent regressions
model_response._hidden_params["vertex_ai_safety_results"] = (
safety_ratings # older approach - maintaining to prevent regressions
)
## ADD CITATION METADATA ##
setattr(model_response, "vertex_ai_citation_metadata", citation_metadata)
model_response._hidden_params[
"vertex_ai_citation_metadata"
] = citation_metadata # older approach - maintaining to prevent regressions
model_response._hidden_params["vertex_ai_citation_metadata"] = (
citation_metadata # older approach - maintaining to prevent regressions
)
except Exception as e:
raise VertexAIError(
@@ -1841,6 +1849,7 @@ class ModelResponseIterator:
args["tool_calls"] = [tool_use]
returned_chunk = ModelResponseStream(
id=chunk.get("responseId", None),
choices=[
StreamingChoices(
index=0,
tests/local_testing/test_amazing_vertex_completion.py
+100
View File
@@ -3839,3 +3839,103 @@ def test_vertex_schema_test():
print(response)
def test_vertex_ai_response_id():
"""Test that litellm preserves the response ID from Vertex AI's API for non-streaming responses"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
load_vertex_ai_credentials()
client = HTTPHandler()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
mock_response.json.return_value = {
"responseId": "vertex_ai_response_123",
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Hello! How can I help you today?"}],
},
"finishReason": "STOP",
"safetyRatings": [
{
"category": "HARM_CATEGORY_HATE_SPEECH",
"probability": "NEGLIGIBLE",
}
],
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 8,
"totalTokenCount": 18,
},
}
with patch.object(client, "post", return_value=mock_response) as mock_post:
response = completion(
model="vertex_ai/gemini-1.5-pro",
messages=[{"role": "user", "content": "Hi!"}],
client=client,
)
# Verify the response ID is preserved
assert response.id == "vertex_ai_response_123"
assert response.choices[0].message.content == "Hello! How can I help you today?"
def test_vertex_ai_streaming_response_id():
"""Test that litellm preserves the response ID from Vertex AI's API for streaming responses"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
make_sync_call,
)
load_vertex_ai_credentials()
client = HTTPHandler()
def mock_post(url, **kwargs):
def stream_response():
chunk = {
"responseId": "vertex_ai_response_stream_123",
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Hello streaming!"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 8,
"totalTokenCount": 18,
},
}
yield json.dumps(chunk)
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.iter_lines = MagicMock(return_value=stream_response())
return mock_response
logging_obj = MagicMock()
with patch.object(client, "post", side_effect=mock_post):
iterator = make_sync_call(
client=client,
gemini_client=None,
api_base="https://mock-vertex-ai-api.com",
headers={},
data="{}",
model="gemini-pro",
messages=[],
logging_obj=logging_obj,
)
iterator = iter(iterator)
first_chunk = next(iterator)
assert first_chunk.id == "vertex_ai_response_stream_123"
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
+3 -2
View File
@@ -1,13 +1,14 @@
import asyncio
import json
from copy import deepcopy
from typing import List, cast
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
import pytest
from pydantic import BaseModel
import litellm
from litellm import ModelResponse
from litellm import ModelResponse, completion
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig,
)