Capstone 1: Enterprise PDF Chatbot with RAG

+-----------------------------------------------------------------------------------+
|                           PROJECT 1 ARCHITECTURE                                  |
+-----------------------------------------------------------------------------------+
|  PDF Upload -> PyMuPDF Text & Table Extraction -> Recursive Semantic Chunking     |
|                                                          |                        |
|                                                          v                        |
|  Query Engine <- Cosine Reranking <- Hybrid Search <- ChromaDB Vector Store       |
|       |                                                                           |
|       v                                                                           |
|  Streaming Answer with Exact Citations & Page References (FastAPI + Streamlit)    |
+-----------------------------------------------------------------------------------+

Objective

Build a production document Q&A assistant capable of ingesting multi-page complex enterprise PDFs (contracts, architecture specs, financial reports) and providing answers with verified page citations.

Complete Implementation

import os
import fitz # PyMuPDF
import chromadb
from chromadb.utils import embedding_functions
from openai import OpenAI

client = OpenAI()

# 1. Extract and Chunk PDF with Page Metadata
def extract_pdf_chunks(pdf_path: str, chunk_size: int = 500, overlap: int = 50) -> list[dict]:
    doc = fitz.open(pdf_path)
    chunks = []

    for page_num in range(len(doc)):
        page = doc[page_num]
        text = page.get_text("text").strip()
        if not text:
            continue

        words = text.split()
        for i in range(0, len(words), chunk_size - overlap):
            chunk_words = words[i:i + chunk_size]
            chunk_text = " ".join(chunk_words)
            chunks.append({
                "text": chunk_text,
                "metadata": {
                    "source": os.path.basename(pdf_path),
                    "page": page_num + 1,
                    "chunk_id": f"{page_num+1}_{i}"
                }
            })
    return chunks

# 2. Ingest into ChromaDB
def ingest_chunks_to_chroma(chunks: list[dict], collection_name: str = "enterprise_docs"):
    chroma_client = chromadb.Client()
    openai_ef = embedding_functions.OpenAIEmbeddingFunction(
        api_key=os.environ.get("OPENAI_API_KEY"),
        model_name="text-embedding-3-small"
    )
    collection = chroma_client.get_or_create_collection(
        name=collection_name,
        embedding_function=openai_ef
    )

    collection.add(
        documents=[c["text"] for c in chunks],
        metadatas=[c["metadata"] for c in chunks],
        ids=[f"doc_{idx}" for idx in range(len(chunks))]
    )
    return collection

# 3. RAG Generation with Citations
def query_rag_engine(collection, question: str) -> str:
    results = collection.query(query_texts=[question], n_results=3)
    retrieved_texts = results["documents"][0]
    retrieved_meta = results["metadatas"][0]

    context_str = ""
    for text, meta in zip(retrieved_texts, retrieved_meta):
        context_str += f"
[Document: {meta['source']} | Page: {meta['page']}]
{text}
"

    prompt = f"""
Answer the user's question using ONLY the provided context snippets.
Always cite the exact Page numbers supporting your answer.

CONTEXT:
{context_str}

QUESTION:
{question}
"""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1
    )
    return response.choices[0].message.content