app.py

import os
from cs50 import SQL
from flask import Flask, flash, redirect, render_template, request, session
from werkzeug.utils import secure_filename
from pypdf import PdfReader
from helpers import apology, split_into_chunks, compute_tfidf

app = Flask(__name__)

app.config["TEMPLATES_AUTO_RELOAD"] = True
app.config["UPLOAD_FOLDER"] = "uploads"
app.config["ALLOWED_EXTENSIONS"] = {"txt", "pdf"}

os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)

db = SQL("sqlite:///assistant.db")

db.execute("""
    CREATE TABLE IF NOT EXISTS documents (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        filename TEXT NOT NULL,
        uploaded_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
""")

db.execute("""
    CREATE TABLE IF NOT EXISTS chunks (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        doc_id INTEGER,
        chunk_text TEXT NOT NULL,
        FOREIGN KEY(doc_id) REFERENCES documents(id)
    )
""")

db.execute("""
    CREATE TABLE IF NOT EXISTS history (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        query TEXT NOT NULL,
        answer TEXT NOT NULL,
        source_doc TEXT,
        timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
""")


def allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in app.config["ALLOWED_EXTENSIONS"]


@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        query = request.form.get("query")
        if not query:
            return apology("Must provide a search query", 400)

        chunks_data = db.execute("""
            SELECT chunks.chunk_text, documents.filename 
            FROM chunks 
            JOIN documents ON chunks.doc_id = documents.id
        """)

        if not chunks_data:
            return apology("No documents uploaded yet. Please upload a document first.", 400)

        chunks = [row["chunk_text"] for row in chunks_data]
        best_chunk, score = compute_tfidf(query, chunks)

        if score == 0:
            best_chunk = "No relevant answer found in the uploaded documents."
            filename = "N/A"
        else:
            filename = next(row["filename"] for row in chunks_data if row["chunk_text"] == best_chunk)

        db.execute(
            "INSERT INTO history (query, answer, source_doc) VALUES (?, ?, ?)",
            query, best_chunk, filename
        )

        return render_template("index.html", result=best_chunk, score=score, filename=filename, query=query)

    return render_template("index.html", result=None)


@app.route("/upload", methods=["GET", "POST"])
def upload():
    if request.method == "POST":
        if "file" not in request.files:
            return apology("No file part", 400)
        
        file = request.files["file"]
        if file.filename == "":
            return apology("No selected file", 400)

        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config["UPLOAD_FOLDER"], filename)
            file.save(filepath)

            text = ""
            if filename.endswith(".txt"):
                with open(filepath, "r", encoding="utf-8") as f:
                    text = f.read()
            elif filename.endswith(".pdf"):
                reader = PdfReader(filepath)
                for page in reader.pages:
                    text += page.extract_text() or ""

            if not text.strip():
                return apology("Could not extract text from file", 400)

            doc_id = db.execute("INSERT INTO documents (filename) VALUES (?)", filename)

            chunks = split_into_chunks(text)
            for chunk in chunks:
                db.execute("INSERT INTO chunks (doc_id, chunk_text) VALUES (?, ?)", doc_id, chunk)

            return redirect("/")

        return apology("Invalid file extension", 400)

    return render_template("upload.html")


@app.route("/history")
def history():
    logs = db.execute("SELECT * FROM history ORDER BY timestamp DESC")
    return render_template("history.html", history=logs)


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