RAG Agent with Knowledge Base

A complete example demonstrating how to build a RAG (Retrieval-Augmented Generation) agent using vector stores and semantic search. This example shows how to create a knowledge-based agent that can answer questions from your own documents.

Check Full Code Here

---

Overview

This example demonstrates:

• ✅ Vector store creation and configuration
• ✅ Document processing and embedding generation
• ✅ KnowledgeBaseSearchTool integration
• ✅ Multiple search strategies (hybrid, vector, keyword)
• ✅ Metadata filtering for targeted results
• ✅ Production-ready error handling

---

What You'll Build

A documentation assistant that can:

• Search through your documents using semantic understanding
• Filter results by category, type, or custom metadata
• Use multiple search strategies for different query types
• Provide accurate, context-aware answers

---

Prerequisites

• Python 3.12 or Python 3.13
• CAIP Agent ID (created via CAIP Portal)
• CAIP API Key
• CAIP Space ID
• Documents to index (text files, markdown, etc.)

---

Quick Start

1. Environment Setup

Create a .env file:

CAIP_API_KEY is a unified key for both Agents API and LLM API access. See CAIP API Key Authentication to obtain your key.

CAIP_API_KEY=your_api_key_here
CAIP_SPACE_ID=your_space_id_here
CAIP_AGENT_ID=your_agent_id_here
CAIP_REGION=ROW  # Optional: ROW (default) or CN

2. Install Dependencies

pip install caip-agents-sdk

3. Complete Example

"""RAG Agent with Knowledge Base - Complete Example."""

import asyncio
from caip_agents_sdk import CAIPAgentsClient
from caip_agents_sdk.tools import KnowledgeBaseSearchTool


async def main():
    """Build a complete RAG agent from scratch."""
    
    # ============================================================
    # Step 1: Initialize Client
    # ============================================================
    print("Initializing CAIP client...")
    client = CAIPAgentsClient()
    
    # ============================================================
    # Step 2: Create Vector Store
    # ============================================================
    print("\nCreating vector store...")
    vector_store = await client.create_vector_store({
        "name": "ProductDocs",
        "description": "Product documentation knowledge base",
        "embeddingModel": "text-embedding-3-small",
        "embeddingDimension": 1536,
        "chunkSize": 1000,
        "chunkOverlap": 200
    })

    vector_store_id = vector_store.vectorStoreId
    print(f"✅ Vector store created: {vector_store_id}")
    
    # ============================================================
    # Step 3: Add Documents to Vector Store
    # ============================================================
    print("\nAdding documents to vector store...")
    
    documents = [
        {
            "text": "To authenticate API requests, include your API key in the Authorization header as a Bearer token. Example: Authorization: Bearer YOUR_API_KEY",
            "metadata": {
                "category": "api",
                "topic": "authentication",
                "page": 1,
                "version": "v2"
            }
        },
        {
            "text": "The API has rate limits of 1000 requests per hour for standard tier and 10000 requests per hour for premium tier. Rate limit headers are included in all responses.",
            "metadata": {
                "category": "api",
                "topic": "rate-limits",
                "page": 2,
                "version": "v2"
            }
        },
        {
            "text": "Getting started with the SDK: Install using pip install caip-agents-sdk. Then import CAIPAgentsClient and create your first agent with client.create_agent().",
            "metadata": {
                "category": "tutorials",
                "topic": "getting-started",
                "page": 1,
                "level": "beginner"
            }
        },
        {
            "text": "For production deployments, we recommend using environment variables for configuration. Set CAIP_API_KEY, CAIP_SPACE_ID, and CAIP_AGENT_ID in your .env file. Optionally set CAIP_REGION to ROW or CN (ROW is the default when unset).",
            "metadata": {
                "category": "deployment",
                "topic": "configuration",
                "page": 5,
                "level": "intermediate"
            }
        },
        {
            "text": "Error handling best practices: Always wrap agent operations in try-except blocks. The SDK provides specific exceptions like ValidationError, APIError, and NotFoundError.",
            "metadata": {
                "category": "best-practices",
                "topic": "error-handling",
                "page": 10,
                "level": "advanced"
            }
        }
    ]
    
    # Generate embeddings and store documents
    texts = [doc["text"] for doc in documents]
    metadata_list = [doc["metadata"] for doc in documents]
    
    # Generate embeddings for all documents
    embeddings_response = await client.batch_text_to_embeddings(
        texts=texts,
        vector_store_id=vector_store_id,
        model="text-embedding-3-small",
        dimensions=1536
    )
    
    # Store embeddings in vector store
    await client.create_embeddings(
        vector_store_id=vector_store_id,
        embedding_data={
            "embeddings": embeddings_response,
            "chunks": texts,
            "documentIds": [f"doc_{i}" for i in range(len(texts))],
            "sequences": [0] * len(texts),
            "metadata": metadata_list
        }
    )
    
    for i in range(len(documents)):
        print(f"  ✅ Added document {i + 1}/{len(documents)}")
    
    print(f"✅ All documents indexed successfully")
    
    # ============================================================
    # Step 4: Create Agent with RAG Capabilities
    # ============================================================
    print("\nCreating RAG agent...")
    agent = client.create_agent("pydantic_ai", "your_agent_id")
    await agent.initialize()
    
    # Create conversation thread
    thread = await client.create_thread(
        agent_id="your_agent_id",
        thread_data={"title": "RAG Demo Session"}
    )
    agent.thread_id = thread.threadId
    print(f"✅ Agent ready with thread: {thread.threadId}")
    
    # ============================================================
    # Step 5: Register Knowledge Base Search Tools
    # ============================================================
    print("\nRegistering search tools...")
    
    kb_tool = KnowledgeBaseSearchTool(
        client=client,
        vector_store_id=vector_store_id
    )
    
    # General search tool (hybrid strategy)
    @agent.tool_plain
    async def search_documentation(query: str) -> str:
        """Search all documentation using hybrid search (semantic + keyword)."""
        return await kb_tool.search(
            query=query,
            search_type="hybrid",
            top_k=5
        )
    
    # API-specific search
    @agent.tool_plain
    async def search_api_docs(query: str) -> str:
        """Search API documentation only."""
        return await kb_tool.search(
            query=query,
            search_type="vector",
            metadata_filter={"category": "api"},
            score_threshold=0.7
        )
    
    # Tutorial search
    @agent.tool_plain
    async def search_tutorials(query: str) -> str:
        """Search beginner-friendly tutorials."""
        return await kb_tool.search(
            query=query,
            search_type="vector",
            metadata_filter={
                "category": "tutorials",
                "level": "beginner"
            }
        )
    
    # Advanced topic search
    @agent.tool_plain
    async def search_advanced_topics(query: str) -> str:
        """Search advanced topics and best practices."""
        return await kb_tool.search(
            query=query,
            search_type="hybrid",
            metadata_filter={"level": "advanced"}
        )
    
    print("✅ Search tools registered")
    
    # ============================================================
    # Step 6: Test the RAG Agent
    # ============================================================
    print("\n" + "=" * 70)
    print("Testing RAG Agent")
    print("=" * 70)
    
    test_queries = [
        "How do I authenticate API requests?",
        "What are the rate limits?",
        "How do I get started with the SDK?",
        "What are the best practices for error handling?",
        "How do I configure for production?"
    ]
    
    for i, query in enumerate(test_queries, 1):
        print(f"\n{'─' * 70}")
        print(f"Query {i}: {query}")
        print('─' * 70)
        
        try:
            response = await agent.run(query)
            print(f"\n📝 Agent Response:")
            print(response)
        except Exception as e:
            print(f"❌ Error: {str(e)}")
    
    print("\n" + "=" * 70)
    print("Demo Complete!")
    print("=" * 70)
    
    # ============================================================
    # Step 7: Cleanup Resources
    # ============================================================
    print("\nCleaning up resources...")
    
    try:
        # Delete all embeddings from the vector store
        print(f"Deleting embeddings from vector store: {vector_store_id}")
        await client.delete_embeddings(
        vector_store_id=vector_store.vectorStoreId, delete_request={"deleteAll": True}
        )
        await client.delete_vector_store(vector_store.vectorStoreId)
        print("✓ Deleted all embeddings and vector store\n")
        
        print("\n✅ Cleanup complete!")
        
    except Exception as e:
        print(f"❌ Cleanup error: {str(e)}")


if __name__ == "__main__":
    asyncio.run(main())

---

Key Features Explained

1. Vector Store Configuration

vector_store = await client.create_vector_store({
    "name": "ProductDocs",
    "embeddingModel": "text-embedding-3-small",  # Fast, cost-effective
    "embeddingDimension": 1536,  # Matches model dimension
    "chunkSize": 1000,  # Good balance for most docs
    "chunkOverlap": 200  # 20% overlap maintains context
})

Why these settings?

• text-embedding-3-small: Best balance of speed, quality, and cost
• 1536 dimension: Standard for this model
• 1000 chunk size: Captures full context without being too large
• 200 overlap: Ensures important information at boundaries isn't lost

2. Document Metadata

metadata = {
    "category": "api",      # High-level category
    "topic": "authentication",  # Specific topic
    "page": 1,              # Numeric reference
    "version": "v2"         # Version tracking
}

Best Practices:

• Use consistent naming conventions
• Include searchable attributes
• Add version information
• Use proper types (integers for numbers, not strings)

3. Multiple Search Tools

The example registers different tools for different search needs:

# General purpose
@agent.tool_plain
async def search_documentation(query: str) -> str:
    return await kb_tool.search(query=query, search_type="hybrid")

# Category-filtered
@agent.tool_plain
async def search_api_docs(query: str) -> str:
    return await kb_tool.search(
        query=query,
        search_type="vector",
        metadata_filter={"category": "api"}
    )

Strategy:

• Provide specific tools for common categories
• Use metadata filters to narrow results
• Choose appropriate search type for each use case

---

Search Strategies

Hybrid Search (Recommended for Most Cases)

await kb_tool.search(
    query="authentication",
    search_type="hybrid",
    vector_weight=0.7,  # 70% semantic, 30% keyword
    top_k=5
)

When to use: General queries, balanced relevance

Vector Search (Semantic Understanding)

await kb_tool.search(
    query="how to secure my API",
    search_type="vector",
    score_threshold=0.7,
    top_k=5
)

When to use: Conceptual questions, related topics

Keyword Filter (Strict Matching)

await kb_tool.search(
    query="authentication authorization token",
    search_type="keyword_filter",
    metadata_filter={"category": "api"},
    filter_logic="OR"
)

When to use: Exact term matching, filtered searches

---

Advanced Features

1. Dynamic Metadata Filtering

@agent.tool_plain
async def search_by_level(query: str, level: str) -> str:
    """Search documentation by skill level."""
    return await kb_tool.search(
        query=query,
        search_type="hybrid",
        metadata_filter={"level": level}
    )

2. Score Threshold Tuning

# High precision (fewer, more relevant results)
await kb_tool.search(
    query=query,
    score_threshold=0.8,
    top_k=3
)

# High recall (more results, lower threshold)
await kb_tool.search(
    query=query,
    score_threshold=0.5,
    top_k=10
)

3. Error Handling

@agent.tool_plain
async def safe_search(query: str) -> str:
    """Search with comprehensive error handling."""
    try:
        return await kb_tool.search(query=query, search_type="hybrid")
    except ValidationError as e:
        return f"Invalid query: {str(e)}"
    except APIError as e:
        return f"Search failed: {str(e)}"
    except Exception as e:
        return f"Unexpected error: {str(e)}"

---

Production Considerations

1. Batch Document Processing

For large document sets:

async def process_documents_batch(
    documents: list,
    vector_store_id: str,
    batch_size: int = 10
):
    """
    Process documents in batches to avoid API limits.
    
    Args:
        documents: List of dicts with 'text' and 'metadata' keys
        vector_store_id: ID of the vector store
        batch_size: Number of documents to process at once (default: 10)
    """
    for i in range(0, len(documents), batch_size):
        batch = documents[i:i + batch_size]
        texts = [doc["text"] for doc in batch]
        metadata_list = [doc["metadata"] for doc in batch]
        doc_ids = [f"doc_{i + j}" for j in range(len(batch))]
        
        # Generate embeddings using batch API
        # Supported models: text-embedding-3-small, text-embedding-3-large, titan-text-embeddings-v2
        # Supported dimensions: 1024, 1536, 3072 (must match vector store configuration)
        embeddings_response = await client.batch_text_to_embeddings(
            texts=texts,
            vector_store_id=vector_store_id,
            model="text-embedding-3-small",  # Must match vector store's embedding model
            dimensions=1536  # Must match vector store's dimension
        )
        
        # Store embeddings in vector store
        await client.create_embeddings(
            vector_store_id=vector_store_id,
            embedding_data={
                "embeddings": embeddings_response,
                "chunks": texts,
                "documentIds": doc_ids,
                "sequences": [0] * len(batch),
                "metadata": metadata_list
            }
        )
        
        print(f"✅ Processed batch {i // batch_size + 1} ({len(batch)} documents)")

2. Document Updates

Handle document updates properly:

async def update_document(
    vector_store_id: str,
    doc_id: str,
    new_text: str,
    metadata: dict
):
    """
    Update a document in the vector store by deleting old and adding new embedding.
    
    Args:
        vector_store_id: ID of the vector store
        doc_id: Document ID to update
        new_text: New text content
        metadata: Updated metadata
    """
    # Step 1: Delete old embeddings for this document
    # Note: delete_embeddings with documentIds filter deletes all embeddings for that document
    await client.delete_embeddings(
        vector_store_id=vector_store_id,
        delete_request={
            "documentIds": [doc_id]  # Deletes all embeddings with this documentId
        }
    )
    print(f"🗑️ Deleted old embeddings for document: {doc_id}")
    
    # Step 2: Generate new embedding
    # Supported models: text-embedding-3-small, text-embedding-3-large, titan-text-embeddings-v2
    # Supported dimensions: 1024, 1536, 3072 (must match vector store configuration)
    query_embedding = await client.text_to_embedding(
        text=new_text,
        vector_store_id=vector_store_id
    )
    print(f"🔄 Generated new embedding for document: {doc_id}")
    
    # Step 3: Add new embedding to vector store
    await client.create_embeddings(
        vector_store_id=vector_store_id,
        embedding_data={
            "embeddings": [query_embedding],
            "chunks": [new_text],
            "documentIds": [doc_id],
            "sequences": [0],
            "metadata": [metadata]
        }
    )
    print(f"✅ Updated document: {doc_id}")

3. Monitoring and Logging

import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@agent.tool_plain
async def monitored_search(query: str) -> str:
    """Search with logging for monitoring."""
    logger.info(f"Search query: {query}")
    
    try:
        result = await kb_tool.search(query=query, search_type="hybrid")
        logger.info(f"Search successful, returned results")
        return result
    except Exception as e:
        logger.error(f"Search failed: {str(e)}", exc_info=True)
        raise

---

Troubleshooting

No Results Found

Problem: Agent returns "No relevant documentation found"

Solutions:

1. Check that documents were indexed successfully
2. Lower the score_threshold
3. Try search_type="hybrid" for better coverage
4. Verify metadata filters aren't too restrictive

Irrelevant Results

Problem: Search returns unrelated documents

Solutions:

1. Increase score_threshold (e.g., from 0.5 to 0.7)
2. Reduce top_k to get only best matches
3. Add metadata filters to narrow search scope
4. Use search_type="vector" for semantic relevance

Slow Performance

Problem: Searches take too long

Solutions:

1. Reduce top_k value
2. Use specific metadata filters
3. Consider search_type="keyword" for simple queries
4. Implement caching for common queries

---

Next Steps

• RAG & Vector Stores - Complete RAG documentation
• Search & Filtering - Learn about search strategies
• Tools - Learn about KnowledgeBaseSearchTool in detail
• Quick Start - Build your first agent
• Deployment Guide - Deploy your agent to production
• Documentation Assistant - See another complete example

---

Additional Resources

• CAIP Agents SDK: Repository