205 lines
7.6 KiB
Python
205 lines
7.6 KiB
Python
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])
|
|
|
|
|
|
# ── Diagnostics ─────────────────────────────────────────────────────────────
|
|
|
|
@chat_bp.route("/ping-llm", methods=["GET"])
|
|
@login_required
|
|
def ping_llm():
|
|
"""Test LM Studio connectivity. Returns config + a short completion."""
|
|
url = current_app.config.get("LM_STUDIO_URL", "")
|
|
model = current_app.config.get("LM_STUDIO_MODEL", "")
|
|
embed_model = current_app.config.get("LM_STUDIO_EMBEDDING_MODEL", "")
|
|
try:
|
|
reply = llm_service.ask(user_message="Reply with exactly: OK", context_chunks=[], history=[])
|
|
return jsonify({"status": "ok", "reply": reply, "url": url, "model": model, "embed_model": embed_model})
|
|
except Exception as e:
|
|
return jsonify({"status": "error", "error": str(e), "url": url, "model": model, "embed_model": embed_model}), 502
|
|
|
|
|
|
@chat_bp.route("/debug-rag", methods=["GET"])
|
|
@login_required
|
|
def debug_rag():
|
|
"""Diagnose RAG: show collection size, run a test query, return raw distances."""
|
|
query = request.args.get("q", "test")
|
|
source_id = request.args.get("source_id", type=int)
|
|
source_type = request.args.get("source_type", "doc")
|
|
try:
|
|
collection = rag_service._get_collection()
|
|
total_chunks = collection.count()
|
|
source_ids = [source_id] if source_id else None
|
|
raw = collection.query(
|
|
query_texts=[query],
|
|
n_results=min(5, max(1, total_chunks)),
|
|
where=rag_service._build_where(current_user.id, source_ids, source_type if source_id else None),
|
|
include=["documents", "distances", "ids"],
|
|
)
|
|
hits = [
|
|
{"id": i, "distance": round(d, 4), "preview": t[:120]}
|
|
for i, d, t in zip(
|
|
(raw.get("ids") or [[]])[0],
|
|
(raw.get("distances") or [[]])[0],
|
|
(raw.get("documents") or [[]])[0],
|
|
)
|
|
]
|
|
return jsonify({"total_chunks": total_chunks, "query": query, "hits": hits})
|
|
except Exception as e:
|
|
return jsonify({"error": str(e)}), 500
|
|
|
|
|
|
# ── 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 — failures are non-fatal (chat continues without context)
|
|
chunks = []
|
|
try:
|
|
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 not context_refs:
|
|
chunks = rag_service.similarity_search(
|
|
query=message,
|
|
user_id=current_user.id,
|
|
top_k=current_app.config["RAG_TOP_K"],
|
|
)
|
|
except Exception as e:
|
|
current_app.logger.warning(f"RAG lookup failed, continuing without context: {e}")
|
|
|
|
current_app.logger.info(f"RAG: found {len(chunks)} chunks for session {session_id}")
|
|
|
|
# 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})
|