assistant/assistant.py

import os
import re
import pdfplumber
from flask import Blueprint, request, jsonify
from werkzeug.utils import secure_filename

from assistant.storage import (
    generate_doc_id,
    save_document,
    load_document,
    save_metadata,
    get_all_docs,
)

assistant_bp = Blueprint('assistant_bp', __name__)
UPLOAD_FOLDER = './uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

FALLBACK = "This information is not available in your uploaded notes."

STOP_WORDS = {
    "what", "is", "a", "an", "the", "of", "to", "in", "and", "why",
    "how", "who", "when", "where", "are", "do", "does", "did", "for",
    "on", "it", "this", "that", "tell", "me", "about", "define",
    "explain", "give", "meaning", "can", "you", "please", "with",
    "was", "will", "has", "have", "had", "be", "been", "by", "or",
    "at", "from", "as", "its", "my", "your"
}

# ---------------------------------------------------------------------------
# POST /upload_notes
# ---------------------------------------------------------------------------
@assistant_bp.route('/upload_notes', methods=['POST'])
def upload_notes():
    if 'file' not in request.files:
        return jsonify({"error": "No file part in request"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    filename = secure_filename(file.filename)
    filepath = os.path.join(UPLOAD_FOLDER, filename)
    file.save(filepath)

    try:
        text = ""
        with pdfplumber.open(filepath) as pdf:
            for page in pdf.pages:
                extracted = page.extract_text()
                if extracted:
                    text += extracted + "\n"

        if not text.strip():
            return jsonify({"error": "Could not extract any text from PDF"}), 422

        doc_id = generate_doc_id()
        save_document(doc_id, text)
        save_metadata(doc_id, filename)

        return jsonify({"doc_id": doc_id, "message": "Saved successfully"}), 200

    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        if os.path.exists(filepath):
            os.remove(filepath)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _extract_keywords(question: str) -> list:
    words = re.findall(r'\b\w+\b', question.lower())
    keywords = [w for w in words if w not in STOP_WORDS and len(w) > 1]
    return keywords if keywords else words

# ---------------------------------------------------------------------------
# POST /ask
# ---------------------------------------------------------------------------
@assistant_bp.route('/ask', methods=['POST'])
def ask_question():
    data = request.get_json()
    if not data:
        return jsonify({"error": "Request body must be JSON"}), 400
    if 'doc_id' not in data or 'question' not in data:
        return jsonify({"error": "Missing 'doc_id' or 'question'"}), 400

    doc_id   = data['doc_id']
    question = data['question'].strip()

    if not question:
        return jsonify({"answer": "Please provide a non-empty question."}), 200

    # STEP 1 — Load document
    text = load_document(doc_id)
    if text is None:
        return jsonify({"error": f"Document '{doc_id}' not found"}), 404

    # STEP 2 — Extract Keywords
    keywords = _extract_keywords(question)
    
    # STEP 3 — Match context with line tracking (Pure RAG)
    lines = text.split("\n")
    matched_data = []

    for i, line in enumerate(lines):
        line_lower = line.lower()
        for kw in keywords:
            if kw in line_lower:
                matched_data.append({
                    "line": i + 1,
                    "text": line.strip()
                })
                break
        if len(matched_data) >= 3: # Keep it concise for direct RAG
            break

    # Build context string
    context = " ".join([item["text"] for item in matched_data])

    # STEP 4 — Guard: check context length
    if len(context.strip()) < 30:
        return jsonify({
            "answer": FALLBACK,
            "source": None
        }), 200

    # STEP 5 — Determine source metadata (first match)
    first_match = matched_data[0]
    # Find the column index of the first keyword in the first matched line
    column_index = first_match["text"].lower().find(keywords[0]) if keywords else 0

    # STEP 6 — Return the raw extracted context natively as the answer
    answer = f"According to your notes: {context}"

    return jsonify({
        "answer": answer,
        "source": {
            "line": first_match["line"],
            "column": column_index
        }
    }), 200

# ---------------------------------------------------------------------------
# GET /documents
# ---------------------------------------------------------------------------
@assistant_bp.route('/documents', methods=['GET'])
def list_documents():
    return jsonify(get_all_docs()), 200