notes/blueprints/chat.py

import json
from flask import Blueprint, request, jsonify, current_app
from flask_login import login_required, current_user
from models import db, ChatSession, ChatMessage
from services import rag_service, llm_service

chat_bp = Blueprint("chat", __name__, url_prefix="/api/chat")


# ── Session management ──────────────────────────────────────────────────────

@chat_bp.route("/sessions", methods=["GET"])
@login_required
def list_sessions():
    sessions = (
        ChatSession.query.filter_by(user_id=current_user.id)
        .order_by(ChatSession.updated_at.desc())
        .limit(50)
        .all()
    )
    return jsonify([s.to_dict() for s in sessions])


@chat_bp.route("/sessions", methods=["POST"])
@login_required
def create_session():
    data = request.get_json(silent=True) or {}
    session = ChatSession(user_id=current_user.id, title=data.get("title", "New Chat"))
    db.session.add(session)
    db.session.commit()
    return jsonify(session.to_dict()), 201


@chat_bp.route("/sessions/<int:session_id>", methods=["DELETE"])
@login_required
def delete_session(session_id):
    session = ChatSession.query.filter_by(id=session_id, user_id=current_user.id).first_or_404()
    db.session.delete(session)
    db.session.commit()
    return jsonify({"success": True})


@chat_bp.route("/sessions/<int:session_id>/messages", methods=["GET"])
@login_required
def get_messages(session_id):
    session = ChatSession.query.filter_by(id=session_id, user_id=current_user.id).first_or_404()
    return jsonify([m.to_dict() for m in session.messages])


# ── Main chat ────────────────────────────────────────────────────────────────

@chat_bp.route("/sessions/<int:session_id>/ask", methods=["POST"])
@login_required
def ask(session_id):
    session = ChatSession.query.filter_by(id=session_id, user_id=current_user.id).first_or_404()
    data = request.get_json(silent=True) or {}

    message = (data.get("message") or "").strip()
    if not message:
        return jsonify({"error": "No message provided"}), 400

    # context_ids: list of objects {"id": int, "type": "doc"|"url"}
    context_refs = data.get("context_ids", [])
    doc_ids = [r["id"] for r in context_refs if r.get("type") == "doc"]
    url_ids = [r["id"] for r in context_refs if r.get("type") == "url"]

    # RAG lookup
    chunks = []
    if doc_ids:
        chunks += rag_service.similarity_search(
            query=message,
            user_id=current_user.id,
            source_ids=doc_ids,
            source_type="doc",
            top_k=current_app.config["RAG_TOP_K"],
        )
    if url_ids:
        chunks += rag_service.similarity_search(
            query=message,
            user_id=current_user.id,
            source_ids=url_ids,
            source_type="url",
            top_k=current_app.config["RAG_TOP_K"],
        )
    # If no specific ids given, search all user context
    if not context_refs:
        chunks = rag_service.similarity_search(
            query=message,
            user_id=current_user.id,
            top_k=current_app.config["RAG_TOP_K"],
        )

    # Build history (last 10 messages for context window)
    history = [
        {"role": m.role, "content": m.content}
        for m in session.messages[-10:]
    ]

    try:
        reply = llm_service.ask(
            user_message=message,
            context_chunks=chunks,
            history=history,
        )
    except Exception as e:
        current_app.logger.error(f"LLM error: {e}")
        return jsonify({"error": f"LLM request failed: {str(e)}"}), 502

    # Persist messages
    user_msg = ChatMessage(
        session_id=session.id,
        role="user",
        content=message,
        context_ids=json.dumps(context_refs),
    )
    assistant_msg = ChatMessage(
        session_id=session.id,
        role="assistant",
        content=reply,
    )
    db.session.add_all([user_msg, assistant_msg])

    # Update session title after first user message
    if len(session.messages) == 0:
        session.title = message[:60] + ("…" if len(message) > 60 else "")

    db.session.commit()

    return jsonify({
        "reply": reply,
        "context_used": len(chunks),
        "message_id": assistant_msg.id,
    })


# ── Inline chat ──────────────────────────────────────────────────────────────

@chat_bp.route("/inline", methods=["POST"])
@login_required
def inline():
    data = request.get_json(silent=True) or {}
    selected_text = (data.get("selected_text") or "").strip()
    question = (data.get("question") or "").strip()

    if not selected_text or not question:
        return jsonify({"error": "selected_text and question are required"}), 400

    try:
        reply = llm_service.ask_inline(selected_text=selected_text, question=question)
    except Exception as e:
        current_app.logger.error(f"Inline LLM error: {e}")
        return jsonify({"error": f"LLM request failed: {str(e)}"}), 502

    return jsonify({"reply": reply})