app.py

import os
import uuid
import json
import hashlib
import sqlite3

import fitz  # PyMuPDF
import requests
from flask import Flask, render_template, request, jsonify, session, Response

app = Flask(__name__)
app.secret_key = os.urandom(24)

OLLAMA_URL = "http://localhost:11434/api/chat"
MODEL = os.environ.get("OLLAMA_MODEL", "qwen3.5:35b-a3b")
UPLOAD_FOLDER = os.path.join(os.path.dirname(__file__), "uploads")
DB_PATH = os.path.join(os.path.dirname(__file__), "cache.db")
MAX_CONTEXT_CHARS = 120_000  # ~60K tokens, leaves room for conversation

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# In-memory storage: session_id -> {"text": str, "history": list}
store = {}


def init_db():
    conn = sqlite3.connect(DB_PATH)
    conn.execute(
        """CREATE TABLE IF NOT EXISTS pdf_cache (
            hash TEXT PRIMARY KEY,
            filename TEXT,
            text TEXT NOT NULL,
            summary TEXT NOT NULL
        )"""
    )
    conn.commit()
    conn.close()


init_db()


def get_cache(pdf_hash):
    conn = sqlite3.connect(DB_PATH)
    row = conn.execute(
        "SELECT text, summary FROM pdf_cache WHERE hash = ?", (pdf_hash,)
    ).fetchone()
    conn.close()
    if row:
        return {"text": row[0], "summary": row[1]}
    return None


def set_cache(pdf_hash, filename, text, summary):
    conn = sqlite3.connect(DB_PATH)
    conn.execute(
        "INSERT OR REPLACE INTO pdf_cache (hash, filename, text, summary) VALUES (?, ?, ?, ?)",
        (pdf_hash, filename, text, summary),
    )
    conn.commit()
    conn.close()


def extract_text(pdf_path):
    doc = fitz.open(pdf_path)
    text = ""
    for page in doc:
        text += page.get_text()
    doc.close()
    return text


def hash_file(filepath):
    h = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def ollama_stream(messages):
    resp = requests.post(
        OLLAMA_URL,
        json={"model": MODEL, "messages": messages, "stream": True},
        stream=True,
        timeout=300,
    )
    resp.raise_for_status()
    for line in resp.iter_lines():
        if line:
            data = json.loads(line)
            content = data.get("message", {}).get("content", "")
            if content:
                yield content
            if data.get("done"):
                break


def get_session_id():
    if "sid" not in session:
        session["sid"] = str(uuid.uuid4())
    return session["sid"]


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/upload", methods=["POST"])
def upload():
    if "file" not in request.files:
        return jsonify({"error": "No file provided"}), 400

    file = request.files["file"]
    if not file.filename.lower().endswith(".pdf"):
        return jsonify({"error": "Only PDF files are supported"}), 400

    sid = get_session_id()
    filepath = os.path.join(UPLOAD_FOLDER, f"{sid}.pdf")
    file.save(filepath)

    pdf_hash = hash_file(filepath)
    cached = get_cache(pdf_hash)

    if cached:
        os.remove(filepath)
        truncated = cached["text"]
        summary = cached["summary"]
        store[sid] = {
            "text": truncated,
            "history": [{"role": "assistant", "content": summary}],
        }

        def generate_cached():
            yield f"data: {json.dumps({'stage': 'cache_hit'})}\n\n"
            yield f"data: {json.dumps({'content': summary})}\n\n"
            yield "data: [DONE]\n\n"

        return Response(generate_cached(), mimetype="text/event-stream")

    # Not cached — extract and summarize
    def generate():
        yield f"data: {json.dumps({'stage': 'extracting'})}\n\n"

        text = extract_text(filepath)
        os.remove(filepath)

        if not text.strip():
            yield f"data: {json.dumps({'error': 'Could not extract text from this PDF. It may be scanned/image-based.'})}\n\n"
            return

        truncated = text[:MAX_CONTEXT_CHARS]
        store[sid] = {"text": truncated, "history": []}

        yield f"data: {json.dumps({'stage': 'summarizing'})}\n\n"

        messages = [
            {
                "role": "system",
                "content": "You are a helpful assistant. The user has uploaded a PDF document. Provide a clear, well-structured summary of its contents.",
            },
            {
                "role": "user",
                "content": f"Please summarize this document:\n\n{truncated}",
            },
        ]
        full_response = []
        for chunk in ollama_stream(messages):
            full_response.append(chunk)
            yield f"data: {json.dumps({'content': chunk})}\n\n"

        summary = "".join(full_response)
        store[sid]["history"] = [
            {"role": "assistant", "content": summary},
        ]

        set_cache(pdf_hash, file.filename, truncated, summary)

        yield "data: [DONE]\n\n"

    return Response(generate(), mimetype="text/event-stream")


@app.route("/chat", methods=["POST"])
def chat():
    sid = get_session_id()
    data = request.get_json()
    user_message = data.get("message", "").strip()

    if not user_message:
        return jsonify({"error": "Empty message"}), 400

    if sid not in store:
        return jsonify({"error": "No PDF uploaded yet. Please upload a PDF first."}), 400

    pdf_text = store[sid]["text"]
    history = store[sid]["history"]

    history.append({"role": "user", "content": user_message})

    messages = [
        {
            "role": "system",
            "content": f"You are a helpful assistant answering questions about the following document. Base your answers on the document content.\n\n--- DOCUMENT ---\n{pdf_text}\n--- END DOCUMENT ---",
        },
    ] + history

    def generate():
        full_response = []
        for chunk in ollama_stream(messages):
            full_response.append(chunk)
            yield f"data: {json.dumps({'content': chunk})}\n\n"

        history.append({"role": "assistant", "content": "".join(full_response)})
        yield "data: [DONE]\n\n"

    return Response(generate(), mimetype="text/event-stream")


if __name__ == "__main__":
    app.run(debug=True, port=5005)