From f0cb80ec5034e6239ae53a7f4e83cb8198b99149 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 5 Jun 2025 20:18:55 -0700 Subject: [PATCH] [Feat] Return response_id == upstream response ID for VertexAI + Google AI studio (Stream+Non stream) (#11456) * 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 --- .../litellm_core_utils/streaming_handler.py | 20 ++-- .../vertex_and_google_ai_studio_gemini.py | 61 ++++++----- .../test_amazing_vertex_completion.py | 100 ++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 5 +- 4 files changed, 147 insertions(+), 39 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ea9d867256..2b7acbd905 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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}" ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index ba86ecf612..6ae64fccdc 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -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, diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 25b7a9dd9d..0efdc07217 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -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" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 65ff99976a..96d77a1d9b 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -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, )