Claude AI Tutorial #22: Build a Document Q&A System with RAG

Upload PDF → Parse → Chunk → Embed → Store in Vector DB
                                         ↓
Question → Embed → Search → Top Chunks → Claude → Cited Answer

mkdir document-qa && cd document-qa
pip install anthropic openai chromadb pymupdf

mkdir document-qa && cd document-qa
npm init -y
npm install @anthropic-ai/sdk openai chromadb pdf-parse

import fitz  # PyMuPDF

def parse_pdf(pdf_path: str) -> list[dict]:
    """Parse a PDF and return a list of pages with text."""
    doc = fitz.open(pdf_path)
    pages = []

    for page_num in range(len(doc)):
        page = doc[page_num]
        text = page.get_text()

        if text.strip():
            pages.append({
                "page": page_num + 1,
                "text": text.strip(),
                "source": pdf_path,
            })

    doc.close()
    return pages

# Parse a PDF
pages = parse_pdf("company-handbook.pdf")
print(f"Parsed {len(pages)} pages")

import pdf from "pdf-parse";
import { readFileSync } from "fs";

interface Page {
  page: number;
  text: string;
  source: string;
}

async function parsePdf(pdfPath: string): Promise<Page[]> {
  const buffer = readFileSync(pdfPath);
  const data = await pdf(buffer);

  // pdf-parse returns all text at once; split by page breaks
  const pageTexts = data.text.split("\f").filter((t) => t.trim());

  return pageTexts.map((text, i) => ({
    page: i + 1,
    text: text.trim(),
    source: pdfPath,
  }));
}

from openai import OpenAI

openai_client = OpenAI()

def chunk_text(text: str, chunk_size: int = 500, overlap: int = 100) -> list[str]:
    """Split text into overlapping chunks."""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]

        if end < len(text):
            last_period = chunk.rfind(".")
            last_newline = chunk.rfind("\n")
            break_point = max(last_period, last_newline)
            if break_point > chunk_size * 0.3:
                chunk = chunk[:break_point + 1]
                end = start + break_point + 1

        chunks.append(chunk.strip())
        start = end - overlap

    return chunks

def get_embeddings(texts: list[str]) -> list[list[float]]:
    """Get embeddings using OpenAI."""
    response = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=texts,
    )
    return [item.embedding for item in response.data]

def process_pages(pages: list[dict]) -> list[dict]:
    """Chunk pages and create embeddings."""
    all_chunks = []

    for page in pages:
        chunks = chunk_text(page["text"])
        for i, chunk in enumerate(chunks):
            all_chunks.append({
                "text": chunk,
                "source": page["source"],
                "page": page["page"],
                "chunk_index": i,
            })

    # Batch embed all chunks
    texts = [c["text"] for c in all_chunks]
    embeddings = get_embeddings(texts)

    for i, embedding in enumerate(embeddings):
        all_chunks[i]["embedding"] = embedding

    return all_chunks

import chromadb

chroma = chromadb.PersistentClient(path="./chroma_db")
collection = chroma.get_or_create_collection(
    name="documents",
    metadata={"hnsw:space": "cosine"},
)

def ingest_pdf(pdf_path: str):
    """Full ingestion pipeline: parse → chunk → embed → store."""
    print(f"Parsing {pdf_path}...")
    pages = parse_pdf(pdf_path)

    print(f"Chunking and embedding {len(pages)} pages...")
    chunks = process_pages(pages)

    print(f"Storing {len(chunks)} chunks...")
    collection.add(
        ids=[f"{pdf_path}_p{c['page']}_c{c['chunk_index']}" for c in chunks],
        embeddings=[c["embedding"] for c in chunks],
        documents=[c["text"] for c in chunks],
        metadatas=[{
            "source": c["source"],
            "page": c["page"],
            "chunk_index": c["chunk_index"],
        } for c in chunks],
    )
    print(f"Done. Ingested {len(chunks)} chunks from {pdf_path}")

# Ingest documents
ingest_pdf("company-handbook.pdf")
ingest_pdf("product-manual.pdf")

import psycopg2
from pgvector.psycopg2 import register_vector

conn = psycopg2.connect("postgresql://user:pass@localhost/docqa")
register_vector(conn)

cur = conn.cursor()

# Create table
cur.execute("""
    CREATE TABLE IF NOT EXISTS chunks (
        id SERIAL PRIMARY KEY,
        text TEXT NOT NULL,
        source TEXT NOT NULL,
        page INTEGER,
        chunk_index INTEGER,
        embedding vector(1536)
    )
""")

# Create index
cur.execute("""
    CREATE INDEX IF NOT EXISTS chunks_embedding_idx
    ON chunks USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 100)
""")

def store_chunks_pgvector(chunks: list[dict]):
    """Store chunks in pgvector."""
    for chunk in chunks:
        cur.execute(
            "INSERT INTO chunks (text, source, page, chunk_index, embedding) VALUES (%s, %s, %s, %s, %s)",
            (chunk["text"], chunk["source"], chunk["page"], chunk["chunk_index"], chunk["embedding"]),
        )
    conn.commit()

def search_pgvector(query_embedding: list[float], top_k: int = 5) -> list[dict]:
    """Search for similar chunks in pgvector."""
    cur.execute(
        "SELECT text, source, page, 1 - (embedding <=> %s) AS similarity FROM chunks ORDER BY embedding <=> %s LIMIT %s",
        (query_embedding, query_embedding, top_k),
    )
    rows = cur.fetchall()
    return [{"text": r[0], "source": r[1], "page": r[2], "similarity": r[3]} for r in rows]

import anthropic

claude = anthropic.Anthropic()

def ask(question: str, top_k: int = 5) -> dict:
    """Ask a question about your documents."""
    # Step 1: Embed the question
    query_embedding = get_embeddings([question])[0]

    # Step 2: Find relevant chunks
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
    )

    chunks = results["documents"][0]
    metadatas = results["metadatas"][0]

    # Step 3: Build context with source references
    context_parts = []
    for i, (chunk, meta) in enumerate(zip(chunks, metadatas)):
        context_parts.append(
            f"[Source {i+1}: {meta['source']}, Page {meta['page']}]\n{chunk}"
        )
    context = "\n\n---\n\n".join(context_parts)

    # Step 4: Ask Claude
    response = claude.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        system="""You are a document Q&A assistant. Answer questions based ONLY on the provided context.

<rules>
- Answer only from the provided sources
- Cite sources using [Source N] notation
- If the answer is not in the context, say "I could not find this in the provided documents"
- Be concise and direct
- If multiple sources support the answer, cite all of them
</rules>""",
        messages=[
            {
                "role": "user",
                "content": f"<context>\n{context}\n</context>\n\n<question>\n{question}\n</question>",
            }
        ],
    )

    answer = response.content[0].text

    return {
        "answer": answer,
        "sources": [
            {"source": m["source"], "page": m["page"]}
            for m in metadatas
        ],
        "tokens": {
            "input": response.usage.input_tokens,
            "output": response.usage.output_tokens,
        },
    }

# Ask questions
result = ask("What is the company vacation policy?")
print(f"Answer: {result['answer']}")
print(f"\nSources:")
for s in result["sources"]:
    print(f"  - {s['source']}, Page {s['page']}")

import Anthropic from "@anthropic-ai/sdk";

const claude = new Anthropic();

interface QAResult {
  answer: string;
  sources: Array<{ source: string; page: number }>;
  tokens: { input: number; output: number };
}

async function ask(question: string, topK: number = 5): Promise<QAResult> {
  const queryEmbedding = (await getEmbeddings([question]))[0];

  const collection = await chroma.getCollection({ name: "documents" });
  const results = await collection.query({
    queryEmbeddings: [queryEmbedding],
    nResults: topK,
  });

  const chunks = results.documents[0];
  const metadatas = results.metadatas[0];

  const context = chunks
    .map(
      (chunk, i) =>
        `[Source ${i + 1}: ${metadatas[i]?.source}, Page ${metadatas[i]?.page}]\n${chunk}`
    )
    .join("\n\n---\n\n");

  const response = await claude.messages.create({
    model: "claude-sonnet-4-6",
    max_tokens: 2048,
    system: `You are a document Q&A assistant. Answer questions based ONLY on the provided context.

<rules>
- Answer only from the provided sources
- Cite sources using [Source N] notation
- If the answer is not in the context, say "I could not find this in the provided documents"
- Be concise and direct
</rules>`,
    messages: [
      {
        role: "user",
        content: `<context>\n${context}\n</context>\n\n<question>\n${question}\n</question>`,
      },
    ],
  });

  const answer =
    response.content[0].type === "text" ? response.content[0].text : "";

  return {
    answer,
    sources: metadatas.map((m: any) => ({ source: m.source, page: m.page })),
    tokens: {
      input: response.usage.input_tokens,
      output: response.usage.output_tokens,
    },
  };
}

class DocumentChat:
    """Conversational document Q&A with memory."""

    def __init__(self):
        self.messages: list[dict] = []
        self.system = """You are a document Q&A assistant.
Answer based on the provided context. Cite sources with [Source N].
If you cannot find the answer, say so."""

    def ask(self, question: str, top_k: int = 5) -> str:
        """Ask a question with conversation history."""
        # Get relevant chunks
        query_embedding = get_embeddings([question])[0]
        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
        )

        chunks = results["documents"][0]
        metadatas = results["metadatas"][0]

        context = "\n\n---\n\n".join(
            f"[Source {i+1}: {m['source']}, Page {m['page']}]\n{c}"
            for i, (c, m) in enumerate(zip(chunks, metadatas))
        )

        # Add user message with context
        self.messages.append({
            "role": "user",
            "content": f"<context>\n{context}\n</context>\n\n{question}",
        })

        response = claude.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=2048,
            system=self.system,
            messages=self.messages,
        )

        answer = response.content[0].text
        self.messages.append({"role": "assistant", "content": answer})

        return answer

# Usage
chat = DocumentChat()
print(chat.ask("What is the vacation policy?"))
print(chat.ask("How many days for senior employees?"))  # Follow-up
print(chat.ask("Does this apply to contractors?"))  # Another follow-up

def ask_with_caching(question: str, document_context: str) -> str:
    """Ask with prompt caching for cheaper repeated queries."""
    response = claude.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        system=[
            {
                "type": "text",
                "text": "You are a document Q&A assistant. Answer only from context. Cite with [Source N].",
            },
            {
                "type": "text",
                "text": f"<documents>\n{document_context}\n</documents>",
                "cache_control": {"type": "ephemeral"},
            },
        ],
        messages=[{"role": "user", "content": question}],
    )

    return response.content[0].text