Vertex AI RAG Engine tool for ADK

ADK 的 Vertex AI RAG 引擎工具

Supported in ADKPython v0.1.0Java v0.2.0

The vertex_ai_rag_retrieval tool allows the agent to perform private data retrieval using Vertex AI RAG Engine.

vertex_ai_rag_retrieval 工具允许智能体使用 Vertex AI RAG 引擎执行私有数据检索。

When you use grounding with Vertex AI RAG Engine, you need to prepare a RAG corpus beforehand. Please refer to the RAG ADK agent sample or Vertex AI RAG Engine page for setting it up.

当您使用 Vertex AI RAG 引擎进行基础时,您需要预先准备 RAG 语料库。 请参阅 RAG ADK 智能体示例或 Vertex AI RAG 引擎页面进行设置。

Warning: Single tool per agent limitation

This tool can only be used by itself within an agent instance. For more information about this limitation and workarounds, see Limitations for ADK tools.

警告:每个智能体单个工具限制

此工具只能在智能体实例中单独使用。 有关此限制和变通方法的更多信息,请参阅 ADK 工具的限制。

Python

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

from google.adk.agents import Agent
from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
from vertexai.preview import rag

from dotenv import load_dotenv
from .prompts import return_instructions_root

load_dotenv()

ask_vertex_retrieval = VertexAiRagRetrieval(
    name='retrieve_rag_documentation',
    description=(
        'Use this tool to retrieve documentation and reference materials for the question from the RAG corpus,'
    ),
    rag_resources=[
        rag.RagResource(
            # please fill in your own rag corpus
            # here is a sample rag corpus for testing purpose
            # e.g. projects/123/locations/us-central1/ragCorpora/456
            rag_corpus=os.environ.get("RAG_CORPUS")
        )
    ],
    similarity_top_k=10,
    vector_distance_threshold=0.6,
)

root_agent = Agent(
    model='gemini-2.0-flash-001',
    name='ask_rag_agent',
    instruction=return_instructions_root(),
    tools=[
        ask_vertex_retrieval,
    ]
)