Case Study: Document Q&A Bot

Let's build a complete Q&A system that lets users upload documents and ask questions. This combines everything we've learned: LLM APIs, prompt engineering, embeddings, and RAG.

---

Architecture

---

Project Structure

doc-qa-bot/
├── app/
│   ├── __init__.py
│   ├── main.py           # FastAPI app
│   ├── models.py         # Pydantic models
│   ├── database.py       # Postgres connection
│   ├── embeddings.py     # OpenAI embedding client
│   ├── chunker.py        # Document chunking
│   └── rag.py            # RAG pipeline
├── requirements.txt
└── .env

---

Step 1: Database Setup

# app/database.py
import psycopg2
from psycopg2.extras import RealDictCursor
import os

def get_connection():
    return psycopg2.connect(
        os.getenv("DATABASE_URL"),
        cursor_factory=RealDictCursor
    )

def init_db():
    conn = get_connection()
    with conn.cursor() as cur:
        # Enable pgvector
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector")

        # Documents table
        cur.execute("""
            CREATE TABLE IF NOT EXISTS documents (
                id SERIAL PRIMARY KEY,
                filename TEXT NOT NULL,
                status TEXT DEFAULT 'processing',
                created_at TIMESTAMP DEFAULT NOW()
            )
        """)

        # Chunks table with embeddings
        cur.execute("""
            CREATE TABLE IF NOT EXISTS chunks (
                id SERIAL PRIMARY KEY,
                document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE,
                content TEXT NOT NULL,
                embedding vector(1536),
                chunk_index INTEGER
            )
        """)

        # Index for fast similarity search
        cur.execute("""
            CREATE INDEX IF NOT EXISTS chunks_embedding_idx
            ON chunks USING ivfflat (embedding vector_cosine_ops)
            WITH (lists = 100)
        """)

    conn.commit()
    conn.close()

---

Step 2: Document Chunker

# app/chunker.py
import re

def chunk_document(text: str, chunk_size: int = 400, overlap: int = 50) -> list[str]:
    """
    Split document into chunks, trying to respect paragraph boundaries.
    """
    # Clean the text
    text = re.sub(r'\n{3,}', '\n\n', text)  # Normalize multiple newlines
    text = re.sub(r' {2,}', ' ', text)       # Normalize spaces

    # Split by paragraphs first
    paragraphs = text.split('\n\n')

    chunks = []
    current_chunk = []
    current_length = 0

    for para in paragraphs:
        para_words = para.split()
        para_length = len(para_words)

        if current_length + para_length <= chunk_size:
            current_chunk.append(para)
            current_length += para_length
        else:
            # Save current chunk
            if current_chunk:
                chunks.append('\n\n'.join(current_chunk))

            # Start new chunk (with overlap from previous)
            if chunks and overlap > 0:
                # Get last N words from previous chunk
                prev_words = chunks[-1].split()[-overlap:]
                current_chunk = [' '.join(prev_words), para]
                current_length = overlap + para_length
            else:
                current_chunk = [para]
                current_length = para_length

    # Don't forget the last chunk
    if current_chunk:
        chunks.append('\n\n'.join(current_chunk))

    return chunks

---

Step 3: Embeddings Client

# app/embeddings.py
from openai import OpenAI
import os

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def get_embedding(text: str) -> list[float]:
    """Generate embedding for a single text."""
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return response.data[0].embedding

def get_embeddings_batch(texts: list[str]) -> list[list[float]]:
    """Generate embeddings for multiple texts (more efficient)."""
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    return [item.embedding for item in response.data]

---

Step 4: RAG Pipeline

# app/rag.py
from openai import OpenAI
from app.database import get_connection
from app.embeddings import get_embedding
import os

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def retrieve_context(query: str, top_k: int = 5, threshold: float = 0.7) -> list[dict]:
    """Find relevant chunks for a query."""
    query_embedding = get_embedding(query)
    conn = get_connection()

    with conn.cursor() as cur:
        cur.execute("""
            SELECT
                c.content,
                d.filename,
                1 - (c.embedding <=> %s::vector) AS similarity
            FROM chunks c
            JOIN documents d ON c.document_id = d.id
            WHERE d.status = 'ready'
              AND 1 - (c.embedding <=> %s::vector) > %s
            ORDER BY c.embedding <=> %s::vector
            LIMIT %s
        """, (query_embedding, query_embedding, threshold, query_embedding, top_k))

        results = [
            {"content": row["content"], "source": row["filename"], "similarity": row["similarity"]}
            for row in cur.fetchall()
        ]

    conn.close()
    return results

def generate_answer(query: str, context_chunks: list[dict], chat_history: list = None) -> str:
    """Generate answer using retrieved context."""

    if not context_chunks:
        return "I couldn't find relevant information in the uploaded documents to answer your question."

    # Build context string
    context = "\n\n---\n\n".join([
        f"[Source: {c['source']}]\n{c['content']}"
        for c in context_chunks
    ])

    # Build messages
    messages = [
        {
            "role": "system",
            "content": """You are a helpful assistant that answers questions based on the provided documents.

RULES:
- Only answer based on the provided context
- If the context doesn't contain the answer, say so
- Be concise but complete
- Mention which document(s) the information came from
- If asked a follow-up question, use the conversation history for context"""
        }
    ]

    # Add chat history for follow-ups
    if chat_history:
        messages.extend(chat_history[-4:])  # Last 2 exchanges

    # Add current query with context
    messages.append({
        "role": "user",
        "content": f"""Based on these documents:

{context}

---

Question: {query}"""
    })

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        temperature=0.3,  # Lower = more focused
        max_tokens=500
    )

    return response.choices[0].message.content

def ask(query: str, chat_history: list = None) -> dict:
    """Main RAG function: retrieve and generate."""
    # Retrieve
    chunks = retrieve_context(query)

    # Generate
    answer = generate_answer(query, chunks, chat_history)

    # Return with metadata
    return {
        "answer": answer,
        "sources": list(set([c["source"] for c in chunks])),
        "chunks_used": len(chunks)
    }

---

Step 5: FastAPI Application

# app/main.py
from fastapi import FastAPI, UploadFile, HTTPException, BackgroundTasks
from pydantic import BaseModel
from app.database import get_connection, init_db
from app.chunker import chunk_document
from app.embeddings import get_embeddings_batch
from app.rag import ask
import io

app = FastAPI(title="Document Q&A Bot")

# Initialize database on startup
@app.on_event("startup")
def startup():
    init_db()

# Request/Response models
class QuestionRequest(BaseModel):
    question: str
    chat_history: list = None

class AnswerResponse(BaseModel):
    answer: str
    sources: list[str]

# Background task for document processing
def process_document(doc_id: int, content: str, filename: str):
    conn = get_connection()

    try:
        # Chunk the document
        chunks = chunk_document(content)

        # Generate embeddings in batch
        embeddings = get_embeddings_batch(chunks)

        # Store chunks with embeddings
        with conn.cursor() as cur:
            for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
                cur.execute(
                    """INSERT INTO chunks (document_id, content, embedding, chunk_index)
                       VALUES (%s, %s, %s, %s)""",
                    (doc_id, chunk, embedding, i)
                )

            # Mark document as ready
            cur.execute(
                "UPDATE documents SET status = 'ready' WHERE id = %s",
                (doc_id,)
            )

        conn.commit()
        print(f"Processed {filename}: {len(chunks)} chunks")

    except Exception as e:
        # Mark as failed
        with conn.cursor() as cur:
            cur.execute(
                "UPDATE documents SET status = 'failed' WHERE id = %s",
                (doc_id,)
            )
        conn.commit()
        print(f"Failed to process {filename}: {e}")

    finally:
        conn.close()

@app.post("/upload")
async def upload_document(file: UploadFile, background_tasks: BackgroundTasks):
    """Upload a document for processing."""

    # Validate file type
    allowed_types = [".txt", ".md", ".pdf"]
    if not any(file.filename.endswith(t) for t in allowed_types):
        raise HTTPException(400, f"File type not supported. Use: {allowed_types}")

    # Read content
    content = await file.read()

    # Handle different file types
    if file.filename.endswith(".pdf"):
        # You'd use PyPDF2 or pdfplumber here
        # For simplicity, assuming text extraction is done
        text = content.decode("utf-8", errors="ignore")
    else:
        text = content.decode("utf-8")

    # Create document record
    conn = get_connection()
    with conn.cursor() as cur:
        cur.execute(
            "INSERT INTO documents (filename) VALUES (%s) RETURNING id",
            (file.filename,)
        )
        doc_id = cur.fetchone()["id"]
    conn.commit()
    conn.close()

    # Process in background
    background_tasks.add_task(process_document, doc_id, text, file.filename)

    return {"message": "Document uploaded", "document_id": doc_id, "status": "processing"}

@app.get("/documents")
async def list_documents():
    """List all uploaded documents and their status."""
    conn = get_connection()
    with conn.cursor() as cur:
        cur.execute("SELECT id, filename, status, created_at FROM documents ORDER BY created_at DESC")
        docs = cur.fetchall()
    conn.close()
    return docs

@app.post("/ask", response_model=AnswerResponse)
async def ask_question(request: QuestionRequest):
    """Ask a question about the uploaded documents."""
    result = ask(request.question, request.chat_history)
    return AnswerResponse(
        answer=result["answer"],
        sources=result["sources"]
    )

---

Running the App

# Install dependencies
pip install fastapi uvicorn psycopg2-binary openai python-dotenv

# Set environment variables
export DATABASE_URL="postgresql://localhost/docqa"
export OPENAI_API_KEY="sk-..."

# Run the server
uvicorn app.main:app --reload

---

Testing the API

# Upload a document
curl -X POST "http://localhost:8000/upload" \
  -F "file=@docs/handbook.md"

# Check document status
curl "http://localhost:8000/documents"

# Ask a question
curl -X POST "http://localhost:8000/ask" \
  -H "Content-Type: application/json" \
  -d '{"question": "What is the refund policy?"}'

# Follow-up question with history
curl -X POST "http://localhost:8000/ask" \
  -H "Content-Type: application/json" \
  -d '{
    "question": "What about for monthly plans?",
    "chat_history": [
      {"role": "user", "content": "What is the refund policy?"},
      {"role": "assistant", "content": "Annual plans have a 14-day refund policy..."}
    ]
  }'

---

Production Considerations

---

What You've Built

Congratulations! You now have a working Document Q&A system that:

---

Course Complete!

You've learned the essential skills for integrating AI into backend applications:

• Add PDF parsing with pdfplumber
• Implement streaming responses for better UX
• Add a frontend with React or Next.js
• Explore AI agents for more complex workflows

Happy building!