GOOGLE_GENAI_USE_VERTEXAI=1
GOOGLE_CLOUD_PROJECT=your-project-id
GOOGLE_CLOUD_LOCATION=us-central1
RAG_CORPUS=projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{CORPUS_ID}

from google.adk import Agent
from google.adk.tools import VertexAiRagRetrieval

# Configure RAG retrieval tool
rag_tool = VertexAiRagRetrieval(
    name="retrieve_docs",
    description="Retrieve relevant documentation for the question",
    rag_corpus=os.environ["RAG_CORPUS"],
    similarity_top_k=10,
    vector_distance_threshold=0.6,
)

# Create agent with RAG tool
agent = Agent(
    name="rag_agent",
    model="gemini-2.0-flash-001",
    instruction=INSTRUCTION_PROMPT,
    tools=[rag_tool],
)

INSTRUCTION_PROMPT = """
You are an AI assistant with access to a specialized document corpus.

RETRIEVAL:
- Use retrieve_docs for specific knowledge questions
- Skip retrieval for casual conversation
- Ask clarifying questions when intent is unclear

SCOPE:
- Only answer questions related to the corpus
- Say "I don't have information about that" for out-of-scope queries

CITATIONS:
- Always cite sources at the end of responses
- Format: [Title](url) or [Document Section](url)
- Consolidate multiple citations from the same source
"""

from vertexai.preview import rag

# Create corpus
corpus = rag.create_corpus(display_name="my-corpus")

# Import documents (PDF, TXT, HTML)
rag.import_files(
    corpus_name=corpus.name,
    paths=["gs://bucket/doc.pdf"],  # or local files
    chunk_size=512,
    chunk_overlap=100,
)