diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 3a6c36624d..396ff310e3 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -1675,6 +1675,29 @@ 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, ) + + # LiteLLM exposes Gemini's grounding metadata on the ModelResponse itself + # rather than inside the message. Mirror the native Gemini path so that + # downstream consumers (event.grounding_metadata, after_model_callback, + # citation pipelines, ...) can rely on it for both model paths. + raw_grounding = getattr(response, "vertex_ai_grounding_metadata", None) + if raw_grounding: + # 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): + llm_response.grounding_metadata = raw_grounding + elif isinstance(raw_grounding, dict): + try: + llm_response.grounding_metadata = ( + 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 llm_response diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index c195076349..5ffdd966bf 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -2358,6 +2358,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 = {