From d791df59b67ea487e4c7200dd3a627eb48f25dd8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Jun 2026 09:49:31 +0530 Subject: [PATCH] fix(vertex): propagate Vertex AI metadata in streaming success callbacks Streaming calls assembled via stream_chunk_builder were missing vertex_ai_grounding_metadata and vertex_ai_url_context_metadata in standard_logging_object.response. Merge metadata from chunks into the assembled response and mirror non-streaming hidden_params on Gemini chunks. Co-authored-by: Cursor --- .../streaming_chunk_builder_utils.py | 44 ++++++++++ .../vertex_and_google_ai_studio_gemini.py | 8 ++ litellm/main.py | 2 + .../test_litellm_logging.py | 35 ++++++++ .../test_streaming_chunk_builder_utils.py | 82 +++++++++++++++++++ 5 files changed, 171 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c3842..99e610d69ad 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -32,6 +32,14 @@ if TYPE_CHECKING: ) +VERTEX_AI_STREAM_METADATA_FIELDS = ( + "vertex_ai_grounding_metadata", + "vertex_ai_url_context_metadata", + "vertex_ai_safety_ratings", + "vertex_ai_citation_metadata", +) + + class ChunkProcessor: def __init__(self, chunks: List, messages: Optional[list] = None): self.chunks = self._sort_chunks(chunks) @@ -79,6 +87,42 @@ class ChunkProcessor: model_response._hidden_params = chunk.get("_hidden_params", {}) return model_response + @staticmethod + def _get_chunk_attr(chunk: Any, field_name: str) -> Any: + if isinstance(chunk, dict): + value = chunk.get(field_name) + if value is not None: + return value + model_extra = chunk.get("model_extra") + if isinstance(model_extra, dict): + return model_extra.get(field_name) + return None + return getattr(chunk, field_name, None) + + @staticmethod + def propagate_vertex_ai_metadata_from_chunks( + response: ModelResponse, chunks: List[Any] + ) -> None: + """ + Merge Vertex AI metadata from streaming chunks into the assembled response. + + Gemini/Vertex streaming sets these fields on individual chunks but + stream_chunk_builder must propagate them for logging callbacks. + """ + for field_name in VERTEX_AI_STREAM_METADATA_FIELDS: + merged: List[Any] = [] + for chunk in chunks: + value = ChunkProcessor._get_chunk_attr(chunk, field_name) + if not value: + continue + if isinstance(value, list): + merged.extend(value) + else: + merged.append(value) + if merged: + setattr(response, field_name, merged) + response._hidden_params[field_name] = merged + @staticmethod def _get_chunk_id(chunks: List[Dict[str, Any]]) -> str: """ 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 5cd02293f14..d6436f744c4 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 @@ -3386,9 +3386,17 @@ class ModelResponseIterator: choice.finish_reason = "tool_calls" setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore + model_response._hidden_params["vertex_ai_grounding_metadata"] = ( + grounding_metadata + ) setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore + model_response._hidden_params["vertex_ai_url_context_metadata"] = ( + url_context_metadata + ) setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore + model_response._hidden_params["vertex_ai_safety_ratings"] = safety_ratings setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore + model_response._hidden_params["vertex_ai_citation_metadata"] = citation_metadata return ( grounding_metadata, diff --git a/litellm/main.py b/litellm/main.py index 64891e2def9..c7f0b4a2091 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7761,6 +7761,7 @@ def stream_chunk_builder( # noqa: PLR0915 "cost", logging_obj._response_cost_calculator(result=response), ) + processor.propagate_vertex_ai_metadata_from_chunks(response, chunks) return response tool_call_chunks = [ @@ -7940,6 +7941,7 @@ def stream_chunk_builder( # noqa: PLR0915 usage, "cost", logging_obj._response_cost_calculator(result=response) ) + processor.propagate_vertex_ai_metadata_from_chunks(response, chunks) return response except Exception as e: verbose_logger.exception( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index d57d8dafdbd..34edd6eccf3 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2165,6 +2165,41 @@ def test_get_assembled_streaming_response_returns_result_for_streaming(): assert assembled is result +def test_streaming_success_handler_includes_vertex_ai_metadata_in_standard_logging(): + """Assembled streaming responses should include Vertex AI metadata in logging payload.""" + import datetime + + from litellm.types.utils import Choices, Message + + logging_obj = _make_logging_obj(stream=True) + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + url_context_metadata = [{"urlMetadata": [{"retrievedUrl": "https://example.com"}]}] + result = ModelResponse( + id="resp-1", + choices=[ + Choices( + index=0, + message=Message(role="assistant", content="hello"), + finish_reason="stop", + ) + ], + model="gemini-2.5-flash", + ) + setattr(result, "vertex_ai_grounding_metadata", grounding_metadata) + setattr(result, "vertex_ai_url_context_metadata", url_context_metadata) + result._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + result._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata + + start = datetime.datetime.now() + end = datetime.datetime.now() + logging_obj.success_handler(result=result, start_time=start, end_time=end) + + payload = logging_obj.model_call_details.get("standard_logging_object") + assert payload is not None + assert payload["response"]["vertex_ai_grounding_metadata"] == grounding_metadata + assert payload["response"]["vertex_ai_url_context_metadata"] == url_context_metadata + + def test_get_assembled_streaming_response_returns_none_for_non_streaming_text_completion(): """Non-streaming TextCompletionResponse should also return None.""" import datetime diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd9..22383deb569 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -613,3 +613,85 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert ( response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" ) + + +def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): + """Vertex AI metadata on streaming chunks must appear on assembled response.""" + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + url_context_metadata = [{"urlMetadata": [{"retrievedUrl": "https://example.com"}]}] + + chunk1 = ModelResponseStream( + id="chatcmpl-vertex-1", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="The weather", role="assistant"), + ) + ], + ) + setattr(chunk1, "vertex_ai_grounding_metadata", grounding_metadata) + chunk1._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + + chunk2 = ModelResponseStream( + id="chatcmpl-vertex-1", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" is sunny.", role="assistant"), + ) + ], + ) + setattr(chunk2, "vertex_ai_url_context_metadata", url_context_metadata) + chunk2._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata + + response = stream_chunk_builder(chunks=[chunk1, chunk2]) + assert response is not None + assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata + assert getattr(response, "vertex_ai_url_context_metadata") == url_context_metadata + assert response._hidden_params["vertex_ai_grounding_metadata"] == grounding_metadata + assert ( + response._hidden_params["vertex_ai_url_context_metadata"] + == url_context_metadata + ) + + dumped = response.model_dump() + assert dumped["vertex_ai_grounding_metadata"] == grounding_metadata + assert dumped["vertex_ai_url_context_metadata"] == url_context_metadata + + +def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): + """Dict snapshot chunks (model_dump) should also propagate Vertex AI metadata.""" + chunk_dict = ModelResponseStream( + id="chatcmpl-vertex-2", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ).model_dump() + chunk_dict["vertex_ai_grounding_metadata"] = [{"webSearchQueries": ["test query"]}] + chunk_dict["_hidden_params"] = { + "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}] + } + + response = stream_chunk_builder(chunks=[chunk_dict]) + assert response is not None + assert getattr(response, "vertex_ai_grounding_metadata") == [ + {"webSearchQueries": ["test query"]} + ] + assert response.model_dump()["vertex_ai_grounding_metadata"] == [ + {"webSearchQueries": ["test query"]} + ]