feat(rag): PDF/docx ingest + chat citations

- Upload endpoint (base64 JSON, no multipart dep): extracts text from
  pdf/docx/txt/md via pypdf + python-docx, then runs the existing
  chunk/embed pipeline. Documents page uploads files straight through.
- Citations: the chat stream emits an SSE `sources` event listing the
  documents that fed the answer; the UI shows them as chips under the reply.
- Deps: pypdf, python-docx (both pure-Python, Windows-safe).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
jon
2026-07-23 13:48:32 -05:00
co-authored by Claude Opus 4.8
parent f509034fdd
commit ac0eb1e5f6
6 changed files with 123 additions and 9 deletions
+46
View File
@@ -287,10 +287,12 @@ async def chat_stream_endpoint(payload: Dict[str, Any]):
# Retrieve relevant uploaded documents (RAG) and inject the top chunks.
doc_hits = await store.search_documents(message, get_ollama_manager().embed, limit=3)
doc_titles: list = []
if doc_hits:
doc_block = "\n\n".join(f"[{d['title']}]\n{d['text']}" for d in doc_hits)
separator = "\n\n---\nRelevant documents (cite as source material):\n\n"
system_prompt = (system_prompt + separator + doc_block) if system_prompt else doc_block
doc_titles = list(dict.fromkeys(d["title"] for d in doc_hits)) # unique, order-preserving
# Fetch web search results for time-sensitive queries
search_results = ""
@@ -373,6 +375,10 @@ async def chat_stream_endpoint(payload: Dict[str, Any]):
meta: dict = {}
final_model = model
# Tell the client which documents fed this answer (RAG citations).
if doc_titles:
yield f"event: sources\ndata: {_json.dumps({'sources': doc_titles})}\n\n"
# ── Phase 1: stream primary model response ────────────────────
try:
async for chunk in stream_chat_response(
@@ -953,6 +959,46 @@ async def add_document(payload: Dict[str, Any] = Body(...)):
return result
def _extract_text(filename: str, data: bytes) -> str:
"""Pull plain text from an uploaded file by extension. PDF/DOCX use
pure-Python parsers; anything else is decoded as UTF-8."""
import io
name = (filename or "").lower()
if name.endswith(".pdf"):
from pypdf import PdfReader
reader = PdfReader(io.BytesIO(data))
return "\n\n".join((page.extract_text() or "") for page in reader.pages)
if name.endswith(".docx"):
import docx
doc = docx.Document(io.BytesIO(data))
return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip())
return data.decode("utf-8", errors="replace")
@app.post("/documents/upload")
async def upload_document(payload: Dict[str, Any] = Body(...)):
"""Ingest a file (pdf/docx/txt/md) sent as base64. Extracts text, then runs
the same chunk/embed pipeline as a pasted document."""
import base64
import os as _os
filename = (payload.get("filename") or "").strip()
b64 = payload.get("data") or ""
if not filename or not b64:
raise HTTPException(status_code=400, detail="filename and data are required")
try:
raw = base64.b64decode(b64)
except Exception:
raise HTTPException(status_code=400, detail="data must be base64")
try:
text = _extract_text(filename, raw).strip()
except Exception as e:
raise HTTPException(status_code=400, detail=f"could not read {filename}: {e}")
if not text:
raise HTTPException(status_code=400, detail="no extractable text in file")
title = _os.path.splitext(_os.path.basename(filename))[0] or filename
return await store.add_document(title, text, get_ollama_manager().embed)
@app.get("/documents/{doc_id}")
async def get_document(doc_id: str):
chunks = store.get_document(doc_id)