#!/usr/bin/env python
"""Chat with your PDF β upload a PDF and ask questions; answers are grounded in the
document with page citations (real RAG: chunk -> embed -> retrieve -> answer).
Embeddings: fastembed (ONNX, CPU-friendly, no torch). LLM: Anthropic Claude API if
ANTHROPIC_API_KEY is set, else a local Ollama model.
pip install flask anthropic pypdf fastembed numpy
python app.py # http://127.0.0.1:8500
Project #3 of the "30 Projects in 15 Days" challenge β GritAI.
"""
import os, io, json, re, uuid
import numpy as np
from flask import Flask, request, Response, jsonify
HERE = os.path.dirname(os.path.abspath(__file__))
PORT = int(os.environ.get("PORT", "8500"))
ANTHROPIC_KEY = os.environ.get("ANTHROPIC_API_KEY")
ANTHROPIC_MODEL = os.environ.get("ANTHROPIC_MODEL", "claude-sonnet-5")
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "qwen2.5:7b")
MAX_PAGES = 120
CHUNK_CHARS = 1200
CHUNK_OVERLAP = 160
TOP_K = 6
MAX_DOCS = 24
SAMPLE_DIR = os.path.join(HERE, "sample_pdfs")
SAMPLES = {
"memo": ("simple_memo.pdf", "π", "Office memo", "1-page internal memo"),
"manual": ("product_manual.pdf", "π", "Product manual", "Smart thermostat guide"),
"report": ("research_report.pdf", "π", "Market report", "Data, tables, findings"),
"handbook": ("employee_handbook.pdf", "π", "Employee handbook", "13 pages, many sections"),
}
# ---- embedding model (lazy singleton) ----------------------------------------
_EMB = None
def emb_model():
global _EMB
if _EMB is None:
from fastembed import TextEmbedding
_EMB = TextEmbedding("BAAI/bge-small-en-v1.5")
return _EMB
def embed_passages(texts):
return np.array(list(emb_model().passage_embed(texts)), dtype=np.float32)
def embed_query(text):
return np.array(list(emb_model().query_embed([text]))[0], dtype=np.float32)
def _norm(m):
n = np.linalg.norm(m, axis=-1, keepdims=True)
return m / np.clip(n, 1e-8, None)
# ---- pdf -> chunks -----------------------------------------------------------
def extract_pages(data):
from pypdf import PdfReader
reader = PdfReader(io.BytesIO(data))
title = (reader.metadata.title if reader.metadata and reader.metadata.title else "").strip()
pages = []
for i, pg in enumerate(reader.pages[:MAX_PAGES], start=1):
txt = (pg.extract_text() or "").strip()
if txt:
pages.append((i, re.sub(r"[ \t]+", " ", txt)))
return title, len(reader.pages), pages
def chunk_pages(pages):
chunks = []
for pageno, text in pages:
i = 0
while i < len(text):
piece = text[i:i + CHUNK_CHARS].strip()
if len(piece) > 40:
chunks.append({"page": pageno, "text": piece})
i += CHUNK_CHARS - CHUNK_OVERLAP
return chunks
DOCS = {}
DOC_ORDER = []
def index_pdf(data, filename):
title, npages, pages = extract_pages(data)
if not pages:
raise ValueError("No extractable text found β this PDF may be scanned images (needs OCR).")
chunks = chunk_pages(pages)
emb = _norm(embed_passages([c["text"] for c in chunks]))
doc_id = uuid.uuid4().hex[:12]
disp = re.sub(r"\.(html?|pdf)$", "", (title or ""), flags=re.I).strip()
if not disp:
disp = re.sub(r"\.pdf$", "", filename or "document", flags=re.I)
DOCS[doc_id] = {"title": disp, "pages": npages, "chunks": chunks, "emb": emb}
DOC_ORDER.append(doc_id)
while len(DOC_ORDER) > MAX_DOCS:
DOCS.pop(DOC_ORDER.pop(0), None)
return doc_id, DOCS[doc_id]
def retrieve(doc, question, k=TOP_K):
q = _norm(embed_query(question).reshape(1, -1))[0]
sims = doc["emb"] @ q
idx = np.argsort(-sims)[:k]
return [doc["chunks"][i] for i in idx]
# ---- answer (streaming) ------------------------------------------------------
def system_prompt():
return (
"You answer questions about a specific document using ONLY the excerpts provided.\n"
"- Cite the page number(s) you used, in parentheses like (p. 4).\n"
"- If the answer is not in the excerpts, say you couldn't find it in the document β do not guess.\n"
"- Be concise and faithful. Write plain text; no markdown symbols like ** or #."
)
def build_context(chunks):
return "\n\n".join(f"[p. {c['page']}] {c['text']}" for c in chunks)
def answer_from(chunks, question):
ctx = build_context(chunks)
user = f"Document excerpts:\n\n{ctx}\n\n---\nQuestion: {question}"
if ANTHROPIC_KEY:
import anthropic
client = anthropic.Anthropic(api_key=ANTHROPIC_KEY)
with client.messages.stream(model=ANTHROPIC_MODEL, max_tokens=700,
system=system_prompt(),
messages=[{"role": "user", "content": user}]) as s:
for t in s.text_stream:
yield t
return
import requests
with requests.post(f"{OLLAMA_URL}/api/chat", stream=True, timeout=120, json={
"model": OLLAMA_MODEL, "stream": True,
"messages": [{"role": "system", "content": system_prompt()},
{"role": "user", "content": user}],
"options": {"temperature": 0.2}}) as r:
r.raise_for_status()
for line in r.iter_lines():
if not line:
continue
d = json.loads(line)
if d.get("message", {}).get("content"):
yield d["message"]["content"]
if d.get("done"):
break
app = Flask(__name__)
app.config["MAX_CONTENT_LENGTH"] = 20 * 1024 * 1024
@app.route("/")
def home():
return Response(PAGE, mimetype="text/html")
@app.route("/api/upload", methods=["POST"])
def api_upload():
f = request.files.get("pdf")
if not f:
return jsonify(error="No file uploaded."), 400
try:
doc_id, doc = index_pdf(f.read(), f.filename or "document.pdf")
return jsonify(doc_id=doc_id, title=doc["title"], pages=doc["pages"], chunks=len(doc["chunks"]))
except Exception as e:
return jsonify(error="Could not process that β please try again."), 400
@app.route("/api/sample", methods=["POST"])
def api_sample():
name = request.get_json(force=True).get("name")
meta = SAMPLES.get(name)
if not meta:
return jsonify(error="unknown sample"), 400
path = os.path.join(SAMPLE_DIR, meta[0])
if not os.path.exists(path):
return jsonify(error="sample not available"), 404
try:
with open(path, "rb") as fh:
doc_id, doc = index_pdf(fh.read(), meta[0])
return jsonify(doc_id=doc_id, title=doc["title"], pages=doc["pages"], chunks=len(doc["chunks"]))
except Exception as e:
return jsonify(error="Could not process that β please try again."), 400
@app.route("/api/ask", methods=["POST"])
def api_ask():
b = request.get_json(force=True)
doc = DOCS.get(b.get("doc_id", ""))
question = (b.get("question") or "").strip()
def gen():
import sys as _sys
if not doc:
yield "data: " + json.dumps({"t": "That document isn't loaded anymore β please re-upload it."}) + "\n\n"
elif not question:
yield "data: " + json.dumps({"t": "Please ask a question."}) + "\n\n"
else:
try:
chunks = retrieve(doc, question)
src = {}
for c in chunks:
src.setdefault(str(c["page"]), c["text"][:420])
yield "data: " + json.dumps({"sources": src}) + "\n\n"
for chunk in answer_from(chunks, question):
yield "data: " + json.dumps({"t": chunk}) + "\n\n"
except Exception as e:
print("ask error:", type(e).__name__, file=_sys.stderr, flush=True)
yield "data: " + json.dumps({"t": "Sorry β I hit an error answering. Please try again."}) + "\n\n"
yield "data: [DONE]\n\n"
return Response(gen(), mimetype="text/event-stream",
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
PAGE = """
Chat with your PDF
GRITAIChat with your PDF
▸ LiveSDVOSBLawton OK
Ask your documents anything
Upload a PDF and get answers grounded in its contents — every answer cites the exact page it came from, and it won’t guess.
01Upload
+
Drop a PDF here or click to upload
Up to 20 MB · text-based PDFs
Instant answersCites the pageWon’t guess
02Or try a sample
01Ask
PDF
document
"""
if __name__ == "__main__":
backend = "Claude API" if ANTHROPIC_KEY else f"Ollama ({OLLAMA_MODEL} @ {OLLAMA_URL})"
print(f"Chat with your PDF on http://127.0.0.1:{PORT} [backend: {backend}]", flush=True)
app.run(host="0.0.0.0", port=PORT, threaded=True)