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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user