fix: propagate Gemini grounding metadata from ModelResponse

Merge #5661

### Link to Issue
- Closes: #5659

### Problem
LiteLLM exposes Gemini's grounding metadata on `ModelResponse.vertex_ai_grounding_metadata` rather than inside the message, and `_model_response_to_generate_content_response` in `lite_llm.py` never reads it. As a result `LlmResponse.grounding_metadata` is always `None` when Gemini is called through `LiteLlm`, breaking `event.grounding_metadata`, `after_model_callback`, and citation pipelines for the entire LiteLlm path. The native `Gemini()` path picks grounding up from `candidate.grounding_metadata` (`llm_response.py:190`).

### Solution
Pull `vertex_ai_grounding_metadata` off the `ModelResponse` after the message is converted, and attach it to `LlmResponse.grounding_metadata`. Three shapes are accepted:

- `types.GroundingMetadata` instance → used as-is.
- `dict` → validated via `types.GroundingMetadata.model_validate`.
- `list` (one entry per candidate) → first entry, then the same dispatch.

If validation fails the handler logs a warning and leaves grounding unset, so a single malformed payload doesn't break the rest of the response.

### Testing Plan

**Unit Tests** in `tests/unittests/models/test_litellm.py`:
- `test_model_response_to_generate_content_response_grounding_metadata_dict` — dict payload is propagated.
- `test_model_response_to_generate_content_response_grounding_metadata_list` — list payload uses the first entry.
- `test_model_response_to_generate_content_response_no_grounding_metadata` — missing attribute leaves grounding as `None` (no regression for non-Gemini LiteLlm models).

Full `test_litellm.py` suite passes locally:

```
259 passed in 9.36s
```

**Manual E2E**: Live Gemini API verification was not possible locally (no API key configured). The unit tests cover the conversion function the bug report points to.

### Checklist

- [x] I have read the CONTRIBUTING.md document.
- [x] I have performed a self-review of my own code.
- [x] I have commented my code, particularly in hard-to-understand areas.
- [x] I have added tests that prove my fix is effective.
- [x] New and existing unit tests pass locally with my changes.
- [ ] I have manually tested my changes end-to-end. (See above — requires live Gemini access.)

Co-authored-by: George Weale <gweale@google.com>
COPYBARA_INTEGRATE_REVIEW=#5661 from 1wos:fix-litellm-grounding 4657c9d
PiperOrigin-RevId: 934010043

Author: 2 people

src/google/adk/models/lite_llm.py

@@ -1741,6 +1741,38 @@ def _has_meaningful_signal(message: Message | Delta | None) -> bool:
       ) from e
 
 
+def _extract_grounding_metadata(
+    response: ModelResponse | ModelResponseStream,
+) -> types.GroundingMetadata | None:
+  """Pulls Gemini grounding metadata off a LiteLLM response or stream chunk.
+
+  LiteLLM exposes Gemini's grounding metadata on the response/chunk object
+  rather than inside the message, so the native Gemini path
+  (`candidate.grounding_metadata`) misses it. Mirroring it here lets downstream
+  consumers (event.grounding_metadata, after_model_callback, citation
+  pipelines, ...) rely on it for both model paths.
+
+  Returns the parsed metadata, or None when it is absent or malformed.
+  """
+  raw_grounding = getattr(response, "vertex_ai_grounding_metadata", None)
+  if not raw_grounding:
+    return None
+  # LiteLLM may emit a list (one entry per candidate) or a single value.
+  if isinstance(raw_grounding, list):
+    raw_grounding = raw_grounding[0] if raw_grounding else None
+  if isinstance(raw_grounding, types.GroundingMetadata):
+    return raw_grounding
+  if isinstance(raw_grounding, dict):
+    try:
+      return types.GroundingMetadata.model_validate(raw_grounding)
+    except Exception:  # pragma: no cover
+      logger.warning(
+          "LiteLlm: vertex_ai_grounding_metadata did not match the"
+          " GroundingMetadata schema and was dropped."
+      )
+  return None
+
+
 def _model_response_to_generate_content_response(
     response: ModelResponse,
 ) -> LlmResponse:
@@ -1792,6 +1824,11 @@ def _model_response_to_generate_content_response(
         cached_content_token_count=_extract_cached_prompt_tokens(usage_dict),
         thoughts_token_count=reasoning_tokens if reasoning_tokens else None,
     )
+
+  grounding_metadata = _extract_grounding_metadata(response)
+  if grounding_metadata:
+    llm_response.grounding_metadata = grounding_metadata
+
   return llm_response
 
 
@@ -2374,6 +2411,7 @@ async def generate_content_async(
       aggregated_llm_response = None
       aggregated_llm_response_with_tool_call = None
       usage_metadata = None
+      grounding_metadata = None
       fallback_index = 0
 
       def _finalize_tool_call_response(
@@ -2459,6 +2497,11 @@ def _reset_stream_buffers() -> None:
         function_calls.clear()
 
       async for part in await self.llm_client.acompletion(**completion_args):
+        # Grounding metadata can arrive on the first chunk (search queries) or
+        # the final chunk (supports); keep the latest non-empty one.
+        part_grounding = _extract_grounding_metadata(part)
+        if part_grounding:
+          grounding_metadata = part_grounding
         for chunk, finish_reason in _model_response_to_chunk(part):
           if isinstance(chunk, FunctionChunk):
             index = chunk.index or fallback_index
@@ -2560,11 +2603,17 @@ def _reset_stream_buffers() -> None:
         if usage_metadata:
           aggregated_llm_response.usage_metadata = usage_metadata
           usage_metadata = None
+        if grounding_metadata:
+          aggregated_llm_response.grounding_metadata = grounding_metadata
         yield aggregated_llm_response
 
       if aggregated_llm_response_with_tool_call:
         if usage_metadata:
           aggregated_llm_response_with_tool_call.usage_metadata = usage_metadata
+        if grounding_metadata:
+          aggregated_llm_response_with_tool_call.grounding_metadata = (
+              grounding_metadata
+          )
         yield aggregated_llm_response_with_tool_call
 
     else:

tests/unittests/models/test_litellm.py

@@ -2473,6 +2473,68 @@ def test_model_response_to_generate_content_response_reasoning_field():
   assert response.content.parts[1].text == "Result"
 
 
+def test_model_response_to_generate_content_response_grounding_metadata_dict():
+  """vertex_ai_grounding_metadata as a dict is propagated to the LlmResponse."""
+  model_response = ModelResponse(
+      model="gemini/gemini-2.5-flash",
+      choices=[{
+          "message": {"role": "assistant", "content": "Answer"},
+          "finish_reason": "stop",
+      }],
+  )
+  model_response.vertex_ai_grounding_metadata = {
+      "grounding_chunks": [
+          {"web": {"uri": "https://example.com", "title": "Example"}}
+      ],
+  }
+
+  response = _model_response_to_generate_content_response(model_response)
+
+  assert response.grounding_metadata is not None
+  assert (
+      response.grounding_metadata.grounding_chunks[0].web.uri
+      == "https://example.com"
+  )
+
+
+def test_model_response_to_generate_content_response_grounding_metadata_list():
+  """LiteLLM may emit a list (per candidate); the first entry is used."""
+  model_response = ModelResponse(
+      model="gemini/gemini-2.5-flash",
+      choices=[{
+          "message": {"role": "assistant", "content": "Answer"},
+          "finish_reason": "stop",
+      }],
+  )
+  model_response.vertex_ai_grounding_metadata = [
+      {"grounding_chunks": [{"web": {"uri": "https://a.test", "title": "A"}}]},
+      {"grounding_chunks": [{"web": {"uri": "https://b.test", "title": "B"}}]},
+  ]
+
+  response = _model_response_to_generate_content_response(model_response)
+
+  assert response.grounding_metadata is not None
+  assert (
+      response.grounding_metadata.grounding_chunks[0].web.uri
+      == "https://a.test"
+  )
+
+
+def test_model_response_to_generate_content_response_no_grounding_metadata():
+  """Without vertex_ai_grounding_metadata, grounding_metadata stays None."""
+  model_response = ModelResponse(
+      model="gemini/gemini-2.5-flash",
+      choices=[{
+          "message": {"role": "assistant", "content": "Answer"},
+          "finish_reason": "stop",
+      }],
+  )
+
+  response = _model_response_to_generate_content_response(model_response)
+
+  assert response.grounding_metadata is None
+
+
 def test_reasoning_content_takes_precedence_over_reasoning():
   """Test that 'reasoning_content' is prioritized over 'reasoning'."""
   message = {
@@ -3775,7 +3837,56 @@ async def test_completion_with_drop_params(mock_completion, mock_client):
 
 
 @pytest.mark.asyncio
-async def test_generate_content_async_stream(
+async def test_generate_content_async_stream_grounding_metadata(
+    mock_completion, lite_llm_instance
+):
+  final_chunk = ModelResponseStream(
+      model="test_model",
+      choices=[StreamingChoices(finish_reason="stop", delta=Delta())],
+  )
+  final_chunk.vertex_ai_grounding_metadata = {
+      "grounding_chunks": [
+          {"web": {"uri": "https://example.com", "title": "Example"}}
+      ],
+  }
+  mock_completion.return_value = iter([
+      ModelResponseStream(
+          model="test_model",
+          choices=[
+              StreamingChoices(
+                  finish_reason=None,
+                  delta=Delta(role="assistant", content="Grounded answer"),
+              )
+          ],
+      ),
+      final_chunk,
+  ])
+
+  llm_request = LlmRequest(
+      contents=[
+          types.Content(
+              role="user", parts=[types.Part.from_text(text="Test prompt")]
+          )
+      ],
+  )
+
+  responses = [
+      response
+      async for response in lite_llm_instance.generate_content_async(
+          llm_request, stream=True
+      )
+  ]
+
+  assert responses[-1].partial is False
+  assert responses[-1].grounding_metadata is not None
+  assert (
+      responses[-1].grounding_metadata.grounding_chunks[0].web.uri
+      == "https://example.com"
+  )
+
+
+@pytest.mark.asyncio
+async def test_generate_content_async_stream_with_usage_metadata(
     mock_completion, lite_llm_instance
 ):