Claude RAG: Build a Retrieval-Augmented Generation App

import anthropic
import chromadb
from sentence_transformers import SentenceTransformer
from pathlib import Path
import pypdf
import json
import hashlib
import re
from typing import Optional

# Initialise clients
anthropic_client = anthropic.Anthropic()
chroma_client = chromadb.PersistentClient(path="./chroma_db")
embedder = SentenceTransformer("all-MiniLM-L6-v2")  # Fast, accurate local model


# ─── Document Ingestion ───────────────────────────────────────────────────────

def extract_text_from_pdf(pdf_path: str) -> str:
    """Extract text from a PDF file."""
    reader = pypdf.PdfReader(pdf_path)
    text = ""
    for page in reader.pages:
        text += page.extract_text() + "\n\n"
    return text


def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> list[str]:
    """
    Split text into overlapping chunks.
    
    Args:
        text: Full document text
        chunk_size: Approximate words per chunk
        overlap: Words of overlap between consecutive chunks
    """
    # Clean and normalise whitespace
    text = re.sub(r'\s+', ' ', text).strip()
    words = text.split()
    
    chunks = []
    start = 0
    
    while start < len(words):
        end = min(start + chunk_size, len(words))
        chunk = " ".join(words[start:end])
        
        # Try to end at a sentence boundary
        if end < len(words):
            last_period = chunk.rfind('. ')
            if last_period > chunk_size * 0.6:  # Only if not too short
                chunk = chunk[:last_period + 1]
        
        chunks.append(chunk)
        start += chunk_size - overlap
    
    return chunks


def ingest_document(
    file_path: str,
    collection_name: str,
    document_metadata: dict = None
) -> int:
    """
    Ingest a document into the vector database.
    
    Returns the number of chunks added.
    """
    path = Path(file_path)
    metadata = document_metadata or {}
    
    # Extract text
    if path.suffix.lower() == ".pdf":
        text = extract_text_from_pdf(file_path)
    elif path.suffix.lower() in [".txt", ".md"]:
        text = path.read_text(encoding="utf-8")
    else:
        raise ValueError(f"Unsupported file type: {path.suffix}")
    
    # Split into chunks
    chunks = chunk_text(text)
    print(f"  {path.name}: {len(chunks)} chunks from {len(text)} characters")
    
    # Get or create collection
    collection = chroma_client.get_or_create_collection(
        name=collection_name,
        metadata={"hnsw:space": "cosine"}
    )
    
    # Generate embeddings
    embeddings = embedder.encode(chunks, batch_size=32, show_progress_bar=False)
    
    # Prepare metadata for each chunk
    doc_hash = hashlib.md5(path.name.encode()).hexdigest()[:8]
    
    chunk_ids = [f"{doc_hash}_{i}" for i in range(len(chunks))]
    chunk_metadata = [
        {
            **metadata,
            "source": path.name,
            "chunk_index": i,
            "total_chunks": len(chunks)
        }
        for i in range(len(chunks))
    ]
    
    # Store in ChromaDB
    collection.add(
        ids=chunk_ids,
        embeddings=embeddings.tolist(),
        documents=chunks,
        metadatas=chunk_metadata
    )
    
    return len(chunks)


def ingest_directory(
    directory: str,
    collection_name: str,
    extensions: list = None
) -> int:
    """Ingest all documents in a directory."""
    exts = extensions or [".pdf", ".txt", ".md"]
    total_chunks = 0
    
    for path in Path(directory).iterdir():
        if path.suffix.lower() in exts:
            print(f"Ingesting: {path.name}")
            chunks = ingest_document(str(path), collection_name)
            total_chunks += chunks
    
    return total_chunks


# ─── Retrieval ────────────────────────────────────────────────────────────────

def retrieve_relevant_chunks(
    query: str,
    collection_name: str,
    n_results: int = 5,
    min_relevance_score: float = 0.3
) -> list[dict]:
    """
    Find the most relevant document chunks for a query.
    
    Returns list of dicts with content, source, and relevance score.
    """
    collection = chroma_client.get_or_create_collection(collection_name)
    
    if collection.count() == 0:
        return []
    
    # Embed query
    query_embedding = embedder.encode([query])[0]
    
    # Search
    results = collection.query(
        query_embeddings=[query_embedding.tolist()],
        n_results=min(n_results, collection.count()),
        include=["documents", "metadatas", "distances"]
    )
    
    # Convert distances to similarity scores (cosine distance to similarity)
    chunks = []
    for doc, metadata, distance in zip(
        results["documents"][0],
        results["metadatas"][0],
        results["distances"][0]
    ):
        similarity = 1 - distance  # Cosine similarity
        
        if similarity >= min_relevance_score:
            chunks.append({
                "content": doc,
                "source": metadata.get("source", "Unknown"),
                "chunk_index": metadata.get("chunk_index", 0),
                "similarity": round(similarity, 3)
            })
    
    # Sort by similarity (highest first)
    chunks.sort(key=lambda x: x["similarity"], reverse=True)
    return chunks


# ─── Generation (RAG Answer) ─────────────────────────────────────────────────

def answer_question(
    question: str,
    collection_name: str,
    n_chunks: int = 5,
    model: str = "claude-sonnet-4-6"
) -> dict:
    """
    Answer a question using RAG.
    
    Returns the answer and the source documents used.
    """
    # Retrieve relevant chunks
    chunks = retrieve_relevant_chunks(question, collection_name, n_results=n_chunks)
    
    if not chunks:
        return {
            "answer": "I could not find relevant information in the document library to answer this question.",
            "sources": [],
            "chunks_used": 0
        }
    
    # Build context block
    context_sections = []
    for i, chunk in enumerate(chunks, 1):
        context_sections.append(
            f"[Document: {chunk['source']}, Relevance: {chunk['similarity']}]\n{chunk['content']}"
        )
    
    context = "\n\n---\n\n".join(context_sections)
    
    # Generate answer with Claude
    response = anthropic_client.messages.create(
        model=model,
        max_tokens=2048,
        system="""You are a helpful assistant that answers questions based strictly on provided document excerpts.

Rules:
1. Answer ONLY from the provided context below. Do not use knowledge from your training.
2. If the context does not contain enough information to answer the question, say so clearly.
3. Always cite which document(s) you used in your answer using [Source: filename] notation.
4. If information from multiple documents is relevant, synthesise it clearly.
""",
        messages=[
            {
                "role": "user",
                "content": f"""RETRIEVED DOCUMENT CONTEXT:
{context}

---

QUESTION: {question}

Answer the question based solely on the context above."""
            }
        ]
    )
    
    answer = response.content[0].text
    unique_sources = list(set(chunk["source"] for chunk in chunks))
    
    return {
        "answer": answer,
        "sources": unique_sources,
        "chunks_used": len(chunks),
        "retrieved_chunks": chunks
    }


# ─── Interactive RAG Interface ────────────────────────────────────────────────

class RAGApp:
    """Simple interactive RAG application."""
    
    def __init__(self, collection_name: str):
        self.collection_name = collection_name
    
    def ingest(self, path: str):
        """Ingest a file or directory."""
        p = Path(path)
        if p.is_dir():
            count = ingest_directory(str(p), self.collection_name)
        else:
            count = ingest_document(str(p), self.collection_name)
        print(f"Ingested {count} chunks total.")
    
    def ask(self, question: str) -> str:
        """Ask a question and get a grounded answer."""
        result = answer_question(question, self.collection_name)
        
        print(f"\nAnswer: {result['answer']}")
        print(f"\nSources used: {', '.join(result['sources'])}")
        print(f"Chunks retrieved: {result['chunks_used']}")
        
        return result["answer"]
    
    def run_interactive(self):
        """Run an interactive Q&A session."""
        print(f"\nRAG System ready. Collection: {self.collection_name}")
        print("Type your question or 'quit' to exit.\n")
        
        while True:
            question = input("Question: ").strip()
            if question.lower() in ("quit", "exit"):
                break
            if question:
                self.ask(question)
                print()


# ─── Example Usage ────────────────────────────────────────────────────────────

if __name__ == "__main__":
    # Create RAG app for company documentation
    app = RAGApp(collection_name="company_docs")
    
    # Ingest documents
    app.ingest("./documents/")  # Point to your document folder
    
    # Run interactive Q&A
    app.run_interactive()