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.
|
Hi @VedantMadane , Thank you for your contribution! We appreciate you taking the time to submit this pull request. Can you fix the formatting errors before we can proceed with the review. You can use autoformat.sh. |
ryanaiagent
added
the
request clarification
Run autoformat.sh to fix formatting issues as requested in PR review.
tool_call is part of an ongoing turn, not a terminal event. Removing the premature reset of last_grounding_metadata ensures subsequent messages in the same turn (like another tool_call or turn_complete) retain the grounding information.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request correctly implements the extraction of grounding_metadata from Live API responses, fixing a key instability issue. The changes are logical and well-supported by new unit tests. My review identified two potential high-severity issues. First, the interrupted flag from the server may be lost when there is pending text to be flushed, which could lead to incorrect state handling by the client. Second, the grounding_metadata is not reset after a tool_call, which could cause stale metadata to leak into subsequent, unrelated model turns. I've provided detailed comments and suggestions for fixing these issues to improve the robustness of the implementation.
- Add interrupted parameter to __build_full_text_response to preserve interrupted signal when flushing pending text - Pass interrupted flag in turn_complete and interrupted blocks - Remove premature reset of last_grounding_metadata after interrupted (not a terminal event) - Add documentation for tool_call metadata persistence design decision Addresses review comments: - HIGH: Lost interrupted signal in full text response - HIGH: Premature reset after interrupted - MEDIUM: Duplicate reset logic (simplified by removing premature reset)
Add test_receive_interrupted_with_pending_text_preserves_flag to verify: - interrupted flag is preserved when flushing pending text - grounding_metadata is carried through to the flushed response - accumulated text is properly merged before interruption Addresses MEDIUM priority review comment about missing test coverage for edge cases.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively fixes an issue where grounding_metadata was not being extracted from Live API responses. The changes correctly track and propagate this metadata across various message types, and the new unit tests provide good coverage for these scenarios. I have one suggestion to address a potential logic bug where the grounding metadata could be reset prematurely within a turn, ensuring the fix is robust.
| text, last_grounding_metadata | ||
| ) | ||
| text = '' | ||
| last_grounding_metadata = None |
There was a problem hiding this comment.
This line prematurely resets last_grounding_metadata. The grounding metadata should persist for the entire duration of a turn and only be reset when the turn is complete, as is correctly handled in the turn_complete block.
Other parts of the code that flush buffered text (e.g., when handling tool_call or interrupted signals) correctly preserve the metadata for the remainder of the turn. Resetting it here is inconsistent with that logic and could cause grounding information to be lost before the turn ends.
To ensure consistent behavior, this line should be removed.
3 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.