rag-patterns

What this skill does

# RAG Patterns: Vertex AI RAG Engine Integration

## Core Principles

Retrieval-Augmented Generation (RAG) grounds AI agent responses in authoritative data sources, dramatically reducing hallucinations and enabling agents to answer questions about private, domain-specific information not in their training data.

**Vertex AI RAG Engine** provides a managed service for the complete RAG pipeline: ingestion, embedding, indexing, retrieval, and citation generation.

## RAG Process Overview

### The 5-Stage RAG Pipeline

```
1. INGEST
   └─> Upload documents (PDF, TXT, HTML, etc.) to GCS

2. TRANSFORM & CHUNK
   └─> Split documents into semantic chunks (typically 500-1000 tokens)

3. EMBED
   └─> Generate vector embeddings for each chunk using embedding model

4. INDEX
   └─> Store embeddings in vector database (Vertex AI RAG Corpus)

5. RETRIEVE & GENERATE
   └─> At query time:
       a. Embed user query
       b. Find semantically similar chunks (vector search)
       c. Pass chunks to LLM as context
       d. Generate grounded response with citations
```

## Corpus Management (Required Pattern)

### Creating a RAG Corpus

```python
"""
Create and manage Vertex AI RAG corpus for private knowledge base.
"""
import os
from google.cloud import aiplatform
from google.cloud.aiplatform import rag

# Initialize Vertex AI
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
LOCATION = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")

aiplatform.init(project=PROJECT_ID, location=LOCATION)

async def create_rag_corpus(
    display_name: str,
    description: str
) -> rag.RagCorpus:
    """
    Create a new RAG corpus for knowledge base.

    Args:
        display_name: Human-readable name for the corpus
        description: Purpose and content description

    Returns:
        Created RagCorpus instance
    """
    # Create corpus
    corpus = rag.RagCorpus.create(
        display_name=display_name,
        description=description,
        # Optional: Configure embedding model
        embedding_model_config=rag.EmbeddingModelConfig(
            publisher_model="publishers/google/models/text-embedding-004"
        )
    )

    print(f"Created corpus: {corpus.name}")
    print(f"Resource name: {corpus.resource_name}")

    return corpus
```

### Listing Existing Corpora

```python
def list_rag_corpora() -> list[rag.RagCorpus]:
    """
    List all RAG corpora in the project.

    Returns:
        List of RagCorpus instances
    """
    corpora = rag.RagCorpus.list()

    for corpus in corpora:
        print(f"Name: {corpus.display_name}")
        print(f"ID: {corpus.name}")
        print(f"Description: {corpus.description}")
        print("---")

    return list(corpora)
```

### Deleting a Corpus

```python
def delete_rag_corpus(corpus_name: str) -> None:
    """
    Delete a RAG corpus.

    Args:
        corpus_name: Full resource name of corpus
                     Format: projects/{project}/locations/{location}/ragCorpora/{corpus_id}
    """
    corpus = rag.RagCorpus(corpus_name)
    corpus.delete()
    print(f"Deleted corpus: {corpus_name}")
```

## Document Ingestion (Required Pattern)

### Upload Documents from Google Cloud Storage

```python
from google.cloud.aiplatform import rag
from google.cloud.aiplatform.rag.utils import resources

async def import_files_to_corpus(
    corpus_name: str,
    gcs_uris: list[str],
    chunk_size: int = 1024,
    chunk_overlap: int = 200
) -> None:
    """
    Import documents from GCS into RAG corpus.

    Args:
        corpus_name: Full resource name of target corpus
        gcs_uris: List of GCS URIs (gs://bucket/path/file.pdf)
        chunk_size: Size of text chunks in tokens (default 1024)
        chunk_overlap: Overlap between chunks in tokens (default 200)

    Example:
        await import_files_to_corpus(
            corpus_name="projects/my-proj/locations/us-central1/ragCorpora/123",
            gcs_uris=[
                "gs://my-bucket/docs/manual.pdf",
                "gs://my-bucket/docs/guide.pdf"
            ]
        )
    """
    corpus = rag.RagCorpus(corpus_name)

    # Import files with chunking configuration
    response = rag.import_files(
        corpus_name=corpus.resource_name,
        paths=gcs_uris,
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        max_embedding_requests_per_min=1000  # Rate limiting
    )

    print(f"Import started. Imported {len(gcs_uris)} files.")
    print(f"Response: {response}")
```

### Upload Local Files (via GCS)

```python
from google.cloud import storage
import asyncio

async def upload_local_files_to_rag(
    corpus_name: str,
    local_file_paths: list[str],
    gcs_bucket: str,
    gcs_prefix: str = "rag-docs/"
) -> None:
    """
    Upload local files to GCS, then import to RAG corpus.

    Args:
        corpus_name: RAG corpus resource name
        local_file_paths: Paths to local files
        gcs_bucket: GCS bucket name (without gs://)
        gcs_prefix: Prefix path in bucket for uploaded files

    Example:
        await upload_local_files_to_rag(
            corpus_name="projects/.../ragCorpora/123",
            local_file_paths=["./docs/manual.pdf", "./docs/faq.txt"],
            gcs_bucket="my-rag-bucket",
            gcs_prefix="company-docs/"
        )
    """
    # Initialize GCS client
    storage_client = storage.Client(project=PROJECT_ID)
    bucket = storage_client.bucket(gcs_bucket)

    gcs_uris = []

    # Upload each file to GCS
    for local_path in local_file_paths:
        filename = os.path.basename(local_path)
        gcs_path = f"{gcs_prefix}{filename}"

        blob = bucket.blob(gcs_path)
        blob.upload_from_filename(local_path)

        gcs_uri = f"gs://{gcs_bucket}/{gcs_path}"
        gcs_uris.append(gcs_uri)
        print(f"Uploaded {filename} to {gcs_uri}")

    # Import into RAG corpus
    await import_files_to_corpus(corpus_name, gcs_uris)
```

### Listing Files in Corpus

```python
def list_corpus_files(corpus_name: str) -> None:
    """
    List all files in a RAG corpus.

    Args:
        corpus_name: Full resource name of corpus
    """
    corpus = rag.RagCorpus(corpus_name)
    files = rag.list_files(corpus_name=corpus.resource_name)

    for file in files:
        print(f"File: {file.display_name}")
        print(f"  URI: {file.rag_file_source.gcs_source.uris}")
        print(f"  Status: {file.state}")
        print("---")
```

## ADK Agent Integration (Production Pattern)

### Creating RAG Tool for Agent

```python
from google import genai
from google.genai import types
from google.cloud.aiplatform import rag

async def create_rag_agent(
    corpus_name: str,
    system_instruction: str
) -> types.LlmAgent:
    """
    Create ADK agent with RAG retrieval capabilities.

    Args:
        corpus_name: Full resource name of RAG corpus
        system_instruction: Agent's system prompt

    Returns:
        LlmAgent with RAG tool

    Example:
        agent = await create_rag_agent(
            corpus_name="projects/.../ragCorpora/123",
            system_instruction="You are a helpful assistant that answers questions about our company policies."
        )
    """
    client = genai.Client(vertexai=True)

    # Create RAG retrieval tool
    rag_tool = types.Tool.from_retrieval(
        retrieval=types.VertexRagStore(
            rag_resources=[
                types.RagResource(
                    rag_corpus=corpus_name
                )
            ],
            similarity_top_k=10,  # Number of chunks to retrieve
            vector_distance_threshold=0.3  # Relevance threshold (0-1)
        )
    )

    # Create agent with RAG tool
    agent = types.LlmAgent(
        model="gemini-2.0-flash-exp",
        system_instruction=system_instruction,
        tools=[rag_tool]
    )

    return agent
```

### Configuring Retrieval Parameters

```python
def create_advanced_rag_tool(
    corpus_names: list[str],
    top_k: int = 10,
    distance_threshold: float = 0.3
) -> types.Tool:
    """
    Create R