-
Notifications
You must be signed in to change notification settings - Fork 2.5k
Description
Describe the bug
I found this bug while making some tests with my diagnostic agent earlier today. I have a Vertex AI RAG corpus to store some schema information which my agent retrieves before making queries. The setup for the tool is as follows:
from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
from vertexai.preview import rag
schema_discovery = VertexAiRagRetrieval(
name='schema_discovery',
description=(
'Use this tool to retrieve osquery table schema documentation,'
),
rag_resources=[
rag.RagResource(
rag_corpus=os.environ.get("RAG_CORPORA_URI")
)
],
similarity_top_k=10,
vector_distance_threshold=0.6,
)This code was working until Saturday, but when trying today it doesn't work and returns 429. I was surprised and thought it was really a 429 problem (but as a Googler I have very generous quotas), but I decided to use my old implementation which uses a regular function call instead, and to my surprise it is working normally with no 429 errors:
def discover_schema(search_phrase: str) -> str:
"""Discovers osquery table names and schemas based on a descriptive search phrase.
Args:
search_phrase: A phrase describing the kind of information you're looking for.
For example: 'user login events' or 'network traffic'.
Returns:
Table names and schema information for tables related to the search phrase.
"""
rag_corpora_uri = os.environ.get('RAG_CORPORA_URI')
response = rag.retrieval_query(
rag_resources=[
rag.RagResource(
rag_corpus=rag_corpora_uri,
)
],
text=search_phrase,
)
return json.dumps(MessageToDict(response._pb))I could not find any obvious reasons why calling the vertex ai SDK directly doesn't have any problems, but using the ADK tool has.
To Reproduce
- Create a rag corpus in Vertex AI RAG Engine (can be empty)
- Create a new agent with
adk create agent - Add the following code to
agent.py
from google.adk.agents.llm_agent import Agent
from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
from vertexai.preview import rag
import os
rag_query = VertexAiRagRetrieval(
name='rag_query',
description=(
'Use this tool to query the rag,'
),
rag_resources=[
rag.RagResource(
rag_corpus=os.environ.get("RAG_CORPORA_URI")
)
],
similarity_top_k=10,
vector_distance_threshold=0.6,
)
root_agent = Agent(
model='gemini-2.5-flash',
name='root_agent',
description='A helpful assistant for user questions.',
instruction='Answer user questions to the best of your knowledge',
tools=[
rag_query,
]
)- Run
adk weband say "hello" to the agent
You should get the 429 error.
This implementation directly calling Vertex AI works perfectly:
from vertexai.preview import rag
import os
import json
from google.protobuf.json_format import MessageToDict
def rag_query(search_phrase: str) -> str:
"""Use this tool to query the rag."""
rag_corpora_uri = os.environ.get('RAG_CORPORA_URI')
response = rag.retrieval_query(
rag_resources=[
rag.RagResource(
rag_corpus=rag_corpora_uri,
)
],
text=search_phrase,
)
return json.dumps(MessageToDict(response._pb))Expected behavior
The agent should not error and greet the user
Desktop (please complete the following information):
- OS: Linux
- Python version(python -V): 3.12.3
- ADK version(pip show google-adk): 1.16.0
Model Information:
- Are you using LiteLLM: No
- Which model is being used: gemini-2.5-flash
Additional context
It was working without any errors just two days ago. Downgrading to ADK 1.15.0 didn't solve the issue. Calling Vertex AI RAG directly doesn't have the problem.