VertexAiSearchTool query instability

[dominique-lambert-OBS]

Hello,
As part of the Gemini Enterprise POC, we are primarily implementing three agents: orchectrator, ao_extractor, and cv_retriever.
For cv_retriever, we have chosen to use ADK VertexAiSearchTool to search for resumes.
The prompts to be implemented are satisfactory, even if they still need improvement.

However, we have encountered instability issues that regularly (10% of the time, sometimes even more) prevent VertexAiSearchTool from returning any results.

We have scripted the agent call to iterate using the same prompt, which we know it should always returns at least one document.

However, we have encountered instances of instability that regularly (10% of the time, sometimes even more) result in VertexAiSearchTool not returning any results.

We have scripted the call to the agent to iterate using the same prompt, which we know always returns at least one document.

The problem:

• It is the LLM that returns responses based on the “grounding” returned by VertexAiSearchTool
• The LLM does not explicitly indicate that there is a “connection” problem to the CV Data Store, but may indicate this:
  o [response from root_agent Type: Complete Text Message]: I am sorry, but I encountered an error trying to connect to the search tool. Please try again later.
  o [response from root_agent Type: Complete Text Message]: I am unable to directly list URIs from the document store as the search tool only provides snippets of information, not the direct URIs of the documents. To retrieve URIs, I would need a different functionality from the search tool.
• In the event of an anomaly, VertexAiSearchTool returns the grounding_metadata field but not the grounding_chunks, grounding_supports, and retrieval_queries fields as expected by the LLM:

I am attaching the associated logs and http_response. Cases 2 and 7 correspond to the anomaly, while case 10 corresponds to a case that works.

Context :
google-cloud-aiplatform[adk,agent_engines]==1.125.0
google-adk==1.18.0
python 3.12.12

Only one agent : root_agent
Only one tool : VertexAiSearchTool(data_store_id=DATASTORE_PATH)
The data datastore (location global) contains 151 unstructured documents (.pdf, .docx), Digital Parser with no advanced chunking

myapp_2.log
myapp_7.log
myapp_10.log
sdk_http_reponse_7.json
sdk_http_response_2.json
sdk_http_response_10.json

root_agent.py
main.py
agent.py

Could you suggest some avenues of research that would help us understand the problem?
Thank you in advance.

[DeanChensj]

Hi @dominique-lambert-OBS , is this an issue with Vertex AI Search itself?

[dominique-lambert-OBS]

Hi @DeanChensj

It is difficult to answer this question because we can't see any log in the tool VertexAiSearchTool
See the issue 1121

“the after_tool_callback and before_tool_callback functions do not trigger for VertexAiSearchTool (and other built-in model tools such as GoogleSearchTool and VertexAiRagRetrieval), as these tools run on the server side in the LLM model rather than on the client side.”

We do not know if the grounding phase is performed in VertexAiSearchTool.

But VertexAiSearchTool returns the grounding_metadata field but not others fields grounding_chunks, grounding_supports, and retrieval_queries

Would it be possible to see the code that uses VertexAiSearch to query the data store?

Thank you for your help.

[surajksharma07]

@dominique-lambert-OBS I found the issue - there's a bug in ADK's Live API that doesn't extract grounding_metadata from VertexAiSearchTool responses.
As a workaround, use non-streaming mode by adding run_config=RunConfig(streaming_mode=StreamingMode.NONE) to your runner calls (import from google.adk.agents.run_config).
This forces the Unary API which correctly handles grounding data.
We're looking more on fixing the Live API bug in gemini_llm_connection.py - for now this workaround should resolve your intermittent failures.

[dominique-lambert-OBS]

Thank you @surajksharma07 , I will have a look on this workaround