Grounding Metadata when Using Google Search Tool as AgentTool

Gemini API
, ,
1

I’m currently building a Vertex AI Agent with access to both the built in RAG search tool and built in Google Search tool as AgentTools, since that is the only work-around to using multiple built in tools with one agent. The code is as follows for building the AgentTools and root agent:

google_search_agent = Agent(
    name="google_search_agent",
    model=AGENT_MODEL,
    description="This agent is designed to use the Google Search tool for retrieving information from the web. You must include citations at the end of your response.",
    instruction=return_instructions_google_search(),
    tools=[google_search_tool()],
)

GOOGLE_TOOL = AgentTool(google_search_agent)

search_rag_agent = Agent(
    name="search_rag_agent",
    model=AGENT_MODEL,
    description="This agent is designed to use the RAG tool for unstructured data in RAG Engine",
    instruction=return_instructions_rag_search(),
    tools=[ask_vertex_retrieval()],
)
RAG_TOOL = AgentTool(search_rag_agent)

WEATHER_TOOL = get_weather

root_agent = Agent(
    name=AGENT_NAME,
    model=AGENT_MODEL, 
    generate_content_config=generate_content_config,  # Optional, if you want to use custom generation settings
    description=AGENT_DESCRIPTION,
    instruction=AGENT_INSTRUCTIONS,
    tools=[RAG_TOOL, GOOGLE_TOOL],
)

The tools the sub-agents wrapped in AgentTools use are defined as follows:

def ask_vertex_retrieval():
    return VertexAiRagRetrieval(
        name='corpus',
        description=(
            'Use this tool to retrieve documentation and reference materials for the question from the RAG corpus,'
        ),
        rag_resources=[
            rag.RagResource(
                rag_corpus=os.environ.get("RAG_CORPUS")
            )
        ],
        similarity_top_k=10,
        vector_distance_threshold=0.6,
    )

def google_search_tool():
    """Returns a Google Search Tool instance for abstraction."""
    return google_search

Using the tools works, but I am unable to access grounding metadata (chunks for the RAG tool, and URL sources for the google search tool) when getting responses. I’ve included an image below in the adk where the only data I get in response is the natural language result from the sub-agent wrapped as an AgentTool.

image
image1359×646 59.8 KB

Is there any way to get the grounding metadata from the built in tools programmatically, so that the root agent can see that information? This is especially important for the google search tool, where just prompting the sub-agent/AgentTool to list its citations does not seem to work, and citations are necessary for proper grounding and usage in applications.

Any information or clarification would be greatly appreciated!

3

If anyone runs into this issue, there is currently an open issue on Github, so hopefully this can be resolved sometime soon.

5

a work around you can do is that the parent agent and the agent wrapped in the tool shares artifact service, you can use artifact service as a media to pass this information :

  • you have an after_model_callback for the agent that is wrapped as tool that save the grounding_metadata from the llm_response to artifact service.

  • you have an after_agent_callback / after_tool_callback for the parent agent to read the grounding_metadata from the artifact service