Skip to content

fix(tools): Use function declaration for VertexAiRagRetrieval in Gemini 2.x+ #3912

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
danicat wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix(tools): Use function declaration for VertexAiRagRetrieval in Gemini 2.x+ #3912

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e9cc9da
fix(tools): Use function declaration for VertexAiRagRetrieval in Gemi…
danicat Dec 13, 2025
b7c2ef2
Merge branch 'main' into 3315
danicat Dec 15, 2025
35154f2
Merge branch 'main' into 3315
danicat Dec 17, 2025
6859f8c
Merge branch 'main' into 3315
danicat Dec 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 4 additions & 21 deletions src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
from google.genai import types
from typing_extensions import override

from ...utils.model_name_utils import is_gemini_2_or_above
from ..tool_context import ToolContext
from .base_retrieval_tool import BaseRetrievalTool

Expand Down Expand Up @@ -62,26 +61,10 @@ async def process_llm_request(
tool_context: ToolContext,
llm_request: LlmRequest,
) -> None:
# Use Gemini built-in Vertex AI RAG tool for Gemini 2 models.
if is_gemini_2_or_above(llm_request.model):
llm_request.config = (
types.GenerateContentConfig()
if not llm_request.config
else llm_request.config
)
llm_request.config.tools = (
[] if not llm_request.config.tools else llm_request.config.tools
)
llm_request.config.tools.append(
types.Tool(
retrieval=types.Retrieval(vertex_rag_store=self.vertex_rag_store)
)
)
else:
# Add the function declaration to the tools
await super().process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
# Add the function declaration to the tools
await super().process_llm_request(
tool_context=tool_context, llm_request=llm_request
)

@override
async def run_async(
Expand Down
17 changes: 5 additions & 12 deletions tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,15 +133,8 @@ def test_vertex_rag_retrieval_for_gemini_2_x():
('user', 'test1'),
]
assert len(mockModel.requests[0].config.tools) == 1
assert mockModel.requests[0].config.tools == [
types.Tool(
retrieval=types.Retrieval(
vertex_rag_store=types.VertexRagStore(
rag_corpora=[
'projects/123456789/locations/us-central1/ragCorpora/1234567890'
]
)
)
)
]
assert 'rag_retrieval' not in mockModel.requests[0].tools_dict
assert (
mockModel.requests[0].config.tools[0].function_declarations[0].name
== 'rag_retrieval'
)
assert mockModel.requests[0].tools_dict['rag_retrieval'] is not None
Loading