app.routes.chat

  1"""Chat route handlers for the AIFT Flask application.
  2
  3Handles interactive chat sessions about completed forensic analysis results,
  4including message submission, SSE streaming, and history management.
  5
  6Attributes:
  7    chat_bp: Flask Blueprint for chat routes.
  8"""
  9
 10from __future__ import annotations
 11
 12import copy
 13import threading
 14
 15from flask import Blueprint, Response, current_app, request
 16
 17from ..chat import ChatManager
 18
 19from .state import (
 20    STATE_LOCK,
 21    CHAT_PROGRESS,
 22    error_response,
 23    success_response,
 24    get_case,
 25    new_progress,
 26    stream_sse,
 27)
 28from .tasks import (
 29    run_task_with_case_log_context,
 30    run_chat,
 31    load_case_analysis_results,
 32)
 33
 34__all__ = ["chat_bp"]
 35
 36chat_bp = Blueprint("chat", __name__)
 37
 38
 39@chat_bp.post("/api/cases/<case_id>/chat")
 40def chat_with_case(case_id: str) -> Response | tuple[Response, int]:
 41    """Initiate a chat interaction about completed analysis results.
 42
 43    Args:
 44        case_id: UUID of the case.
 45
 46    Returns:
 47        ``(Response, 202)`` confirming start, or error.
 48    """
 49    case = get_case(case_id)
 50    if case is None:
 51        return error_response(f"Case not found: {case_id}", 404)
 52
 53    payload = request.get_json(silent=True)
 54    if not isinstance(payload, dict):
 55        return error_response("Chat payload must be a JSON object.", 400)
 56
 57    message = str(payload.get("message", "")).strip()
 58    if not message:
 59        return error_response("`message` is required.", 400)
 60
 61    with STATE_LOCK:
 62        case_snapshot_for_check = dict(case)
 63    if not load_case_analysis_results(case_snapshot_for_check):
 64        return error_response("No analysis results available for this case. Run analysis first.", 400)
 65
 66    config = current_app.config.get("AIFT_CONFIG", {})
 67    if not isinstance(config, dict):
 68        return error_response("Invalid in-memory configuration state.", 500)
 69
 70    with STATE_LOCK:
 71        chat_state = CHAT_PROGRESS.setdefault(case_id, new_progress())
 72        if chat_state.get("status") == "running":
 73            return error_response("Chat is already running for this case.", 409)
 74        CHAT_PROGRESS[case_id] = new_progress(status="running")
 75
 76    config_snapshot = copy.deepcopy(config)
 77    threading.Thread(
 78        target=run_task_with_case_log_context,
 79        args=(case_id, run_chat, case_id, message, config_snapshot),
 80        daemon=True,
 81    ).start()
 82    return success_response({"status": "processing"}, 202)
 83
 84
 85@chat_bp.get("/api/cases/<case_id>/chat/stream")
 86def stream_chat_progress(case_id: str) -> Response | tuple[Response, int]:
 87    """Stream chat response tokens via SSE.
 88
 89    Args:
 90        case_id: UUID of the case.
 91
 92    Returns:
 93        SSE Response, or 404 error.
 94    """
 95    if get_case(case_id) is None:
 96        return error_response(f"Case not found: {case_id}", 404)
 97    return stream_sse(CHAT_PROGRESS, case_id)
 98
 99
100@chat_bp.get("/api/cases/<case_id>/chat/history")
101def get_case_chat_history(case_id: str) -> Response | tuple[Response, int]:
102    """Retrieve the full chat message history for a case.
103
104    Args:
105        case_id: UUID of the case.
106
107    Returns:
108        JSON response with chat messages, or 404 error.
109    """
110    case = get_case(case_id)
111    if case is None:
112        return error_response(f"Case not found: {case_id}", 404)
113    with STATE_LOCK:
114        case_dir = case["case_dir"]
115    manager = ChatManager(case_dir)
116    return success_response({"messages": manager.get_history()})
117
118
119@chat_bp.delete("/api/cases/<case_id>/chat/history")
120def clear_case_chat_history(case_id: str) -> Response | tuple[Response, int]:
121    """Clear the chat history for a case.
122
123    Args:
124        case_id: UUID of the case.
125
126    Returns:
127        JSON confirmation, or 404 error.
128    """
129    case = get_case(case_id)
130    if case is None:
131        return error_response(f"Case not found: {case_id}", 404)
132    with STATE_LOCK:
133        case_dir = case["case_dir"]
134        audit_logger = case["audit"]
135    manager = ChatManager(case_dir)
136    manager.clear()
137    audit_logger.log("chat_history_cleared", {"case_id": case_id})
138    return success_response({"status": "cleared", "case_id": case_id})