MCP, RAG & CAG

┌─────────────────────────────────────────────────────────────────┐
│                        RAG Pipeline Flow                        │
└─────────────────────────────────────────────────────────────────┘

    ┌─────────────────┐
    │   User Query    │  "What is our vacation policy?"
    └────────┬────────┘
             │
             ▼
    ┌─────────────────┐
    │ 1. Embedding    │  Convert query → vector representation
    │    Model        │  [0.234, -0.891, 0.456, ..., 0.123]
    └────────┬────────┘
             │
             ▼
    ┌─────────────────┐
    │ 2. Vector       │  Similarity search in vector database
    │    Search       │  (cosine similarity, dot product, etc.)
    └────────┬────────┘
             │
             ▼
    ┌─────────────────┐
    │ 3. Retrieve     │  Get top K most relevant documents
    │    Top K Docs   │  (K = 3-5 typically)
    └────────┬────────┘
             │
             ▼
    ┌─────────────────┐
    │ 4. Build        │  Query + Retrieved Docs → Prompt
    │    Context      │  Format for LLM consumption
    └────────┬────────┘
             │
             ▼
    ┌─────────────────┐
    │ 5. LLM          │  Generate answer using context
    │    Generation   │  (Temperature: 0.0-0.3 for factual)
    └────────┬────────┘
             │
             ▼
    ┌─────────────────┐
    │ 6. Response     │  Return answer + source citations
    │    with Sources │  to user
    └─────────────────┘

User Query:
"What is our company's vacation policy?"

Step 1 - Embedding:
  Input:  "What is our company's vacation policy?"
  Output: [0.234, -0.891, 0.456, 0.123, ..., -0.567]  (384 dimensions)

Step 2 - Vector Search:
  Search ChromaDB for similar document embeddings
  Method: Cosine similarity

Step 3 - Retrieve Top Documents:
  ✓ Document 1: "Employee Handbook - Section 5.2"
    Similarity: 0.94
    Preview: "Vacation Policy: Full-time employees receive 20 days..."

  ✓ Document 2: "HR Policy Update 2024"
    Similarity: 0.87
    Preview: "Updated vacation rollover policy effective Jan 2024..."

  ✓ Document 3: "Benefits Guide"
    Similarity: 0.82
    Preview: "Comprehensive benefits including vacation, health..."

Step 4 - Build Context:
  Prompt = """
  Context:
  [Document 1 content]
  [Document 2 content]
  [Document 3 content]

  Question: What is our company's vacation policy?
  Answer based only on the context above.
  """

Step 5 - LLM Generation:
  Model: claude-sonnet-4-20250514
  Temperature: 0.3 (factual responses)
  Max tokens: 1000

  Output: "Based on the Employee Handbook Section 5.2, full-time
          employees receive 20 days of paid vacation per year,
          accruing at 1.67 days per month. According to the 2024
          HR Policy Update, unused vacation can roll over up to
          10 days annually..."

Step 6 - Return Response:
  Answer: [Generated text above]
  Sources:
    - Employee Handbook - Section 5.2 (94% match)
    - HR Policy Update 2024 (87% match)
    - Benefits Guide (82% match)

  Total latency: ~800ms
    - Embedding: 50ms
    - Vector search: 100ms
    - LLM generation: 650ms

# rag_system.py - Complete RAG implementation

from typing import List, Dict
import numpy as np
from sentence_transformers import SentenceTransformer
import chromadb
from anthropic import Anthropic

class RAGSystem:
    """
    Production RAG system with embeddings, vector search, and LLM generation.
    """

    def __init__(
        self,
        embedding_model: str = "all-MiniLM-L6-v2",
        llm_model: str = "claude-sonnet-4-20250514",
        api_key: str = None
    ):
        # Embedding model for converting text to vectors
        self.embedder = SentenceTransformer(embedding_model)

        # Vector database for similarity search
        self.chroma_client = chromadb.Client()
        self.collection = self.chroma_client.create_collection(
            name="documents",
            metadata={"hnsw:space": "cosine"}
        )

        # LLM for generation
        self.llm = Anthropic(api_key=api_key)
        self.llm_model = llm_model

        self.doc_count = 0

    def add_documents(self, documents: List[Dict[str, str]]):
        """
        Add documents to the RAG system.

        Args:
            documents: List of dicts with 'text', 'metadata' (title, source, etc.)
        """
        texts = [doc['text'] for doc in documents]
        metadatas = [doc.get('metadata', {}) for doc in documents]

        # Generate embeddings
        embeddings = self.embedder.encode(texts, show_progress_bar=True)

        # Add to vector database
        ids = [f"doc_{self.doc_count + i}" for i in range(len(documents))]
        self.collection.add(
            embeddings=embeddings.tolist(),
            documents=texts,
            metadatas=metadatas,
            ids=ids
        )

        self.doc_count += len(documents)
        print(f"Added {len(documents)} documents. Total: {self.doc_count}")

    def retrieve(self, query: str, top_k: int = 3) -> List[Dict]:
        """
        Retrieve most relevant documents for a query.

        Args:
            query: User question
            top_k: Number of documents to retrieve

        Returns:
            List of retrieved documents with metadata and scores
        """
        # Embed query
        query_embedding = self.embedder.encode([query])[0]

        # Search vector database
        results = self.collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=top_k
        )

        # Format results
        retrieved_docs = []
        for i in range(len(results['ids'][0])):
            retrieved_docs.append({
                'id': results['ids'][0][i],
                'text': results['documents'][0][i],
                'metadata': results['metadatas'][0][i],
                'similarity': 1 - results['distances'][0][i]  # Convert distance to similarity
            })

        return retrieved_docs

    def generate_answer(
        self,
        query: str,
        retrieved_docs: List[Dict],
        include_sources: bool = True
    ) -> Dict:
        """
        Generate answer using retrieved context.

        Args:
            query: User question
            retrieved_docs: Documents from retrieve()
            include_sources: Include source citations

        Returns:
            Answer with metadata
        """
        # Build context from retrieved documents
        context = "\n\n".join([
            f"Source {i+1} ({doc['metadata'].get('title', 'Unknown')}):\n{doc['text']}"
            for i, doc in enumerate(retrieved_docs)
        ])

        # Create prompt with context
        system_prompt = """You are a helpful assistant that answers questions based on the provided context.

Guidelines:
1. Answer ONLY using information from the provided context
2. If the context doesn't contain enough information, say so
3. Cite sources when possible (Source 1, Source 2, etc.)
4. Be concise but complete"""

        user_prompt = f"""Context:
{context}

Question: {query}

Answer the question based on the context above."""

        # Call LLM
        response = self.llm.messages.create(
            model=self.llm_model,
            max_tokens=1000,
            temperature=0.3,  # Lower temperature for factual responses
            system=system_prompt,
            messages=[{"role": "user", "content": user_prompt}]
        )

        answer = response.content[0].text

        # Build result
        result = {
            'answer': answer,
            'query': query,
            'sources': [
                {
                    'title': doc['metadata'].get('title', 'Unknown'),
                    'similarity': doc['similarity'],
                    'excerpt': doc['text'][:200] + '...'
                }
                for doc in retrieved_docs
            ] if include_sources else []
        }

        return result

    def query(self, question: str, top_k: int = 3) -> Dict:
        """
        Complete RAG pipeline: retrieve + generate.

        Args:
            question: User question
            top_k: Number of documents to retrieve

        Returns:
            Answer with sources
        """
        # Step 1: Retrieve relevant documents
        retrieved_docs = self.retrieve(question, top_k=top_k)

        # Step 2: Generate answer
        result = self.generate_answer(question, retrieved_docs)

        return result


# Example Usage
if __name__ == "__main__":
    import os

    # Initialize RAG system
    rag = RAGSystem(api_key=os.getenv("ANTHROPIC_API_KEY"))

    # Add company documents
    documents = [
        {
            "text": """Vacation Policy: Full-time employees receive 20 days of paid vacation per year.
            Vacation accrues monthly at 1.67 days per month. Unused vacation rolls over up to 10 days.
            Employees must request vacation at least 2 weeks in advance except for emergencies.""",
            "metadata": {"title": "Employee Handbook - Section 5.2", "source": "HR"}
        },
        {
            "text": """Remote Work Policy: Employees may work remotely up to 3 days per week.
            Remote work requires manager approval and must be scheduled in advance.
            Employees must be available during core hours (10 AM - 3 PM) regardless of location.""",
            "metadata": {"title": "Remote Work Guidelines", "source": "HR"}
        },
        {
            "text": """Health Benefits: Company provides comprehensive health insurance through BlueCross.
            Coverage starts on the first day of the month following hire date.
            Premium is $150/month for individuals, $350/month for families.
            Annual deductible: $1,000 individual, $2,000 family.""",
            "metadata": {"title": "Benefits Guide 2024", "source": "Benefits"}
        },
        {
            "text": """Performance Review Process: Reviews conducted annually in December.
            Employees receive feedback on goals, competencies, and overall performance.
            Reviews determine merit increases and promotion eligibility.
            Self-assessment due by December 1st, manager review by December 15th.""",
            "metadata": {"title": "Performance Management", "source": "HR"}
        }
    ]

    rag.add_documents(documents)

    # Ask questions
    questions = [
        "How many vacation days do I get?",
        "What is the remote work policy?",
        "When are performance reviews?",
        "How much does health insurance cost?"
    ]

    for question in questions:
        print(f"\n{'='*60}")
        print(f"Q: {question}")
        print('='*60)

        result = rag.query(question, top_k=2)

        print(f"\nA: {result['answer']}")
        print(f"\nSources:")
        for i, source in enumerate(result['sources'], 1):
            print(f"  {i}. {source['title']} (similarity: {source['similarity']:.2f})")

# Output example:
# Q: How many vacation days do I get?
# A: Full-time employees receive 20 days of paid vacation per year,
#    which accrues monthly at 1.67 days per month (Source 1).
# Sources:
#   1. Employee Handbook - Section 5.2 (similarity: 0.94)
#   2. Benefits Guide 2024 (similarity: 0.45)

MCP Architecture:

┌─────────────────────────────────────────────────────────┐
│                      AI Application                     │
│              (Claude, custom chatbot, etc.)             │
└────────────────────┬────────────────────────────────────┘
                     │ MCP Protocol
                     │ (Standardized communication)
                     │
     ┌───────────────┼───────────────┐
     │               │               │
     ▼               ▼               ▼
┌─────────┐    ┌─────────┐    ┌─────────┐
│  MCP    │    │  MCP    │    │  MCP    │
│ Server  │    │ Server  │    │ Server  │
│ (DB)    │    │ (Files) │    │ (API)   │
└─────────┘    └─────────┘    └─────────┘
     │               │               │
     ▼               ▼               ▼
┌─────────┐    ┌─────────┐    ┌─────────┐
│Database │    │File     │    │External │
│         │    │System   │    │API      │
└─────────┘    └─────────┘    └─────────┘

Key Concepts:

1. Resources: Read-only data sources (files, DB records, API responses)
2. Tools: Actions the AI can take (create, update, delete, execute)
3. Prompts: Pre-built prompt templates with variables
4. Sampling: Request AI to generate content (for complex workflows)

Benefits:
✓ Standardized protocol (not vendor-specific)
✓ Composable (multiple MCP servers work together)
✓ Secure (fine-grained access control)
✓ Extensible (easy to add new integrations)

# mcp_server.py - MCP server for database access
# Uses the official MCP Python SDK (pip install mcp)

from mcp.server.fastmcp import FastMCP
import sqlite3
import json

# Initialize MCP server with the high-level FastMCP API
mcp = FastMCP("company-db-server")

# Database connection
DB_PATH = "company.db"


def get_connection():
    return sqlite3.connect(DB_PATH)


@mcp.resource("db://tables")
def list_tables() -> str:
    """List all available database tables."""
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = [row[0] for row in cursor.fetchall()]
    conn.close()
    return json.dumps({"tables": tables}, indent=2)


@mcp.resource("db://tables/{table_name}")
def read_table(table_name: str) -> str:
    """Read data from a specific table (limited to 100 rows)."""
    conn = get_connection()
    cursor = conn.cursor()
    # Use parameterized allowlist to prevent SQL injection
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
    if not cursor.fetchone():
        conn.close()
        return json.dumps({"error": f"Table '{table_name}' not found"})

    cursor.execute(f"SELECT * FROM [{table_name}] LIMIT 100")
    columns = [description[0] for description in cursor.description]
    rows = cursor.fetchall()
    conn.close()

    data = [dict(zip(columns, row)) for row in rows]
    return json.dumps(data, indent=2)


@mcp.tool()
def query_database(query: str) -> str:
    """
    Execute a read-only SQL query on the database.

    Args:
        query: SQL SELECT query to execute
    """
    # Safety: only allow SELECT queries
    if not query.strip().upper().startswith("SELECT"):
        return json.dumps({"error": "Only SELECT queries are allowed"})

    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(query)
        columns = [description[0] for description in cursor.description]
        rows = cursor.fetchall()

        data = [dict(zip(columns, row)) for row in rows]
        return json.dumps({
            "query": query,
            "row_count": len(data),
            "results": data
        }, indent=2)

    except Exception as e:
        return json.dumps({"error": f"Query failed: {str(e)}"})
    finally:
        conn.close()


@mcp.tool()
def get_schema(table_name: str) -> str:
    """
    Get the schema (columns, types) of a database table.

    Args:
        table_name: Name of the table to inspect
    """
    conn = get_connection()
    cursor = conn.cursor()
    # Validate table exists first
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
    if not cursor.fetchone():
        conn.close()
        return json.dumps({"error": f"Table '{table_name}' not found"})

    cursor.execute(f"PRAGMA table_info([{table_name}])")
    schema = cursor.fetchall()
    conn.close()

    columns = []
    for col in schema:
        columns.append({
            "name": col[1],
            "type": col[2],
            "nullable": not col[3],
            "primary_key": bool(col[5])
        })

    return json.dumps({"table": table_name, "columns": columns}, indent=2)


# Run the server (uses stdio transport by default)
if __name__ == "__main__":
    mcp.run()

# Claude can now:
# 1. List resources: "Show me all tables"       → reads db://tables
# 2. Read resources: "Get employees data"        → reads db://tables/employees
# 3. Call tools: "Query engineering department"   → calls query_database()
# 4. Inspect schema: "What columns in users?"    → calls get_schema()
#
# Example conversation:
# User: "How many employees in each department?"
# Claude (internally):
#   1. Calls get_schema("employees") to understand structure
#   2. Calls query_database("SELECT department, COUNT(*) FROM employees GROUP BY department")
#   3. Returns formatted answer to user
#
# Response: "Based on the database:
#   - Engineering: 45 employees
#   - Sales: 23 employees
#   - Marketing: 12 employees"
#
# Works with any MCP-compatible AI system (Claude Desktop, Claude Code, etc.)

# cag_system.py - Comprehensive Context-Augmented Generation

from typing import Dict, List, Optional
from datetime import datetime
import anthropic

class CAGSystem:
    """
    Complete CAG system with multiple context sources.
    """

    def __init__(self, api_key: str):
        self.llm = anthropic.Anthropic(api_key=api_key)
        self.conversation_history = {}  # Per-user history

    def generate_response(
        self,
        user_id: str,
        query: str,
        user_profile: Dict,
        include_history: bool = True
    ) -> str:
        """
        Generate response with full context augmentation.

        Args:
            user_id: User identifier
            query: User question
            user_profile: User metadata (name, role, preferences)
            include_history: Include conversation history

        Returns:
            AI response with context
        """
        # Build comprehensive context
        context = self._build_context(user_id, query, user_profile)

        # Build messages with history
        messages = []

        if include_history and user_id in self.conversation_history:
            messages.extend(self.conversation_history[user_id])

        messages.append({
            "role": "user",
            "content": f"{context}\n\nUser question: {query}"
        })

        # Generate response
        response = self.llm.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=2000,
            temperature=0.7,
            system=self._build_system_prompt(user_profile),
            messages=messages
        )

        assistant_message = response.content[0].text

        # Update conversation history
        self._update_history(user_id, query, assistant_message)

        return assistant_message

    def _build_context(
        self,
        user_id: str,
        query: str,
        user_profile: Dict
    ) -> str:
        """
        Intelligently build context from multiple sources.
        """
        context_sections = []

        # 1. Temporal context
        context_sections.append(f"""
=== TEMPORAL CONTEXT ===
Current date/time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
Day of week: {datetime.now().strftime('%A')}
""")

        # 2. User context
        context_sections.append(f"""
=== USER CONTEXT ===
Name: {user_profile.get('name')}
Role: {user_profile.get('role')}
Department: {user_profile.get('department')}
Preferences: {user_profile.get('preferences', {})}
""")

        # 3. Domain-specific context (example: project management)
        if 'current_project' in user_profile:
            context_sections.append(f"""
=== PROJECT CONTEXT ===
Active project: {user_profile['current_project']['name']}
Status: {user_profile['current_project']['status']}
Team members: {', '.join(user_profile['current_project']['team'])}
Deadline: {user_profile['current_project']['deadline']}
""")

        # 4. Retrieved knowledge (RAG-style)
        if self._needs_knowledge_retrieval(query):
            relevant_docs = self._retrieve_knowledge(query)
            context_sections.append(f"""
=== RELEVANT KNOWLEDGE ===
{relevant_docs}
""")

        # 5. Real-time data (if applicable)
        if self._needs_realtime_data(query):
            realtime_data = self._fetch_realtime_data(query)
            context_sections.append(f"""
=== REAL-TIME DATA ===
{realtime_data}
""")

        return "\n".join(context_sections)

    def _build_system_prompt(self, user_profile: Dict) -> str:
        """Build personalized system prompt."""
        return f"""You are a helpful AI assistant for {user_profile.get('name')}.

Your role:
- Provide accurate, helpful responses
- Use the provided context to ground your answers
- Be concise but complete
- Cite sources when using retrieved information
- Respect user's role and permissions

User's role: {user_profile.get('role')}
Tone: {user_profile.get('preferences', {}).get('tone', 'professional')}"""

    def _update_history(self, user_id: str, user_msg: str, assistant_msg: str):
        """Maintain conversation history."""
        if user_id not in self.conversation_history:
            self.conversation_history[user_id] = []

        self.conversation_history[user_id].extend([
            {"role": "user", "content": user_msg},
            {"role": "assistant", "content": assistant_msg}
        ])

        # Keep only last 10 turns
        if len(self.conversation_history[user_id]) > 20:  # 20 messages = 10 turns
            self.conversation_history[user_id] = self.conversation_history[user_id][-20:]

    def _needs_knowledge_retrieval(self, query: str) -> bool:
        """Determine if query needs knowledge base search."""
        knowledge_keywords = ['policy', 'procedure', 'guideline', 'documentation', 'how to']
        return any(kw in query.lower() for kw in knowledge_keywords)

    def _retrieve_knowledge(self, query: str) -> str:
        """Retrieve relevant knowledge (simplified)."""
        # In production: use RAG system
        return "[Knowledge base results would go here]"

    def _needs_realtime_data(self, query: str) -> bool:
        """Determine if query needs real-time data."""
        realtime_keywords = ['current', 'latest', 'now', 'today', 'status']
        return any(kw in query.lower() for kw in realtime_keywords)

    def _fetch_realtime_data(self, query: str) -> str:
        """Fetch real-time data (simplified)."""
        # In production: call MCP servers or APIs
        return "[Real-time data would go here]"


# Example Usage
if __name__ == "__main__":
    import os

    cag = CAGSystem(api_key=os.getenv("ANTHROPIC_API_KEY"))

    # User profile
    user_profile = {
        "name": "Alice Johnson",
        "role": "Senior Engineer",
        "department": "Engineering",
        "current_project": {
            "name": "API Redesign",
            "status": "In Progress",
            "team": ["Alice", "Bob", "Carol"],
            "deadline": "2024-03-15"
        },
        "preferences": {
            "tone": "professional"
        }
    }

    # Example conversation
    queries = [
        "What's my current project status?",
        "Who's on my team?",
        "When is the deadline?"
    ]

    for query in queries:
        print(f"\nUser: {query}")
        response = cag.generate_response(
            user_id="alice_123",
            query=query,
            user_profile=user_profile
        )
        print(f"AI: {response}")

# Output:
# User: What's my current project status?
# AI: Your current project "API Redesign" is In Progress. The deadline
#     is March 15, 2024, and you're working with Bob and Carol.
#
# User: Who's on my team?
# AI: Your team for the API Redesign project includes: Alice (you),
#     Bob, and Carol.
#
# User: When is the deadline?
# AI: The deadline for your API Redesign project is March 15, 2024.