test(gemini): cover every web source for a grounded passage (#48861)

* fix(gemini): preserve every grounding web source

* test(gemini): cover all grounding sources per passage

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Zabir Raihan <zabirraihan@icloud.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/providers/test_gemini_grounding_annotations_wire.py

@@ -0,0 +1,65 @@
+import json
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.wire import Reply, Request, wire_server
+
+
+def _grounded_reply(request: Request) -> Reply:
+    body: Final = {
+        "candidates": [
+            {
+                "content": {"role": "model", "parts": [{"text": "The passage"}]},
+                "finishReason": "STOP",
+                "groundingMetadata": {
+                    "groundingChunks": [
+                        {"retrievedContext": {"uri": "gs://documents/source"}},
+                        {"web": {"uri": "https://example.com/first", "title": "First"}},
+                        {"web": {"uri": "https://example.com/second", "title": "Second"}},
+                    ],
+                    "groundingSupports": [
+                        {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": [0, 99, 2, 1]}
+                    ],
+                },
+            }
+        ],
+        "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 2, "totalTokenCount": 7},
+    }
+    assert json.loads(request.body)["contents"] == [{"role": "user", "parts": [{"text": "Find sources"}]}]
+    if ":streamGenerateContent" in request.target:
+        return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(body)}\n\n".encode(),))
+    return Reply(body=json.dumps(body).encode())
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
+def test_chat_grounding_annotations_include_every_web_source(gateway: Gateway, stream: bool) -> None:
+    with wire_server(_grounded_reply) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(model="gemini/gemini-3.8-flash", api_base=wire.url, api_key="synthetic-key")
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {"model": model, "messages": [{"role": "user", "content": "Find sources"}], "stream": stream},
+        )
+        assert response.status_code == 200, response.text
+        payloads: Final = (
+            [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")]
+            if stream
+            else [response.json()]
+        )
+        messages: Final = tuple(payload["choices"][0]["delta" if stream else "message"] for payload in payloads)
+        annotations: Final = tuple(message["annotations"] for message in messages if message.get("annotations"))
+        assert annotations == (
+            [
+                {
+                    "type": "url_citation",
+                    "url_citation": {
+                        "start_index": 4,
+                        "end_index": 11,
+                        "url": f"https://example.com/{name.lower()}",
+                        "title": name,
+                    },
+                }
+                for name in ("Second", "First")
+            ],
+        ), response.text

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

@@ -3514,6 +3514,39 @@ def test_vertex_ai_annotation_conversion():
     assert "Current time" in time_annotation["url_citation"]["title"]
 
 
+@pytest.mark.parametrize(
+    "indices, expected_indices",
+    [([2, 0], [2, 0]), ([1, 99, -1, 2, 0], [2, 0]), ([0], [0]), ([], []), ([1, 99], [])],
+    ids=["multiple-sources", "non-web-first", "single-source", "empty", "no-web"],
+)
+def test_grounding_annotations_preserve_all_web_sources(indices: list[int], expected_indices: list[int]) -> None:
+    chunks: Final = [
+        {"web": {"uri": "https://example.com/first", "title": "First"}},
+        {"retrievedContext": {"uri": "gs://documents/source"}},
+        {"web": {"uri": "https://example.com/second", "title": "Second"}},
+    ]
+    metadata: Final = {
+        "groundingChunks": chunks,
+        "groundingSupports": [
+            {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": indices}
+        ],
+    }
+    annotations: Final = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata], "The passage")
+    expected: Final = [
+        {
+            "type": "url_citation",
+            "url_citation": {
+                "start_index": 4,
+                "end_index": 11,
+                "url": chunks[index]["web"]["uri"],
+                "title": chunks[index]["web"]["title"],
+            },
+        }
+        for index in expected_indices
+    ]
+    assert json.dumps(annotations) == json.dumps(expected)
+
+
 def test_vertex_ai_annotation_empty_grounding_metadata():
     """
     Test handling of empty or missing grounding metadata.