Merge fb9f3ee24b into f445e466b4

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -2217,6 +2217,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                 continue
 
             image_response: list[ImageURLListItem] | None = None
+            content: str | None = None
             chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"}
             chat_completion_logprobs: ChoiceLogprobs | None = None
             tools: list[ChatCompletionToolCallChunk] | None = None
@@ -2278,13 +2279,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                 if reasoning_content is not None:
                     chat_completion_message["reasoning_content"] = reasoning_content
 
-                if candidate_grounding_metadata:
-                    annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations(
-                        grounding_metadata=candidate_grounding_metadata,
-                        content_text=content,
-                    )
-                    if annotations:
-                        chat_completion_message["annotations"] = annotations
                 (
                     functions,
                     tools,
@@ -2295,6 +2289,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                     is_function_call=is_function_call(standard_optional_params),
                 )
 
+            if candidate_grounding_metadata:
+                annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations(
+                    grounding_metadata=candidate_grounding_metadata,
+                    content_text=content,
+                )
+                if annotations:
+                    chat_completion_message["annotations"] = annotations
+
             if "logprobsResult" in candidate:
                 chat_completion_logprobs = VertexGeminiConfig._transform_logprobs(
                     logprobs_result=candidate["logprobsResult"]

tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py

@@ -16,6 +16,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
 from litellm.llms.vertex_ai.common_utils import VertexAIError
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+    ModelResponseIterator,
     VertexGeminiConfig,
 )
 from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata
@@ -6356,3 +6357,141 @@ def test_gemini_multi_candidate_messages_do_not_share_state():
     assert resp.choices[1].message.tool_calls is None
     assert getattr(resp.choices[1].message, "reasoning_content", None) is None
     assert resp.choices[1].provider_specific_fields["native_finish_reason"] == "STOP"
+
+
+_GROUNDING_METADATA: Final = {
+    "webSearchQueries": ["current price of gold"],
+    "groundingChunks": [
+        {"web": {"uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbC123", "title": "kitco"}},
+        {"web": {"uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/DeF456", "title": "reuters"}},
+    ],
+    "groundingSupports": [
+        {"segment": {"startIndex": 0, "endIndex": 42}, "groundingChunkIndices": [0]},
+        {"segment": {"startIndex": 43, "endIndex": 80}, "groundingChunkIndices": [1]},
+    ],
+}
+_GROUNDED_TEXT: Final = "As of today spot gold trades near $4,270"
+_GROUNDING_URIS: Final = (
+    "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbC123",
+    "https://vertexaisearch.cloud.google.com/grounding-api-redirect/DeF456",
+)
+
+
+def _gemini_stream_iterator() -> ModelResponseIterator:
+    logging_obj = MagicMock()
+    logging_obj.optional_params = {}
+    return ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj)
+
+
+def _delta_annotations(model_response: object) -> tuple[dict, ...]:
+    return tuple(
+        annotation
+        for choice in getattr(model_response, "choices", ())
+        for annotation in (getattr(getattr(choice, "delta", None), "annotations", None) or ())
+    )
+
+
+def _citation_urls(annotations: tuple[dict, ...]) -> tuple[str, ...]:
+    return tuple(annotation["url_citation"]["url"] for annotation in annotations)
+
+
+def test_streaming_grounding_on_the_final_chunk_produces_annotations() -> None:
+    iterator = _gemini_stream_iterator()
+    iterator.chunk_parser({"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]}}]})
+
+    final = iterator.chunk_parser(
+        {"candidates": [{"index": 0, "finishReason": "STOP", "groundingMetadata": _GROUNDING_METADATA}]}
+    )
+
+    assert _citation_urls(_delta_annotations(final)) == _GROUNDING_URIS
+
+
+def test_streaming_grounding_alongside_text_produces_annotations() -> None:
+    iterator = _gemini_stream_iterator()
+
+    chunk = iterator.chunk_parser(
+        {
+            "candidates": [
+                {
+                    "index": 0,
+                    "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]},
+                    "groundingMetadata": _GROUNDING_METADATA,
+                    "finishReason": "STOP",
+                }
+            ]
+        }
+    )
+
+    annotations = _delta_annotations(chunk)
+    assert _citation_urls(annotations) == _GROUNDING_URIS
+    assert annotations[0]["type"] == "url_citation"
+    assert annotations[0]["url_citation"]["start_index"] == 0
+
+
+def test_streaming_without_grounding_carries_no_annotations() -> None:
+    iterator = _gemini_stream_iterator()
+
+    chunk = iterator.chunk_parser(
+        {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "Paris."}]}, "finishReason": "STOP"}]}
+    )
+
+    assert _delta_annotations(chunk) == ()
+
+
+def test_streaming_grounding_without_web_chunks_carries_no_annotations() -> None:
+    iterator = _gemini_stream_iterator()
+
+    chunk = iterator.chunk_parser(
+        {
+            "candidates": [
+                {
+                    "index": 0,
+                    "finishReason": "STOP",
+                    "groundingMetadata": {
+                        "groundingChunks": [{"maps": {"placeId": "abc"}}],
+                        "groundingSupports": [
+                            {"segment": {"startIndex": 0, "endIndex": 5}, "groundingChunkIndices": [0]}
+                        ],
+                    },
+                }
+            ]
+        }
+    )
+
+    assert _delta_annotations(chunk) == ()
+
+
+def test_non_streaming_grounding_annotations_are_unchanged() -> None:
+    model_response = ModelResponse()
+    VertexGeminiConfig._process_candidates(
+        [
+            {
+                "index": 0,
+                "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]},
+                "groundingMetadata": _GROUNDING_METADATA,
+                "finishReason": "STOP",
+            }
+        ],
+        model_response,
+        {},
+    )
+
+    annotations = getattr(model_response.choices[-1].message, "annotations", None)
+    assert _citation_urls(tuple(annotations or ())) == _GROUNDING_URIS
+
+
+def test_streamed_grounding_survives_reassembly() -> None:
+    iterator = _gemini_stream_iterator()
+    chunks = [
+        iterator.chunk_parser(
+            {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]}}]}
+        ),
+        iterator.chunk_parser(
+            {"candidates": [{"index": 0, "finishReason": "STOP", "groundingMetadata": _GROUNDING_METADATA}]}
+        ),
+    ]
+
+    rebuilt = litellm.stream_chunk_builder(chunks=[chunk for chunk in chunks if chunk is not None])
+
+    assert rebuilt is not None
+    assert _citation_urls(tuple(getattr(rebuilt.choices[0].message, "annotations", None) or ())) == _GROUNDING_URIS