Add retrieval + chat loop (Session 3)

Backend:
- POST /ask takes {document_id, question}, runs hybrid retrieval over the
document's chunks, and returns an extractive answer plus the source
passage it came from (with chunk index and score).
- retrieval.py: a per-document index combining BM25 (rank_bm25) with a
FAISS cosine search over TF-IDF vectors; scores are min-max normalised
and averaged. The index is built lazily on first /ask and cached.
- Empty question -> 400, unknown document id -> 404.
- Dockerfile installs libgomp1 (faiss-cpu OpenMP runtime dep).

Frontend:
- Chat interface: question input, message history, a typing indicator
while /ask is in flight, and each answer shows a collapsible source
passage underneath. App swaps the uploader for the chat once a document
is ready, with a 'New document' reset.

Tests:
- test_ask.py covers index ranking, sentence extraction, the /ask happy
path, empty-question, and unknown-document paths.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WN4QRr6dTE2W7hQ2SDnLmY

PROJECT_SPEC.md

@@ -84,9 +84,9 @@ Goal: the core question→answer experience.
 
 **Acceptance criteria**
 
-- [ ] Asking a question about an uploaded doc returns a relevant passage.
-- [ ] The source passage is visibly shown under each answer.
-- [ ] Empty question or unknown document id is handled gracefully.
+- [x] Asking a question about an uploaded doc returns a relevant passage.
+- [x] The source passage is visibly shown under each answer.
+- [x] Empty question or unknown document id is handled gracefully.
 
 ## Session 4 — Data-collection & feedback layer (the score-lifting part)
 

README.md

@@ -6,10 +6,11 @@ A full-stack **document Q&A web app**: upload a PDF (or paste text), ask
 questions in a chat box, and get answers with the **source passage** shown —
 plus a live feedback/telemetry dashboard.
 
-> **Status:** Session 2 — upload & parse flow is live. Drag in a PDF (or paste
-> text) and the backend extracts, chunks, and indexes it in memory, returning a
-> document id. Retrieval/chat and the metrics dashboard land in the following
-> sessions (see [`PROJECT_SPEC.md`](./PROJECT_SPEC.md)).
+> **Status:** Session 3 — the chat loop is live. Upload a PDF (or paste text),
+> then ask questions in a chat box; each answer shows the source passage it came
+> from. Retrieval is a lightweight BM25 + FAISS (TF-IDF) hybrid. The
+> feedback/telemetry dashboard lands in the next sessions (see
+> [`PROJECT_SPEC.md`](./PROJECT_SPEC.md)).
 
 ## Run it (one command)
 

backend/Dockerfile

@@ -10,9 +10,10 @@ COPY . .
 
 EXPOSE 8000
 
-# curl is used by the docker-compose healthcheck.
+# curl is used by the docker-compose healthcheck; libgomp1 is required at
+# runtime by faiss-cpu (OpenMP) and is not present in the slim base image.
 RUN apt-get update \
-    && apt-get install -y --no-install-recommends curl \
+    && apt-get install -y --no-install-recommends curl libgomp1 \
     && rm -rf /var/lib/apt/lists/*
 
 CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

backend/app/main.py

@@ -11,9 +11,11 @@ import os
 
 from fastapi import FastAPI, File, Form, HTTPException, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 
 from . import __version__
 from .parsing import DocumentError, chunk_text, extract_pdf_text
+from .retrieval import DocumentIndex, best_sentence
 from .store import store
 
 logger = logging.getLogger("docuask")
@@ -134,3 +136,41 @@ def get_document(document_id: str) -> dict[str, object]:
         "num_chars": doc.num_chars,
         "status": "ready",
     }
+
+
+class AskRequest(BaseModel):
+    document_id: str
+    question: str
+
+
+@app.post("/ask")
+def ask(req: AskRequest) -> dict[str, object]:
+    """Answer a question about a stored document.
+
+    Runs hybrid retrieval over the document's chunks and returns the extractive
+    answer plus the source passage it came from, so the UI can show its work.
+    """
+    question = req.question.strip()
+    if not question:
+        raise HTTPException(status_code=400, detail="Question must not be empty.")
+
+    doc = store.get(req.document_id)
+    if doc is None:
+        raise HTTPException(status_code=404, detail="Document not found.")
+
+    # Build the index on first use and cache it on the document.
+    if doc.index is None:
+        doc.index = DocumentIndex(doc.chunks)
+
+    hits = doc.index.search(question, k=3)
+    top_index, score = hits[0]
+    passage = doc.chunks[top_index]
+
+    return {
+        "document_id": doc.id,
+        "question": question,
+        "answer": best_sentence(passage, question),
+        "source_passage": passage,
+        "chunk_index": top_index,
+        "score": round(score, 4),
+    }

backend/app/retrieval.py

@@ -0,0 +1,89 @@
+"""Hybrid retrieval for the chat loop (Session 3).
+
+Per the project spec the retrieval is intentionally lightweight: a single
+in-memory index per document combining lexical **BM25** with a **FAISS** cosine
+search over TF-IDF vectors. There is no generation model in the stack, so the
+"answer" is extractive — the sentence in the top passage that best overlaps the
+question.
+
+The two signals are min-max normalised across the document's chunks and
+averaged, which keeps a keyword hit (BM25) and a fuzzy/semantic-ish match
+(TF-IDF/FAISS) both able to surface the right passage.
+"""
+
+from __future__ import annotations
+
+import re
+
+import faiss
+import numpy as np
+from rank_bm25 import BM25Okapi
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+_TOKEN_RE = re.compile(r"[a-z0-9]+")
+_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")
+
+
+def _tokenize(text: str) -> list[str]:
+    return _TOKEN_RE.findall(text.lower())
+
+
+def _minmax(scores: np.ndarray) -> np.ndarray:
+    """Scale to [0, 1]; return zeros when every score is equal (no signal)."""
+    lo = float(scores.min())
+    hi = float(scores.max())
+    if hi - lo < 1e-9:
+        return np.zeros_like(scores)
+    return (scores - lo) / (hi - lo)
+
+
+def best_sentence(passage: str, question: str) -> str:
+    """Return the sentence in ``passage`` with the most question-term overlap."""
+    q_terms = set(_tokenize(question))
+    best = passage.strip()
+    best_score = -1
+    for sentence in _SENTENCE_RE.split(passage):
+        overlap = len(q_terms & set(_tokenize(sentence)))
+        if overlap > best_score:
+            best_score = overlap
+            best = sentence.strip()
+    return best
+
+
+class DocumentIndex:
+    """BM25 + TF-IDF/FAISS index over a single document's chunks."""
+
+    def __init__(self, chunks: list[str]) -> None:
+        if not chunks:
+            raise ValueError("cannot index an empty document")
+        self.chunks = chunks
+
+        self._bm25 = BM25Okapi([_tokenize(c) for c in chunks])
+
+        self._vectorizer = TfidfVectorizer()
+        dense = self._vectorizer.fit_transform(chunks).toarray().astype("float32")
+        faiss.normalize_L2(dense)
+        self._faiss = faiss.IndexFlatIP(dense.shape[1])
+        self._faiss.add(dense)
+
+    def _cosine_scores(self, query: str) -> np.ndarray:
+        """Cosine similarity of the query against every chunk (0 if no terms)."""
+        n = len(self.chunks)
+        scores = np.zeros(n, dtype="float32")
+        vec = self._vectorizer.transform([query]).toarray().astype("float32")
+        if float(np.linalg.norm(vec)) == 0.0:
+            return scores  # no overlapping vocabulary
+        faiss.normalize_L2(vec)
+        sims, idx = self._faiss.search(vec, n)
+        scores[idx[0]] = sims[0]
+        return scores
+
+    def search(self, query: str, k: int = 3) -> list[tuple[int, float]]:
+        """Return up to ``k`` ``(chunk_index, combined_score)`` pairs, best first."""
+        bm25 = np.asarray(self._bm25.get_scores(_tokenize(query)), dtype="float32")
+        cosine = self._cosine_scores(query)
+        combined = 0.5 * _minmax(bm25) + 0.5 * _minmax(cosine)
+
+        k = min(k, len(self.chunks))
+        top = np.argsort(-combined)[:k]
+        return [(int(i), float(combined[i])) for i in top]

backend/app/store.py

@@ -20,6 +20,9 @@ class Document:
     chunks: list[str]
     num_chars: int
     filename: str | None = None
+    # Retrieval index, built lazily on the first /ask and cached here. Kept out
+    # of repr/eq so the dataclass stays cheap to print and compare.
+    index: object | None = field(default=None, repr=False, compare=False)
 
     @property
     def num_chunks(self) -> int:

backend/requirements.txt

@@ -2,5 +2,9 @@ fastapi==0.115.6
 uvicorn[standard]==0.34.0
 python-multipart==0.0.20
 pypdf==5.1.0
+numpy==1.26.4
+scikit-learn==1.6.0
+rank-bm25==0.2.2
+faiss-cpu==1.9.0.post1
 httpx==0.28.1
 pytest==8.3.4

backend/tests/test_ask.py

@@ -0,0 +1,57 @@
+"""Tests for the retrieval + chat loop (Session 3)."""
+
+from fastapi.testclient import TestClient
+
+from app.main import app
+from app.retrieval import DocumentIndex, best_sentence
+
+client = TestClient(app)
+
+FACTS = [
+    "The mitochondria is the powerhouse of the cell.",
+    "Python is a popular programming language for data science.",
+    "The Eiffel Tower is located in Paris, France.",
+]
+
+
+def _ingest(text: str) -> str:
+    response = client.post("/documents", data={"text": text})
+    assert response.status_code == 201
+    return response.json()["document_id"]
+
+
+def test_index_ranks_relevant_chunk_first():
+    index = DocumentIndex(FACTS)
+    hits = index.search("Where is the Eiffel Tower?", k=1)
+    assert FACTS[hits[0][0]] == FACTS[2]
+
+
+def test_best_sentence_picks_overlapping_sentence():
+    passage = " ".join(FACTS)
+    assert "Eiffel" in best_sentence(passage, "Where is the Eiffel Tower?")
+
+
+def test_ask_returns_relevant_passage_and_answer():
+    doc_id = _ingest(" ".join(FACTS))
+    response = client.post(
+        "/ask",
+        json={"document_id": doc_id, "question": "Where is the Eiffel Tower?"},
+    )
+    assert response.status_code == 200
+    body = response.json()
+    assert "Paris" in body["source_passage"]
+    assert "Eiffel" in body["answer"] or "Paris" in body["answer"]
+    assert body["chunk_index"] >= 0
+
+
+def test_ask_empty_question_is_rejected():
+    doc_id = _ingest(" ".join(FACTS))
+    response = client.post("/ask", json={"document_id": doc_id, "question": "   "})
+    assert response.status_code == 400
+
+
+def test_ask_unknown_document_returns_404():
+    response = client.post(
+        "/ask", json={"document_id": "does-not-exist", "question": "Hello?"}
+    )
+    assert response.status_code == 404

frontend/src/App.jsx

@@ -1,9 +1,11 @@
 import { useEffect, useState } from "react";
 import DocumentUploader from "./DocumentUploader";
+import Chat from "./Chat";
 import { getHealth } from "./api";
 
 export default function App() {
   const [status, setStatus] = useState("loading");
+  const [doc, setDoc] = useState(null);
 
   useEffect(() => {
     let cancelled = false;
@@ -44,13 +46,19 @@ export default function App() {
         </div>
       </header>
 
-      <main className="mx-auto max-w-md px-6 py-10">
-        <h2 className="mb-1 text-lg font-semibold">Add a document</h2>
-        <p className="mb-5 text-sm text-slate-500">
-          Upload a PDF or paste text. We extract and index it so you can ask
-          questions about it.
-        </p>
-        <DocumentUploader onReady={(doc) => console.log("ready", doc)} />
+      <main className="mx-auto max-w-xl px-6 py-10">
+        {doc ? (
+          <Chat doc={doc} onReset={() => setDoc(null)} />
+        ) : (
+          <div className="mx-auto max-w-md">
+            <h2 className="mb-1 text-lg font-semibold">Add a document</h2>
+            <p className="mb-5 text-sm text-slate-500">
+              Upload a PDF or paste text. We extract and index it so you can ask
+              questions about it.
+            </p>
+            <DocumentUploader onReady={setDoc} />
+          </div>
+        )}
       </main>
     </div>
   );

frontend/src/Chat.jsx

@@ -0,0 +1,137 @@
+import { useEffect, useRef, useState } from "react";
+import { askQuestion } from "./api";
+
+// A single turn. Assistant turns carry the source passage so it can be shown
+// under the answer. `pending` marks the in-flight assistant bubble.
+function makeMessage(role, text, extra = {}) {
+  return { id: crypto.randomUUID(), role, text, ...extra };
+}
+
+export default function Chat({ doc, onReset }) {
+  const [messages, setMessages] = useState([]);
+  const [question, setQuestion] = useState("");
+  const [busy, setBusy] = useState(false);
+  const endRef = useRef(null);
+
+  useEffect(() => {
+    endRef.current?.scrollIntoView({ behavior: "smooth" });
+  }, [messages]);
+
+  async function send(e) {
+    e.preventDefault();
+    const q = question.trim();
+    if (!q || busy) return;
+
+    setMessages((m) => [...m, makeMessage("user", q)]);
+    setQuestion("");
+    setBusy(true);
+    try {
+      const res = await askQuestion(doc.document_id, q);
+      setMessages((m) => [
+        ...m,
+        makeMessage("assistant", res.answer, { source: res.source_passage }),
+      ]);
+    } catch (err) {
+      setMessages((m) => [...m, makeMessage("error", err.message)]);
+    } finally {
+      setBusy(false);
+    }
+  }
+
+  return (
+    <div className="flex h-[70vh] flex-col rounded-2xl bg-white shadow-lg ring-1 ring-slate-200">
+      <div className="flex items-center justify-between border-b border-slate-200 px-5 py-3">
+        <div className="min-w-0">
+          <p className="truncate text-sm font-semibold text-slate-800">
+            {doc.filename || "Pasted text"}
+          </p>
+          <p className="text-xs text-slate-400">{doc.num_chunks} chunks indexed</p>
+        </div>
+        <button
+          onClick={onReset}
+          className="shrink-0 rounded-lg px-3 py-1.5 text-xs font-medium text-slate-500 ring-1 ring-slate-200 transition hover:bg-slate-50"
+        >
+          New document
+        </button>
+      </div>
+
+      <div className="flex-1 space-y-4 overflow-y-auto px-5 py-5">
+        {messages.length === 0 && (
+          <p className="mt-10 text-center text-sm text-slate-400">
+            Ask a question about this document to get started.
+          </p>
+        )}
+
+        {messages.map((m) => {
+          if (m.role === "user") {
+            return (
+              <div key={m.id} className="flex justify-end">
+                <div className="max-w-[80%] rounded-2xl rounded-br-sm bg-slate-800 px-4 py-2 text-sm text-white">
+                  {m.text}
+                </div>
+              </div>
+            );
+          }
+          if (m.role === "error") {
+            return (
+              <div key={m.id} className="flex justify-start">
+                <div className="max-w-[80%] rounded-2xl bg-red-50 px-4 py-2 text-sm text-red-700 ring-1 ring-red-200">
+                  ⚠ {m.text}
+                </div>
+              </div>
+            );
+          }
+          return (
+            <div key={m.id} className="flex flex-col items-start gap-1">
+              <div className="max-w-[85%] rounded-2xl rounded-bl-sm bg-slate-100 px-4 py-2 text-sm text-slate-800">
+                {m.text}
+              </div>
+              {m.source && (
+                <details className="max-w-[85%] rounded-lg bg-emerald-50 px-3 py-2 text-xs text-slate-600 ring-1 ring-emerald-100">
+                  <summary className="cursor-pointer font-medium text-emerald-700">
+                    Source passage
+                  </summary>
+                  <p className="mt-1 whitespace-pre-wrap leading-relaxed">{m.source}</p>
+                </details>
+              )}
+            </div>
+          );
+        })}
+
+        {busy && (
+          <div className="flex justify-start">
+            <div className="flex gap-1 rounded-2xl rounded-bl-sm bg-slate-100 px-4 py-3">
+              <Dot /> <Dot delay="150ms" /> <Dot delay="300ms" />
+            </div>
+          </div>
+        )}
+        <div ref={endRef} />
+      </div>
+
+      <form onSubmit={send} className="flex gap-2 border-t border-slate-200 p-3">
+        <input
+          value={question}
+          onChange={(e) => setQuestion(e.target.value)}
+          placeholder="Ask a question…"
+          className="flex-1 rounded-lg border border-slate-300 bg-slate-50 px-3 py-2 text-sm outline-none transition focus:border-slate-400 focus:bg-white"
+        />
+        <button
+          type="submit"
+          disabled={busy || !question.trim()}
+          className="rounded-lg bg-emerald-600 px-4 py-2 text-sm font-medium text-white transition hover:bg-emerald-500 disabled:cursor-not-allowed disabled:bg-slate-300"
+        >
+          Ask
+        </button>
+      </form>
+    </div>
+  );
+}
+
+function Dot({ delay = "0ms" }) {
+  return (
+    <span
+      className="inline-block h-2 w-2 animate-bounce rounded-full bg-slate-400"
+      style={{ animationDelay: delay }}
+    />
+  );
+}

frontend/src/api.js

@@ -39,3 +39,14 @@ export async function uploadText(text) {
   if (!res.ok) throw new Error(await errorDetail(res));
   return res.json();
 }
+
+/** POST /ask — resolves to { answer, source_passage, score, ... }. */
+export async function askQuestion(documentId, question) {
+  const res = await fetch(`${API_BASE}/ask`, {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ document_id: documentId, question }),
+  });
+  if (!res.ok) throw new Error(await errorDetail(res));
+  return res.json();
+}