fix: Extract grounding_metadata from Live API server_content #4213
Conversation
Fixes google#3542 Extract grounding_metadata from message.server_content.grounding_metadata in the Live API receive() method and include it in LlmResponse events. This allows VertexAiSearchTool grounding data to be accessible to agents.
Summary of ChangesHello @VedantMadane, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves a critical bug that caused instability in Vertex AI Search queries by failing to extract Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
adk-bot
added
the
live
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly identifies and fixes an issue where grounding_metadata was not being extracted from Live API responses. The approach of using last_grounding_metadata to persist the data across messages within a turn is sound.
However, I've identified a few issues with the implementation of the state reset logic:
- A critical bug where
grounding_metadatais prematurely cleared before being added to aturn_completeresponse. - An omission where
grounding_metadatais not added totool_callresponses. - Some code duplication in the state reset logic that can be simplified.
I've left specific comments with suggestions to address these points. Additionally, I've suggested expanding the new tests to cover these edge cases to prevent future regressions. Once these issues are addressed, this will be a solid fix.
- Fix critical bug: Remove premature reset of last_grounding_metadata before turn_complete response to prevent data loss - Simplify duplicate reset logic in interrupted handling - Add grounding_metadata propagation to tool_call responses - Add test for grounding_metadata with text content + turn_complete - Add test for grounding_metadata with tool_call responses All 27 tests pass.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request correctly identifies and fixes an issue where grounding_metadata was not being extracted from Live API responses. The changes introduce logic to track this metadata across messages within a turn. However, there are several places where this new metadata state is reset prematurely, which could lead to it being dropped before the turn is complete. I've added specific comments with suggestions to address this state management issue. Once that is fixed, this will be a solid improvement.
| if message.tool_call: | ||
| if text: | ||
| yield self.__build_full_text_response(text) | ||
| yield self.__build_full_text_response(text, last_grounding_metadata) | ||
| text = '' | ||
| parts = [ | ||
| types.Part(function_call=function_call) | ||
| for function_call in message.tool_call.function_calls | ||
| ] | ||
| yield LlmResponse(content=types.Content(role='model', parts=parts)) | ||
| yield LlmResponse( | ||
| content=types.Content(role='model', parts=parts), | ||
| grounding_metadata=last_grounding_metadata, | ||
| ) | ||
| last_grounding_metadata = None # Reset after yielding |
There was a problem hiding this comment.
The last_grounding_metadata is being reset to None prematurely. A tool_call is part of a turn, not a terminal event. Resetting the metadata here will cause any subsequent messages in the same turn (like another tool_call or the final turn_complete message) to lose the grounding information.
| if message.tool_call: | |
| if text: | |
| yield self.__build_full_text_response(text) | |
| yield self.__build_full_text_response(text, last_grounding_metadata) | |
| text = '' | |
| parts = [ | |
| types.Part(function_call=function_call) | |
| for function_call in message.tool_call.function_calls | |
| ] | |
| yield LlmResponse(content=types.Content(role='model', parts=parts)) | |
| yield LlmResponse( | |
| content=types.Content(role='model', parts=parts), | |
| grounding_metadata=last_grounding_metadata, | |
| ) | |
| last_grounding_metadata = None # Reset after yielding | |
| if message.tool_call: | |
| if text: | |
| yield self.__build_full_text_response(text, last_grounding_metadata) | |
| text = '' | |
| parts = [ | |
| types.Part(function_call=function_call) | |
| for function_call in message.tool_call.function_calls | |
| ] | |
| yield LlmResponse( | |
| content=types.Content(role='model', parts=parts), | |
| grounding_metadata=last_grounding_metadata, | |
| ) | |
2 participants
Summary
Fixes #3542 - VertexAiSearchTool query instability where grounding metadata is intermittently missing from Live API responses.
Root Cause
The Live API's
receive()method ingemini_llm_connection.pyextracted various fields from Live API messages (usage_metadata, server_content, tool_call, session_resumption_update) but never extractedgrounding_metadatafromserver_content. This prevented agents from accessing grounding data from Vertex AI Search, even when the backend provided it.Changes
Modified
src/google/adk/models/gemini_llm_connection.pylast_grounding_metadatato accumulate grounding across messagesgrounding_metadatafrommessage.server_content.grounding_metadatagrounding_metadatainLlmResponsewhen yielding:__build_full_text_response)retrieval_queriesbut missinggrounding_chunks)Updated
tests/unittests/models/test_gemini_llm_connection.pygrounding_metadata = Noneon mock server_content objectstest_receive_extracts_grounding_metadata- verifies grounding_metadata is extracted and included in content responsestest_receive_grounding_metadata_at_turn_complete- verifies grounding_metadata is carried over to turn_complete responseTest Results
All 25 tests pass, including the 2 new grounding_metadata tests.
Impact
This fix ensures that grounding data from Vertex AI Search is properly extracted and attached to
LlmResponseevents, allowing agents to accessevent.grounding_metadata.grounding_chunkswhen available. The ~10% failure rate mentioned in #3542 was due to the Vertex AI Search backend sometimes returning incomplete responses with onlyretrieval_queriesbut nogrounding_chunks. The new warning log will help identify such backend issues.