RAG techniques and benchmark

README.md

@@ -19,6 +19,7 @@ MemoriaLM is a full-stack RAG application inspired by NotebookLM. It allows auth
 - Upload documents (`.pdf`, `.pptx`, `.txt`) and URLs
 - Chat with uploaded content using RAG
 - Generate study artifacts (reports, quizzes, podcasts)
+- Retrieval modes for comparison: `topk` and `rerank`
 
 ### Architecture
 - **Frontend**: Gradio UI (frontend/app.py)
@@ -78,3 +79,9 @@ Run all tests:
 ```sh
 pytest tests/
 ```
+
+## RAG Technique Comparison
+See `docs/rag_techniques.md` for:
+- implemented retrieval modes (`topk` vs `rerank`)
+- benchmark command (`scripts/rag_benchmark.py`)
+- results table template for the project deliverable

backend/api/artifacts.py

@@ -1,39 +1,50 @@
-from fastapi import APIRouter, HTTPException, Query
-from fastapi.responses import FileResponse
-
-from backend.models.schemas import ArtifactGenerateOut, ArtifactGenerateRequest, ArtifactListOut
-from backend.modules.artifacts import (
+from fastapi import APIRouter, Depends, HTTPException, Query
+from fastapi.responses import FileResponse
+
+from backend.models.schemas import ArtifactGenerateOut, ArtifactGenerateRequest, ArtifactListOut
+from backend.modules.artifacts import (
     generate_podcast,
     generate_quiz,
     generate_report,
-    list_artifacts,
-    resolve_artifact_path,
-)
-from backend.services.storage import NotebookStore
+    list_artifacts,
+    resolve_artifact_path,
+)
+from backend.services.auth import User, enforce_user_match, get_current_user
+from backend.services.storage import NotebookStore
 
 router = APIRouter()
 store = NotebookStore(base_dir="data")
 
 
-@router.get("/{notebook_id}/artifacts", response_model=ArtifactListOut)
-def list_notebook_artifacts(notebook_id: str, user_id: str = Query(...)) -> ArtifactListOut:
-    try:
-        return list_artifacts(store, user_id=user_id, notebook_id=notebook_id)
-    except FileNotFoundError:
-        raise HTTPException(status_code=404, detail="Notebook not found")
+@router.get("/{notebook_id}/artifacts", response_model=ArtifactListOut)
+def list_notebook_artifacts(
+    notebook_id: str,
+    user_id: str | None = Query(default=None),
+    current_user: User = Depends(get_current_user),
+) -> ArtifactListOut:
+    try:
+        enforce_user_match(current_user, user_id)
+        return list_artifacts(store, user_id=current_user.user_id, notebook_id=notebook_id)
+    except FileNotFoundError:
+        raise HTTPException(status_code=404, detail="Notebook not found")
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
 
 
-@router.post("/{notebook_id}/artifacts/report", response_model=ArtifactGenerateOut)
-def create_report_artifact(notebook_id: str, payload: ArtifactGenerateRequest) -> ArtifactGenerateOut:
-    try:
-        return generate_report(
-            store,
-            user_id=payload.user_id,
-            notebook_id=notebook_id,
-            prompt=payload.prompt,
-        )
+@router.post("/{notebook_id}/artifacts/report", response_model=ArtifactGenerateOut)
+def create_report_artifact(
+    notebook_id: str,
+    payload: ArtifactGenerateRequest,
+    current_user: User = Depends(get_current_user),
+) -> ArtifactGenerateOut:
+    try:
+        enforce_user_match(current_user, payload.user_id)
+        return generate_report(
+            store,
+            user_id=current_user.user_id,
+            notebook_id=notebook_id,
+            prompt=payload.prompt,
+        )
     except FileNotFoundError:
         raise HTTPException(status_code=404, detail="Notebook not found")
     except ValueError as exc:
@@ -42,15 +53,20 @@ def create_report_artifact(notebook_id: str, payload: ArtifactGenerateRequest) -
         raise HTTPException(status_code=500, detail=f"Report generation failed: {exc}") from exc
 
 
-@router.post("/{notebook_id}/artifacts/quiz", response_model=ArtifactGenerateOut)
-def create_quiz_artifact(notebook_id: str, payload: ArtifactGenerateRequest) -> ArtifactGenerateOut:
-    try:
-        return generate_quiz(
-            store,
-            user_id=payload.user_id,
-            notebook_id=notebook_id,
-            prompt=payload.prompt,
-            num_questions=payload.num_questions,
+@router.post("/{notebook_id}/artifacts/quiz", response_model=ArtifactGenerateOut)
+def create_quiz_artifact(
+    notebook_id: str,
+    payload: ArtifactGenerateRequest,
+    current_user: User = Depends(get_current_user),
+) -> ArtifactGenerateOut:
+    try:
+        enforce_user_match(current_user, payload.user_id)
+        return generate_quiz(
+            store,
+            user_id=current_user.user_id,
+            notebook_id=notebook_id,
+            prompt=payload.prompt,
+            num_questions=payload.num_questions,
         )
     except FileNotFoundError:
         raise HTTPException(status_code=404, detail="Notebook not found")
@@ -60,15 +76,20 @@ def create_quiz_artifact(notebook_id: str, payload: ArtifactGenerateRequest) ->
         raise HTTPException(status_code=500, detail=f"Quiz generation failed: {exc}") from exc
 
 
-@router.post("/{notebook_id}/artifacts/podcast", response_model=ArtifactGenerateOut)
-def create_podcast_artifact(notebook_id: str, payload: ArtifactGenerateRequest) -> ArtifactGenerateOut:
-    try:
-        return generate_podcast(
-            store,
-            user_id=payload.user_id,
-            notebook_id=notebook_id,
-            prompt=payload.prompt,
-        )
+@router.post("/{notebook_id}/artifacts/podcast", response_model=ArtifactGenerateOut)
+def create_podcast_artifact(
+    notebook_id: str,
+    payload: ArtifactGenerateRequest,
+    current_user: User = Depends(get_current_user),
+) -> ArtifactGenerateOut:
+    try:
+        enforce_user_match(current_user, payload.user_id)
+        return generate_podcast(
+            store,
+            user_id=current_user.user_id,
+            notebook_id=notebook_id,
+            prompt=payload.prompt,
+        )
     except FileNotFoundError:
         raise HTTPException(status_code=404, detail="Notebook not found")
     except ValueError as exc:
@@ -79,18 +100,20 @@ def create_podcast_artifact(notebook_id: str, payload: ArtifactGenerateRequest)
 
 @router.get("/{notebook_id}/artifacts/download")
 def download_artifact(
-    notebook_id: str,
-    user_id: str = Query(...),
-    artifact_type: str = Query(...),
-    filename: str = Query(...),
-):
-    try:
-        path = resolve_artifact_path(
-            store,
-            user_id=user_id,
-            notebook_id=notebook_id,
-            artifact_type=artifact_type,
-            filename=filename,
+    notebook_id: str,
+    user_id: str | None = Query(default=None),
+    artifact_type: str = Query(...),
+    filename: str = Query(...),
+    current_user: User = Depends(get_current_user),
+):
+    try:
+        enforce_user_match(current_user, user_id)
+        path = resolve_artifact_path(
+            store,
+            user_id=current_user.user_id,
+            notebook_id=notebook_id,
+            artifact_type=artifact_type,
+            filename=filename,
         )
         media_type = "audio/mpeg" if path.suffix.lower() == ".mp3" else "text/markdown"
         return FileResponse(path=path, filename=path.name, media_type=media_type)

backend/api/chat.py

@@ -1,7 +1,8 @@
-from fastapi import APIRouter, HTTPException, Query
+from fastapi import APIRouter, Depends, HTTPException
 
 from backend.models.schemas import ChatHistoryOut, ChatMessageOut, ChatRequest, ChatResponseOut
 from backend.modules.rag import answer_notebook_question, get_chat_history
+from backend.services.auth import User, enforce_user_match, get_current_user
 from backend.services.storage import NotebookStore
 
 router = APIRouter()
@@ -9,9 +10,9 @@ store = NotebookStore(base_dir="data")
 
 
 @router.get("/{notebook_id}/chat", response_model=ChatHistoryOut)
-def chat_history(notebook_id: str, user_id: str = Query(...)) -> ChatHistoryOut:
+def chat_history(notebook_id: str, current_user: User = Depends(get_current_user)) -> ChatHistoryOut:
     try:
-        messages = get_chat_history(store, user_id=user_id, notebook_id=notebook_id)
+        messages = get_chat_history(store, user_id=current_user.user_id, notebook_id=notebook_id)
         return ChatHistoryOut(messages=[ChatMessageOut(**m) for m in messages])
     except FileNotFoundError:
         raise HTTPException(status_code=404, detail="Notebook not found")
@@ -20,14 +21,20 @@ def chat_history(notebook_id: str, user_id: str = Query(...)) -> ChatHistoryOut:
 
 
 @router.post("/{notebook_id}/chat", response_model=ChatResponseOut)
-def chat(notebook_id: str, payload: ChatRequest) -> ChatResponseOut:
+def chat(
+    notebook_id: str,
+    payload: ChatRequest,
+    current_user: User = Depends(get_current_user),
+) -> ChatResponseOut:
     try:
+        enforce_user_match(current_user, payload.user_id)
         result = answer_notebook_question(
             store,
-            user_id=payload.user_id,
+            user_id=current_user.user_id,
             notebook_id=notebook_id,
             message=payload.message,
             top_k=payload.top_k,
+            retrieval_mode=payload.retrieval_mode,
         )
         return ChatResponseOut(**result)
     except FileNotFoundError:

backend/api/notebooks.py

@@ -1,8 +1,9 @@
 from typing import List
 
-from fastapi import APIRouter, HTTPException, Query
+from fastapi import APIRouter, Depends, HTTPException
 
 from backend.models.schemas import NotebookCreate, NotebookOut, NotebookRename
+from backend.services.auth import User, enforce_user_match, get_current_user
 from backend.services.storage import NotebookStore
 
 router = APIRouter()
@@ -10,25 +11,29 @@ store = NotebookStore(base_dir="data")
 
 
 @router.get("/", response_model=List[NotebookOut])
-def list_notebooks(user_id: str = Query(...)) -> List[NotebookOut]:
+def list_notebooks(current_user: User = Depends(get_current_user)) -> List[NotebookOut]:
     try:
-        return store.list(user_id)
+        return store.list(current_user.user_id)
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
 
 
 @router.post("/", response_model=NotebookOut)
-def create_notebook(payload: NotebookCreate) -> NotebookOut:
+def create_notebook(
+    payload: NotebookCreate,
+    current_user: User = Depends(get_current_user),
+) -> NotebookOut:
     try:
-        return store.create(payload)
+        enforce_user_match(current_user, payload.user_id)
+        return store.create(NotebookCreate(user_id=current_user.user_id, name=payload.name))
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
 
 
 @router.get("/{notebook_id}", response_model=NotebookOut)
-def get_notebook(notebook_id: str, user_id: str = Query(...)) -> NotebookOut:
+def get_notebook(notebook_id: str, current_user: User = Depends(get_current_user)) -> NotebookOut:
     try:
-        notebook = store.get(user_id, notebook_id)
+        notebook = store.get(current_user.user_id, notebook_id)
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
 
@@ -38,9 +43,14 @@ def get_notebook(notebook_id: str, user_id: str = Query(...)) -> NotebookOut:
 
 
 @router.patch("/{notebook_id}", response_model=NotebookOut)
-def rename_notebook(notebook_id: str, payload: NotebookRename) -> NotebookOut:
+def rename_notebook(
+    notebook_id: str,
+    payload: NotebookRename,
+    current_user: User = Depends(get_current_user),
+) -> NotebookOut:
     try:
-        notebook = store.rename(payload.user_id, notebook_id, payload.name)
+        enforce_user_match(current_user, payload.user_id)
+        notebook = store.rename(current_user.user_id, notebook_id, payload.name)
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
 
@@ -50,9 +60,9 @@ def rename_notebook(notebook_id: str, payload: NotebookRename) -> NotebookOut:
 
 
 @router.delete("/{notebook_id}")
-def delete_notebook(notebook_id: str, user_id: str = Query(...)) -> dict:
+def delete_notebook(notebook_id: str, current_user: User = Depends(get_current_user)) -> dict:
     try:
-        deleted = store.delete(user_id, notebook_id)
+        deleted = store.delete(current_user.user_id, notebook_id)
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
 

backend/api/sources.py

@@ -1,7 +1,10 @@
-from fastapi import APIRouter, File, Form, HTTPException, Query, UploadFile
+from typing import Optional
+
+from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
 
 from backend.models.schemas import SourceListOut, SourceOut, UrlIngestRequest
 from backend.modules.ingestion import ingest_uploaded_bytes, ingest_url, list_ingested_sources
+from backend.services.auth import User, enforce_user_match, get_current_user
 from backend.services.storage import NotebookStore
 
 router = APIRouter()
@@ -9,9 +12,9 @@ store = NotebookStore(base_dir="data")
 
 
 @router.get("/{notebook_id}/sources", response_model=SourceListOut)
-def list_sources(notebook_id: str, user_id: str = Query(...)) -> SourceListOut:
+def list_sources(notebook_id: str, current_user: User = Depends(get_current_user)) -> SourceListOut:
     try:
-        items = list_ingested_sources(store, user_id=user_id, notebook_id=notebook_id)
+        items = list_ingested_sources(store, user_id=current_user.user_id, notebook_id=notebook_id)
         return SourceListOut(sources=items)
     except FileNotFoundError:
         raise HTTPException(status_code=404, detail="Notebook not found")
@@ -20,11 +23,16 @@ def list_sources(notebook_id: str, user_id: str = Query(...)) -> SourceListOut:
 
 
 @router.post("/{notebook_id}/sources/url", response_model=SourceOut)
-def ingest_source_url(notebook_id: str, payload: UrlIngestRequest) -> SourceOut:
+def ingest_source_url(
+    notebook_id: str,
+    payload: UrlIngestRequest,
+    current_user: User = Depends(get_current_user),
+) -> SourceOut:
     try:
+        enforce_user_match(current_user, payload.user_id)
         return ingest_url(
             store,
-            user_id=payload.user_id,
+            user_id=current_user.user_id,
             notebook_id=notebook_id,
             url=str(payload.url),
             source_name=payload.source_name,
@@ -40,14 +48,16 @@ def ingest_source_url(notebook_id: str, payload: UrlIngestRequest) -> SourceOut:
 @router.post("/{notebook_id}/sources/upload", response_model=SourceOut)
 async def upload_source_file(
     notebook_id: str,
-    user_id: str = Form(...),
+    user_id: Optional[str] = Form(default=None),
     file: UploadFile = File(...),
+    current_user: User = Depends(get_current_user),
 ) -> SourceOut:
     try:
+        enforce_user_match(current_user, user_id)
         content = await file.read()
         return ingest_uploaded_bytes(
             store,
-            user_id=user_id,
+            user_id=current_user.user_id,
             notebook_id=notebook_id,
             filename=file.filename or "upload.txt",
             content=content,

backend/models/schemas.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Literal, Optional
 
 from pydantic import BaseModel, HttpUrl
 
@@ -46,6 +46,7 @@ class ChatRequest(BaseModel):
     user_id: str
     message: str
     top_k: int = 5
+    retrieval_mode: Literal["topk", "rerank"] = "topk"
 
 
 class CitationOut(BaseModel):

backend/modules/artifacts.py

@@ -1,17 +1,12 @@
-import audioop
-import json
-import math
-import os
-import re
-import wave
-import torch
-import io
-import soundfile as sf
-from transformers import AutoTokenizer, VitsModel
-from datetime import datetime, timezone
-from io import BytesIO
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+import audioop
+import json
+import re
+import wave
+import io
+from datetime import datetime, timezone
+from io import BytesIO
+from pathlib import Path
+from typing import Dict, List, Optional
 
 from backend.models.schemas import (
     ArtifactFileOut,
@@ -27,8 +22,35 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-_vits_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
-_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
+try:
+    import soundfile as sf
+except ImportError:
+    sf = None
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+try:
+    from transformers import AutoTokenizer, VitsModel
+except ImportError:
+    AutoTokenizer = None
+    VitsModel = None
+
+_vits_model = None
+_tokenizer = None
+
+
+def _get_tts_model():
+    global _vits_model, _tokenizer
+    if _vits_model is not None and _tokenizer is not None:
+        return _vits_model, _tokenizer
+    if torch is None or sf is None or AutoTokenizer is None or VitsModel is None:
+        raise RuntimeError("TTS dependencies are not installed")
+    _vits_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+    _tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
+    return _vits_model, _tokenizer
 
 def _now() -> str:
     return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
@@ -209,13 +231,14 @@ def clean_transcript_for_tts(transcript: str) -> str:
     return " ".join(lines)
 
 
-def _synthesize_podcast_mp3(transcript_text: str) -> bytes:
-    tts_text = clean_transcript_for_tts(transcript_text)[:1800]
-
-    inputs = _tokenizer(tts_text, return_tensors="pt")
-
-    with torch.no_grad():
-        waveform = _vits_model(**inputs).waveform.squeeze().cpu().numpy()
+def _synthesize_podcast_mp3(transcript_text: str) -> bytes:
+    tts_text = clean_transcript_for_tts(transcript_text)[:1800]
+    model, tokenizer = _get_tts_model()
+
+    inputs = tokenizer(tts_text, return_tensors="pt")
+
+    with torch.no_grad():
+        waveform = model(**inputs).waveform.squeeze().cpu().numpy()
 
     wav_buffer = io.BytesIO()
     sf.write(wav_buffer, waveform, 16000, format="WAV")

backend/modules/rag.py

@@ -1,4 +1,5 @@
 from datetime import datetime, timezone
+import re
 from typing import Any, Dict, List
 
 from backend.services.embeddings import embedding_service
@@ -51,12 +52,60 @@ def retrieve_notebook_chunks(
     notebook_id: str,
     query: str,
     top_k: int = 5,
+    retrieval_mode: str = "topk",
 ) -> List[Dict[str, Any]]:
     query_vecs = embedding_service.embed_texts([query])
     if not query_vecs:
         return []
     chroma = ChromaNotebookStore(store.chroma_dir(user_id, notebook_id))
-    return chroma.query(query_vecs[0], k=top_k)
+    mode = retrieval_mode.strip().lower()
+    if mode == "topk":
+        return chroma.query(query_vecs[0], k=top_k)
+    if mode == "rerank":
+        # Retrieve a wider pool first, then rerank using lexical overlap + vector signal.
+        pool_size = max(top_k * 4, top_k)
+        pool = chroma.query(query_vecs[0], k=pool_size)
+        return rerank_chunks(query=query, candidate_chunks=pool, top_k=top_k)
+    raise ValueError("Unsupported retrieval_mode")
+
+
+def _tokenize(text: str) -> set[str]:
+    return {tok for tok in re.findall(r"[a-zA-Z0-9]+", text.lower()) if tok}
+
+
+def _vector_relevance_from_distance(distance: Any) -> float:
+    if distance is None:
+        return 0.0
+    try:
+        dist = float(distance)
+    except (TypeError, ValueError):
+        return 0.0
+    # Chroma cosine distance: lower is better; map to 0..1 relevance-like value.
+    return 1.0 / (1.0 + max(0.0, dist))
+
+
+def rerank_chunks(
+    *,
+    query: str,
+    candidate_chunks: List[Dict[str, Any]],
+    top_k: int,
+    alpha: float = 0.65,
+) -> List[Dict[str, Any]]:
+    if not candidate_chunks:
+        return []
+    query_tokens = _tokenize(query)
+    ranked: List[tuple[float, Dict[str, Any]]] = []
+    for chunk in candidate_chunks:
+        text = str(chunk.get("text", ""))
+        chunk_tokens = _tokenize(text)
+        lexical = (len(query_tokens & chunk_tokens) / len(query_tokens)) if query_tokens else 0.0
+        vector_rel = _vector_relevance_from_distance(chunk.get("distance"))
+        score = (alpha * vector_rel) + ((1.0 - alpha) * lexical)
+        enriched = dict(chunk)
+        enriched["rerank_score"] = score
+        ranked.append((score, enriched))
+    ranked.sort(key=lambda x: x[0], reverse=True)
+    return [item[1] for item in ranked[: max(1, top_k)]]
 
 
 def _citation_objects(retrieved_chunks: List[Dict[str, Any]]) -> List[Dict[str, str]]:
@@ -87,6 +136,7 @@ def answer_notebook_question(
     notebook_id: str,
     message: str,
     top_k: int = 5,
+    retrieval_mode: str = "topk",
 ) -> Dict[str, Any]:
     # Ensures the notebook exists and belongs to the user before retrieving.
     store.require_notebook_dir(user_id, notebook_id)
@@ -97,6 +147,7 @@ def answer_notebook_question(
         notebook_id=notebook_id,
         query=message,
         top_k=top_k,
+        retrieval_mode=retrieval_mode,
     )
     prompt = build_rag_prompt(message, retrieved_chunks)
     answer = llm_service.generate(prompt)

backend/services/auth.py

@@ -1,4 +1,8 @@
 from dataclasses import dataclass
+import re
+from typing import Optional
+
+from fastapi import Header, HTTPException, Request
 
 
 @dataclass
@@ -7,6 +11,50 @@ class User:
     email: str
 
 
-def get_current_user() -> User:
-    # Placeholder for HF OAuth integration
-    return User(user_id="demo-user", email="demo@example.com")
+_USER_ID_PATTERN = re.compile(r"[A-Za-z0-9._@-]+$")
+
+
+def _validate_user_id(value: str) -> str:
+    if not _USER_ID_PATTERN.fullmatch(value):
+        raise HTTPException(status_code=400, detail="Invalid authenticated user id")
+    return value
+
+
+def _extract_user_id_from_request(request: Request) -> Optional[str]:
+    # These headers are common in reverse-proxy / OAuth fronted deployments.
+    header_candidates = [
+        request.headers.get("x-user-id"),
+        request.headers.get("x-hf-username"),
+        request.headers.get("x-forwarded-user"),
+        request.headers.get("remote-user"),
+    ]
+    for candidate in header_candidates:
+        if candidate and candidate.strip():
+            return candidate.strip()
+    return None
+
+
+def get_current_user(
+    request: Request,
+    x_user_id: Optional[str] = Header(default=None, alias="X-User-Id"),
+    x_user_email: Optional[str] = Header(default=None, alias="X-User-Email"),
+) -> User:
+    user_id = (x_user_id or "").strip() or _extract_user_id_from_request(request)
+    if not user_id:
+        raise HTTPException(
+            status_code=401,
+            detail="Authentication required. Provide authenticated user context.",
+        )
+
+    user_id = _validate_user_id(user_id)
+    email = (x_user_email or f"{user_id}@local").strip()
+    return User(user_id=user_id, email=email)
+
+
+def enforce_user_match(current_user: User, supplied_user_id: Optional[str]) -> None:
+    supplied = (supplied_user_id or "").strip()
+    if supplied and supplied != current_user.user_id:
+        raise HTTPException(
+            status_code=403,
+            detail="Authenticated user does not match supplied user_id",
+        )

docs/rag_techniques.md

@@ -0,0 +1,50 @@
+# RAG Techniques Comparison
+
+This project currently supports two retrieval modes in chat:
+
+1. `topk`: baseline dense retrieval from Chroma using cosine distance.
+2. `rerank`: retrieves a larger candidate pool, then re-scores candidates using:
+   - vector relevance (from Chroma distance)
+   - lexical overlap with query terms
+
+## Where It Is Implemented
+
+- Request model: `backend/models/schemas.py` (`ChatRequest.retrieval_mode`)
+- Retrieval logic: `backend/modules/rag.py`
+- Chat API: `backend/api/chat.py`
+
+## How To Run Benchmark
+
+Prerequisites:
+- Backend running
+- Notebook has at least one ingested source
+
+Example:
+
+```bash
+python scripts/rag_benchmark.py \
+  --base-url http://127.0.0.1:8000 \
+  --user-id <user_id> \
+  --notebook-id <notebook_id> \
+  --query "Explain the key ideas in my notes" \
+  --top-k 5 \
+  --runs 5
+```
+
+The script prints JSON with average/min/max latency and citation/chunk stats for both modes.
+
+## Report Template
+
+Use this table in your class deliverable:
+
+| Query | Mode | Avg Latency (ms) | Avg Citations | Notes on Retrieved Context |
+|---|---:|---:|---:|---|
+| Q1 | topk |  |  |  |
+| Q1 | rerank |  |  |  |
+| Q2 | topk |  |  |  |
+| Q2 | rerank |  |  |  |
+
+Recommended write-up points:
+- How different the retrieved chunks were between `topk` and `rerank`
+- Which mode produced less redundant context
+- Latency tradeoff (rerank usually slightly slower)

frontend/app.py

@@ -80,7 +80,17 @@ def _maybe_start_local_backend() -> None:
         time.sleep(0.2)
 
 
-def _api_request(method: str, path: str, *, params=None, json_body=None, files=None, data=None, timeout: int = 60):
+def _api_request(
+    method: str,
+    path: str,
+    *,
+    params=None,
+    json_body=None,
+    files=None,
+    data=None,
+    headers=None,
+    timeout: int = 60,
+):
     url = f"{BACKEND_URL}{path}"
     try:
         resp = requests.request(
@@ -90,6 +100,7 @@ def _api_request(method: str, path: str, *, params=None, json_body=None, files=N
             json=json_body,
             files=files,
             data=data,
+            headers=headers,
             timeout=timeout,
         )
     except requests.RequestException as exc:
@@ -109,6 +120,11 @@ def _api_request(method: str, path: str, *, params=None, json_body=None, files=N
     return None
 
 
+def _auth_headers(user_id: str | None) -> dict[str, str]:
+    uid = (user_id or "").strip()
+    return {"X-User-Id": uid} if uid else {}
+
+
 def _format_notebook_choices(items: list[dict[str, Any]]):
     return [
         (f"{item.get('name', 'Untitled')} [{str(item.get('notebook_id', ''))[:8]}]", item.get("notebook_id"))
@@ -153,15 +169,25 @@ def load_notebooks(user_id: str):
     if not user_id:
         return gr.Dropdown(choices=[], value=None), [], gr.JSON(value={"sources": []}), [], "Enter a user ID."
     try:
-        items = _api_request("GET", "/api/notebooks/", params={"user_id": user_id})
+        items = _api_request("GET", "/api/notebooks/", params={"user_id": user_id}, headers=_auth_headers(user_id))
         choices = _format_notebook_choices(items)
         selected = choices[0][1] if choices else None
         chat_history = []
         sources_payload = {"sources": []}
         if selected:
-            chat = _api_request("GET", f"/api/notebooks/{selected}/chat", params={"user_id": user_id})
+            chat = _api_request(
+                "GET",
+                f"/api/notebooks/{selected}/chat",
+                params={"user_id": user_id},
+                headers=_auth_headers(user_id),
+            )
             chat_history = _messages_to_chatbot(chat.get("messages", []))
-            sources_payload = _api_request("GET", f"/api/notebooks/{selected}/sources", params={"user_id": user_id})
+            sources_payload = _api_request(
+                "GET",
+                f"/api/notebooks/{selected}/sources",
+                params={"user_id": user_id},
+                headers=_auth_headers(user_id),
+            )
         return gr.Dropdown(choices=choices, value=selected), items, gr.JSON(value=sources_payload), chat_history, ""
     except Exception as exc:
         return gr.Dropdown(choices=[], value=None), [], gr.JSON(value={"sources": []}), [], str(exc)
@@ -173,7 +199,12 @@ def create_notebook(user_id: str, notebook_name: str):
     if not user_id:
         return gr.Dropdown(choices=[], value=None), [], "", "Enter a user ID first."
     try:
-        _api_request("POST", "/api/notebooks/", json_body={"user_id": user_id, "name": notebook_name})
+        _api_request(
+            "POST",
+            "/api/notebooks/",
+            json_body={"user_id": user_id, "name": notebook_name},
+            headers=_auth_headers(user_id),
+        )
         dropdown, notebooks, _sources, _chat, _status = load_notebooks(user_id)
         return dropdown, notebooks, "", f"Created notebook '{notebook_name}'."
     except Exception as exc:
@@ -190,8 +221,9 @@ def rename_notebook(user_id: str, notebook_id: str, notebook_name: str):
             "PATCH",
             f"/api/notebooks/{notebook_id}",
             json_body={"user_id": user_id, "name": notebook_name},
+            headers=_auth_headers(user_id),
         )
-        items = _api_request("GET", "/api/notebooks/", params={"user_id": user_id})
+        items = _api_request("GET", "/api/notebooks/", params={"user_id": user_id}, headers=_auth_headers(user_id))
         choices = _format_notebook_choices(items)
         return gr.Dropdown(choices=choices, value=notebook_id), items, f"Renamed notebook to '{notebook_name}'."
     except Exception as exc:
@@ -203,7 +235,12 @@ def delete_notebook(user_id: str, notebook_id: str):
     if not user_id or not notebook_id:
         return gr.Dropdown(choices=[], value=None), [], gr.JSON(value={"sources": []}), [], "Select a notebook to delete."
     try:
-        _api_request("DELETE", f"/api/notebooks/{notebook_id}", params={"user_id": user_id})
+        _api_request(
+            "DELETE",
+            f"/api/notebooks/{notebook_id}",
+            params={"user_id": user_id},
+            headers=_auth_headers(user_id),
+        )
         dropdown, notebooks, sources_json, chat_history, _ = load_notebooks(user_id)
         return dropdown, notebooks, sources_json, chat_history, "Deleted notebook."
     except Exception as exc:
@@ -215,8 +252,18 @@ def on_notebook_change(user_id: str, notebook_id: str):
     if not user_id or not notebook_id:
         return gr.JSON(value={"sources": []}), [], ""
     try:
-        sources_payload = _api_request("GET", f"/api/notebooks/{notebook_id}/sources", params={"user_id": user_id})
-        chat = _api_request("GET", f"/api/notebooks/{notebook_id}/chat", params={"user_id": user_id})
+        sources_payload = _api_request(
+            "GET",
+            f"/api/notebooks/{notebook_id}/sources",
+            params={"user_id": user_id},
+            headers=_auth_headers(user_id),
+        )
+        chat = _api_request(
+            "GET",
+            f"/api/notebooks/{notebook_id}/chat",
+            params={"user_id": user_id},
+            headers=_auth_headers(user_id),
+        )
         return gr.JSON(value=sources_payload), _messages_to_chatbot(chat.get("messages", [])), ""
     except Exception as exc:
         return gr.JSON(value={"sources": []}), [], str(exc)
@@ -233,9 +280,15 @@ def upload_source(user_id: str, notebook_id: str, file_path: str):
                 f"/api/notebooks/{notebook_id}/sources/upload",
                 data={"user_id": user_id},
                 files={"file": (os.path.basename(file_path), f)},
+                headers=_auth_headers(user_id),
                 timeout=180,
             )
-        sources_payload = _api_request("GET", f"/api/notebooks/{notebook_id}/sources", params={"user_id": user_id})
+        sources_payload = _api_request(
+            "GET",
+            f"/api/notebooks/{notebook_id}/sources",
+            params={"user_id": user_id},
+            headers=_auth_headers(user_id),
+        )
         return gr.JSON(value=sources_payload), f"Ingested file: {payload.get('source_name', os.path.basename(file_path))}"
     except Exception as exc:
         return gr.JSON(value={"sources": []}), str(exc)
@@ -251,9 +304,15 @@ def ingest_url_source(user_id: str, notebook_id: str, url: str):
             "POST",
             f"/api/notebooks/{notebook_id}/sources/url",
             json_body={"user_id": user_id, "url": url},
+            headers=_auth_headers(user_id),
             timeout=180,
         )
-        sources_payload = _api_request("GET", f"/api/notebooks/{notebook_id}/sources", params={"user_id": user_id})
+        sources_payload = _api_request(
+            "GET",
+            f"/api/notebooks/{notebook_id}/sources",
+            params={"user_id": user_id},
+            headers=_auth_headers(user_id),
+        )
         return gr.JSON(value=sources_payload), f"Ingested URL: {payload.get('source_name', url)}"
     except Exception as exc:
         return gr.JSON(value={"sources": []}), str(exc)
@@ -272,6 +331,7 @@ def send_message(message: str, history, user_id: str, notebook_id: str):
             "POST",
             f"/api/notebooks/{notebook_id}/chat",
             json_body={"user_id": user_id, "message": message, "top_k": 5},
+            headers=_auth_headers(user_id),
             timeout=180,
         )
         assistant_text = str(resp.get("answer", "")) + _format_citations(resp.get("citations"))
@@ -313,7 +373,12 @@ def refresh_artifacts(user_id: str, notebook_id: str):
         payload = _empty_artifacts_payload()
         return gr.JSON(value=payload), None, None, None, None, "Select a notebook first."
     try:
-        payload = _api_request("GET", f"/api/notebooks/{notebook_id}/artifacts", params={"user_id": user_id})
+        payload = _api_request(
+            "GET",
+            f"/api/notebooks/{notebook_id}/artifacts",
+            params={"user_id": user_id},
+            headers=_auth_headers(user_id),
+        )
         report, quiz, transcript, audio = _artifact_outputs_from_payload(payload)
         return gr.JSON(value=payload), report, quiz, transcript, audio, ""
     except Exception as exc:
@@ -336,6 +401,7 @@ def generate_report_artifact(user_id: str, notebook_id: str, artifact_prompt: st
             "POST",
             f"/api/notebooks/{notebook_id}/artifacts/report",
             json_body={"user_id": user_id, "prompt": (artifact_prompt or "").strip() or None},
+            headers=_auth_headers(user_id),
             timeout=180,
         )
         payload_json, report, quiz, transcript, audio, _ = refresh_artifacts(user_id, notebook_id)
@@ -359,6 +425,7 @@ def generate_quiz_artifact(user_id: str, notebook_id: str, artifact_prompt: str,
                 "prompt": (artifact_prompt or "").strip() or None,
                 "num_questions": int(num_questions),
             },
+            headers=_auth_headers(user_id),
             timeout=180,
         )
         payload_json, report, quiz, transcript, audio, _ = refresh_artifacts(user_id, notebook_id)
@@ -378,6 +445,7 @@ def generate_podcast_artifact(user_id: str, notebook_id: str, artifact_prompt: s
             "POST",
             f"/api/notebooks/{notebook_id}/artifacts/podcast",
             json_body={"user_id": user_id, "prompt": (artifact_prompt or "").strip() or None},
+            headers=_auth_headers(user_id),
             timeout=240,
         )
         payload_json, report, quiz, transcript, audio, _ = refresh_artifacts(user_id, notebook_id)

scripts/rag_benchmark.py

@@ -0,0 +1,93 @@
+import argparse
+import json
+import time
+from statistics import mean
+from typing import Any
+
+import requests
+
+
+def _request(method: str, url: str, *, headers=None, params=None, json_body=None, timeout=180) -> dict[str, Any]:
+    resp = requests.request(
+        method=method,
+        url=url,
+        headers=headers,
+        params=params,
+        json=json_body,
+        timeout=timeout,
+    )
+    resp.raise_for_status()
+    if resp.content:
+        return resp.json()
+    return {}
+
+
+def run_benchmark(
+    *,
+    base_url: str,
+    user_id: str,
+    notebook_id: str,
+    query: str,
+    top_k: int,
+    runs: int,
+) -> dict[str, Any]:
+    headers = {"X-User-Id": user_id}
+    modes = ["topk", "rerank"]
+    results: dict[str, Any] = {}
+
+    for mode in modes:
+        latencies_ms: list[float] = []
+        citations_count: list[int] = []
+        used_chunks: list[int] = []
+
+        for _ in range(runs):
+            start = time.perf_counter()
+            payload = _request(
+                "POST",
+                f"{base_url}/api/notebooks/{notebook_id}/chat",
+                headers=headers,
+                json_body={
+                    "user_id": user_id,
+                    "message": query,
+                    "top_k": top_k,
+                    "retrieval_mode": mode,
+                },
+            )
+            elapsed_ms = (time.perf_counter() - start) * 1000.0
+            latencies_ms.append(elapsed_ms)
+            citations_count.append(len(payload.get("citations", [])))
+            used_chunks.append(int(payload.get("used_chunks", 0)))
+
+        results[mode] = {
+            "avg_latency_ms": round(mean(latencies_ms), 2),
+            "min_latency_ms": round(min(latencies_ms), 2),
+            "max_latency_ms": round(max(latencies_ms), 2),
+            "avg_citations": round(mean(citations_count), 2),
+            "avg_used_chunks": round(mean(used_chunks), 2),
+        }
+    return results
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Benchmark topk vs rerank retrieval modes.")
+    parser.add_argument("--base-url", default="http://127.0.0.1:8000")
+    parser.add_argument("--user-id", required=True)
+    parser.add_argument("--notebook-id", required=True)
+    parser.add_argument("--query", required=True)
+    parser.add_argument("--top-k", type=int, default=5)
+    parser.add_argument("--runs", type=int, default=5)
+    args = parser.parse_args()
+
+    results = run_benchmark(
+        base_url=args.base_url.rstrip("/"),
+        user_id=args.user_id,
+        notebook_id=args.notebook_id,
+        query=args.query,
+        top_k=args.top_k,
+        runs=max(1, args.runs),
+    )
+    print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()

tests/test_api_artifacts.py

@@ -12,6 +12,7 @@ from backend.services.storage import NotebookStore
 
 
 client = TestClient(app)
+AUTH_U1 = {"X-User-Id": "u1"}
 
 
 def _seed_source(store: NotebookStore, user_id: str, notebook_id: str, source_id: str = "src_demo") -> None:
@@ -41,17 +42,23 @@ def test_artifact_endpoints_generate_list_and_download(monkeypatch, tmp_path: Pa
     report = client.post(
         f"/api/notebooks/{nb.notebook_id}/artifacts/report",
         json={"user_id": "u1", "prompt": "Focus on definitions"},
+        headers=AUTH_U1,
     )
     assert report.status_code == 200
 
     podcast = client.post(
         f"/api/notebooks/{nb.notebook_id}/artifacts/podcast",
         json={"user_id": "u1"},
+        headers=AUTH_U1,
     )
     assert podcast.status_code == 200
     podcast_audio_name = Path(podcast.json()["audio_path"]).name
 
-    listed = client.get(f"/api/notebooks/{nb.notebook_id}/artifacts", params={"user_id": "u1"})
+    listed = client.get(
+        f"/api/notebooks/{nb.notebook_id}/artifacts",
+        params={"user_id": "u1"},
+        headers=AUTH_U1,
+    )
     assert listed.status_code == 200
     payload = listed.json()
     assert len(payload["reports"]) == 1
@@ -60,6 +67,7 @@ def test_artifact_endpoints_generate_list_and_download(monkeypatch, tmp_path: Pa
     dl = client.get(
         f"/api/notebooks/{nb.notebook_id}/artifacts/download",
         params={"user_id": "u1", "artifact_type": "podcast", "filename": podcast_audio_name},
+        headers=AUTH_U1,
     )
     assert dl.status_code == 200
     assert dl.content.startswith(b"ID3")

tests/test_api_chat.py

@@ -8,6 +8,7 @@ from backend.services.storage import NotebookStore
 
 
 client = TestClient(app)
+AUTH_U1 = {"X-User-Id": "u1"}
 
 
 def test_chat_endpoint_success_with_mock(monkeypatch, tmp_path):
@@ -35,7 +36,8 @@ def test_chat_endpoint_success_with_mock(monkeypatch, tmp_path):
 
     resp = client.post(
         f"/api/notebooks/{created.notebook_id}/chat",
-        json={"user_id": "u1", "message": "Hi", "top_k": 3},
+        json={"user_id": "u1", "message": "Hi", "top_k": 3, "retrieval_mode": "rerank"},
+        headers=AUTH_U1,
     )
     assert resp.status_code == 200
     data = resp.json()
@@ -61,7 +63,11 @@ def test_chat_history_endpoint_reads_jsonl(tmp_path):
         },
     )
 
-    resp = client.get(f"/api/notebooks/{created.notebook_id}/chat", params={"user_id": "u1"})
+    resp = client.get(
+        f"/api/notebooks/{created.notebook_id}/chat",
+        params={"user_id": "u1"},
+        headers=AUTH_U1,
+    )
     assert resp.status_code == 200
     payload = resp.json()
     assert len(payload["messages"]) == 2

tests/test_api_notebooks.py

@@ -6,6 +6,8 @@ from backend.services.storage import NotebookStore
 
 
 client = TestClient(app)
+AUTH_U1 = {"X-User-Id": "u1"}
+AUTH_U2 = {"X-User-Id": "u2"}
 
 
 def setup_function():
@@ -16,7 +18,7 @@ def test_create_list_get_rename_delete_notebook(tmp_path):
     notebooks_api.store = NotebookStore(base_dir=str(tmp_path))
 
     payload = {"user_id": "u1", "name": "Notebook 1"}
-    resp = client.post("/api/notebooks/", json=payload)
+    resp = client.post("/api/notebooks/", json=payload, headers=AUTH_U1)
     assert resp.status_code == 200
     data = resp.json()
     assert data["user_id"] == "u1"
@@ -25,33 +27,44 @@ def test_create_list_get_rename_delete_notebook(tmp_path):
 
     notebook_id = data["notebook_id"]
 
-    resp_list = client.get("/api/notebooks/", params={"user_id": "u1"})
+    resp_list = client.get("/api/notebooks/", params={"user_id": "u1"}, headers=AUTH_U1)
     assert resp_list.status_code == 200
     listed = resp_list.json()
     assert len(listed) == 1
     assert listed[0]["notebook_id"] == notebook_id
 
-    resp_get = client.get(f"/api/notebooks/{notebook_id}", params={"user_id": "u1"})
+    resp_get = client.get(f"/api/notebooks/{notebook_id}", params={"user_id": "u1"}, headers=AUTH_U1)
     assert resp_get.status_code == 200
     data_get = resp_get.json()
     assert data_get["notebook_id"] == notebook_id
 
-    resp_wrong_user = client.get(f"/api/notebooks/{notebook_id}", params={"user_id": "u2"})
+    resp_wrong_user = client.get(f"/api/notebooks/{notebook_id}", params={"user_id": "u2"}, headers=AUTH_U2)
     assert resp_wrong_user.status_code == 404
 
     resp_rename = client.patch(
         f"/api/notebooks/{notebook_id}",
         json={"user_id": "u1", "name": "Renamed"},
+        headers=AUTH_U1,
     )
     assert resp_rename.status_code == 200
     assert resp_rename.json()["name"] == "Renamed"
 
-    resp_delete = client.delete(f"/api/notebooks/{notebook_id}", params={"user_id": "u1"})
+    resp_delete = client.delete(f"/api/notebooks/{notebook_id}", params={"user_id": "u1"}, headers=AUTH_U1)
     assert resp_delete.status_code == 200
     assert resp_delete.json() == {"deleted": True}
 
 
 def test_get_missing_notebook(tmp_path):
     notebooks_api.store = NotebookStore(base_dir=str(tmp_path))
-    resp = client.get("/api/notebooks/missing", params={"user_id": "u1"})
+    resp = client.get("/api/notebooks/missing", params={"user_id": "u1"}, headers=AUTH_U1)
     assert resp.status_code == 404
+
+
+def test_rejects_mismatched_user_id_payload(tmp_path):
+    notebooks_api.store = NotebookStore(base_dir=str(tmp_path))
+    resp = client.post(
+        "/api/notebooks/",
+        json={"user_id": "u2", "name": "Notebook 1"},
+        headers=AUTH_U1,
+    )
+    assert resp.status_code == 403

tests/test_rag.py

@@ -59,3 +59,27 @@ def test_answer_notebook_question_persists_messages(monkeypatch, tmp_path):
     assert len(messages) == 2
     assert messages[0]["role"] == "user"
     assert messages[1]["role"] == "assistant"
+
+
+def test_rerank_chunks_prefers_lexically_relevant_chunk():
+    candidates = [
+        {
+            "chunk_id": "a",
+            "text": "general overview and intro",
+            "metadata": {},
+            "distance": 0.15,
+        },
+        {
+            "chunk_id": "b",
+            "text": "transformer attention heads and query key value details",
+            "metadata": {},
+            "distance": 0.20,
+        },
+    ]
+    reranked = rag.rerank_chunks(
+        query="explain transformer attention",
+        candidate_chunks=candidates,
+        top_k=1,
+    )
+    assert len(reranked) == 1
+    assert reranked[0]["chunk_id"] == "b"