standard_rag_agen.py

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

# 1. Setup
load_dotenv()

# We explicitly override the location to match your Corpus ID (us-east1)
# If we don't do this, it might default to us-central1 and fail to find the corpus.
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
CORPUS_LOCATION = "us-east1" 
YOUR_CORPUS_ID = "projects/693246358019/locations/us-east1/ragCorpora/2305843009213693952"

# Initialize Vertex AI in the correct region for this corpus
vertexai.init(project=PROJECT_ID, location=CORPUS_LOCATION)

# 2. Configure the Official Retrieval Tool
# This is the "Method 2" production-grade tool class
rag_retrieval_tool = VertexAiRagRetrieval(
    name="company_knowledge_retrieval",
    description="Use this tool to search the company knowledge base for information.",
    rag_resources=[
        rag.RagResource(
            rag_corpus=YOUR_CORPUS_ID
        )
    ],
    similarity_top_k=5,         # Get top 5 most relevant chunks
    vector_distance_threshold=0.5 # Only return chunks that are at least 50% relevant
)

# 3. Define the Agent
# We simply give the configured tool to the agent
root_agent = Agent(
    name="rag_specialist",
    model="gemini-2.0-flash-001",
    tools=[rag_retrieval_tool],
    instruction="""
    You are a Knowledge Specialist.
    
    1. Your ONLY source of truth is the 'company_knowledge_retrieval' tool.
    2. When asked a question, ALWAYS use the tool to find the answer.
    3. If the tool returns information, synthesize it into a clear answer.
    4. If the tool returns nothing (or low relevance), state clearly: "I could not find that information in the provided documents."
    5. Do not use your own external training data to answer specific questions about the document.
    """
)