8. Built-in Tools - Vertex AI RAG Engine

This blog is part of the ADK Masterclass - Hands-On Series. In this tutorial, we'll explore the Vertex AI RAG Engine built-in tool, which allows us to build grounded agents that can retrieve information from our own documents.

Vertex AI RAG Engine is a fully managed retrieval-augmented generation service from Google Cloud. It simplifies the process of ingesting data, embedding it, and retrieving relevant chunks for our LLM applications.

1. What is Vertex AI RAG Engine?

Vertex AI RAG Engine handles the heavy lifting of building RAG pipelines. Instead of managing vector databases, embedding models, and retrieval systems ourselves, RAG Engine provides a unified, managed API.

• End-to-End Management: Handles ingestion, chunking, embedding, indexing, and retrieval.
• High-Quality Retrieval: Leverages Google's semantic search technology for better relevance.
• Easy Integration: Directly supported in ADK via the VertexAiRagTool.

The RAG Process

It automates the standard 6-step RAG workflow:

1. Data Ingestion: Ingests files (local, GCS, Google Drive).
2. Transformation: Chunks data for optimal processing.
3. Embedding: Converts text to vector embeddings.
4. Indexing: Stores vectors in a specialized index (Corpus).
5. Retrieval: Finds relevant chunks for a user query.
6. Generation: The LLM generates an answer using the retrieved context.

Underlying Storage: RagManagedDb

By default, Vertex AI RAG Engine uses RagManagedDb, an enterprise-ready, fully-managed Google Spanner instance. It offers two tiers to match our needs:

• Basic Tier: Cost-effective, fixed configuration (0.1 nodes). Ideal for experiments, small data, or latency-insensitive workloads.
• Scaled Tier: Production-grade with autoscaling (1 to 10 nodes). Designed for large datasets and performance-sensitive applications.

2. Tutorial

Building a RAG Agent

Let's build an agent that answers questions based on a PDF document using Vertex AI RAG Engine.

Prerequisites

• Python 3.11 or higher
• Google Cloud Project with Vertex AI API enabled
• gcloud CLI installed and authenticated

Step 1: Create the Project

Create a new directory for our agent:

adk create rag_agent
cd rag_agent

Install required dependencies:

pip install google-adk

Step 2: Prepare RAG Corpus

We need to create a RAG Corpus in Vertex AI and upload our documents. We can do this via the Google Cloud Console or using Python code.

Option A: Using Google Cloud Console

1. Go to the Vertex AI section in Google Cloud Console.
2. Navigate to RAG Engine (or Vector Search).
3. Create a new Corpus.
4. Upload your PDF or text files to the corpus.
5. Note down the Corpus ID (e.g., projects/123.../locations/us-central1/ragCorpora/456...).

Option B: Using Python SDK

We can also create a corpus and upload files programmatically using the Vertex AI SDK. This approach is ideal for automation and CI/CD pipelines.

First, install the required dependency:

pip install google-cloud-aiplatform

Create a file named setup_corpus.py with the following code:

import vertexai
from vertexai.preview import rag
import os

# Initialize Vertex AI
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "your-project-id")
LOCATION = "us-central1"
vertexai.init(project=PROJECT_ID, location=LOCATION)

def create_rag_corpus(display_name: str, description: str = None):
    """Creates a new RAG corpus with the specified embedding model."""
    embedding_model_config = rag.EmbeddingModelConfig(
        publisher_model="publishers/google/models/text-embedding-004"
    )
    
    corpus = rag.create_corpus(
        display_name=display_name,
        description=description or f"RAG corpus: {display_name}",
        embedding_model_config=embedding_model_config,
    )
    
    print(f"Created corpus: {corpus.name}")
    return corpus

def import_documents(corpus_name: str, file_paths: list):
    """Imports documents into the RAG corpus."""
    # Supported file types: .txt, .html, .md, .pdf
    # Paths can be local files or GCS URIs (gs://bucket/file)
    rag.import_files(
        corpus_name=corpus_name,
        paths=file_paths,
        chunk_size=512,
        chunk_overlap=100
    )
    print(f"Imported {len(file_paths)} file(s) to corpus")

def list_corpora():
    """Lists all RAG corpora in the project."""
    corpora = rag.list_corpora()
    for corpus in corpora:
        corpus_id = corpus.name.split('/')[-1]
        print(f"- {corpus.display_name} (ID: {corpus_id})")
    return corpora

def query_corpus(corpus_id: str, query_text: str, top_k: int = 5):
    """Queries the RAG corpus directly."""
    corpus_path = f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{corpus_id}"
    
    response = rag.retrieval_query(
        rag_resources=[rag.RagResource(rag_corpus=corpus_path)],
        text=query_text,
        rag_retrieval_config=rag.RagRetrievalConfig(top_k=top_k)
    )
    
    if hasattr(response, "contexts") and hasattr(response.contexts, "contexts"):
        for ctx in response.contexts.contexts:
            print(f"Score: {ctx.relevance_score:.3f}")
            print(f"Text: {ctx.text[:200]}...")
            print("---")
    return response

if __name__ == "__main__":
    # 1. Create a new corpus
    corpus = create_rag_corpus(
        display_name="adk-knowledge-base",
        description="Knowledge base for our ADK agent"
    )
    
    # 2. Import documents (local files or GCS URIs)
    import_documents(
        corpus_name=corpus.name,
        file_paths=[
            "documents/product-manual.pdf",
            "documents/faq.md",
            # Or use GCS: "gs://my-bucket/documents/guide.pdf"
        ]
    )
    
    # 3. Test a query
    corpus_id = corpus.name.split('/')[-1]
    print(f"\nCorpus ID for agent config: {corpus_id}")
    print(f"Full corpus name: {corpus.name}")
    
    # Optional: Test query
    # query_corpus(corpus_id, "How do I get started?")

Run the script to create our corpus:

python setup_corpus.py

The script will output the Corpus ID that we'll use in the next step. For more advanced corpus management (update, delete, list files), check out the complete corpus management tools in the ADK Vertex AI RAG Engine repository.

Step 3: Configure the Agent

Create or edit agent.py with the following code. Replace YOUR_RAG_CORPUS_ID with the ID we created in the previous step.

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

# Configure the RAG retrieval tool with our corpus
rag_retrieval = VertexAiRagRetrieval(
    name='retrieve_documents',
    description='Retrieve relevant documents from the knowledge base',
    rag_resources=[
        rag.RagResource(
            rag_corpus="projects/YOUR_PROJECT_ID/locations/us-central1/ragCorpora/YOUR_CORPUS_ID"
        )
    ],
    similarity_top_k=10,
    vector_distance_threshold=0.6,
)

root_agent = Agent(
    model="gemini-2.5-flash",
    name="rag_agent",
    instruction="You are a helpful assistant. Use the retrieval tool to answer questions based on the provided documents.",
    tools=[rag_retrieval],
)

Set up our environment variables in .env:

GOOGLE_GENAI_USE_VERTEXAI=True
GOOGLE_CLOUD_PROJECT="your-project-id"
GOOGLE_CLOUD_LOCATION="us-central1"

Step 4: Run the Agent

Authenticate with Google Cloud (if running locally):

gcloud auth application-default login

Run the agent:

adk web

Ask questions related to the documents we uploaded!

Next Steps

Explore more enterprise search capabilities:

• Built-in Tools - Vertex AI Search: Search across websites and unstructured data stores.
• Custom Function Tools: Build custom tools for specific logic.

Resources

• Vertex AI RAG Engine Documentation
• ADK Built-in Tools Docs