Knowledge Engine tool for ADK

Supported in ADKPython v0.1.0Java v0.2.0Kotlin v0.7.0

The vertex_ai_rag_retrieval tool allows the agent to perform private data retrieval using Knowledge Engine.

When you use grounding with Knowledge Engine, you need to prepare a RAG corpus beforehand. Please refer to the Knowledge Engine page for setting it up.

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.

PythonKotlin

# 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,
    ],
)

import com.google.adk.kt.agents.Instruction
import com.google.adk.kt.agents.LlmAgent
import com.google.adk.kt.models.Gemini
import com.google.adk.kt.tools.VertexAiRagRetrieval
import com.google.adk.kt.types.VertexRagStoreRagResource

/**
 * An agent that answers from a Vertex AI RAG corpus.
 *
 * Retrieval happens inside the model through the Gemini-native `vertexRagStore`
 * kind, so the tool never runs locally.
 */
val ragAgent =
    LlmAgent(
        name = "rag_agent",
        model = Gemini(name = "gemini-flash-latest"),
        instruction =
            Instruction(
                "Answer questions using the documents in the RAG corpus. " +
                    "If the corpus does not cover the question, say so.",
            ),
        tools =
            listOf(
                VertexAiRagRetrieval(
                    name = "retrieve_docs",
                    description = "Retrieve reference material from the Vertex AI RAG corpus.",
                    // One corpus, or specific files from one corpus.
                    ragResources =
                        listOf(
                            VertexRagStoreRagResource(
                                ragCorpus =
                                    "projects/PROJECT_ID/locations/LOCATION/" +
                                        "ragCorpora/CORPUS_ID",
                            ),
                        ),
                    similarityTopK = 3,
                    vectorDistanceThreshold = 0.5,
                ),
            ),
    )