"""Chat SSE endpoint, stop, regenerate.""" from __future__ import annotations import asyncio import json from collections.abc import AsyncIterator from typing import Any from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from app.core.config import Settings, get_settings from app.core.logging import get_logger from app.core.ratelimit import MSG_MESSAGES, limiter from app.core.security import AuthUser, get_current_user, hash_user_id from app.llm.agent import Turn, generate_title from app.services import analytics, users from app.services import conversations as conv_service log = get_logger("api.chat") router = APIRouter(tags=["chat"]) ACTIVE: dict[str, Turn] = {} class SendReq(BaseModel): content: str = Field(..., min_length=1, max_length=12000) attachments: list[str] = Field(default_factory=list) deep: bool = False def _sse(event: str, data: dict[str, Any], event_id: int | None = None) -> str: head = f"id: {event_id}\n" if event_id is not None else "" return f"{head}event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n" async def _stream(turn: Turn, user_text: str, attachments: list[str], conv_id: str, first_message: bool, request: Request) -> AsyncIterator[str]: queue: asyncio.Queue[tuple[str, dict[str, Any]] | None] = asyncio.Queue() async def emit(event: str, data: dict[str, Any]) -> None: await queue.put((event, data)) turn.emit = emit ACTIVE[conv_id] = turn async def worker() -> None: try: await asyncio.wait_for(turn.run(user_text, attachments), timeout=turn.settings.LLM_TURN_TIMEOUT_SECONDS) except TimeoutError: await emit("error", {"code": "timeout", "message_fr": "Le tour a pris trop de temps."}) await emit("done", {"message_id": turn.message_id, "finish_reason": "timeout"}) except Exception as exc: # noqa: BLE001 log.exception("turn_failed") await emit("error", {"code": "internal", "message_fr": "Erreur interne."}) await emit("done", {"message_id": turn.message_id, "finish_reason": "error"}) _ = exc finally: if first_message: await generate_title(conv_id, user_text, emit) await queue.put(None) task = asyncio.create_task(worker()) seq = 0 try: while True: try: item = await asyncio.wait_for(queue.get(), timeout=15) except TimeoutError: yield ": ping\n\n" continue if item is None: break seq += 1 yield _sse(item[0], item[1], seq) if await request.is_disconnected(): turn.cancel.set() finally: ACTIVE.pop(conv_id, None) if not task.done(): turn.cancel.set() try: await asyncio.wait_for(task, timeout=5) except (TimeoutError, asyncio.CancelledError): task.cancel() @router.post("/chat/{conv_id}/messages") async def send_message(conv_id: str, req: SendReq, request: Request, auth: AuthUser = Depends(get_current_user), settings: Settings = Depends(get_settings)) -> StreamingResponse: conv = await conv_service.get_conversation(conv_id, auth.id) if not conv: raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Conversation introuvable.") if conv_id in ACTIVE: raise HTTPException(status.HTTP_409_CONFLICT, detail="Une réponse est déjà en cours.") user = await users.get_user(auth.id) if not user: raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail="Compte introuvable.") if user.role == "student": limiter.check(f"msg_h:{auth.id}", settings.RATE_MESSAGES_PER_HOUR, 3600, MSG_MESSAGES) limiter.check(f"msg_d:{auth.id}", settings.RATE_MESSAGES_PER_DAY, 86400, MSG_MESSAGES) prior = await conv_service.history(conv_id, limit=1) first = len(prior) == 0 await conv_service.add_message(conv_id, "user", req.content, req.attachments) asyncio.create_task(analytics.classify_and_record(hash_user_id(auth.id), conv.course_code, req.content)) deep = req.deep or bool((user.preferences or {}).get("deep")) turn = Turn(conv, user, emit=None, deep=deep, settings=settings) # type: ignore[arg-type] headers = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "keep-alive"} return StreamingResponse(_stream(turn, req.content, req.attachments, conv_id, first, request), media_type="text/event-stream", headers=headers) @router.post("/chat/{conv_id}/stop") async def stop(conv_id: str, auth: AuthUser = Depends(get_current_user)) -> dict: turn = ACTIVE.get(conv_id) if turn and turn.user.id == auth.id: turn.cancel.set() return {"ok": True, "stopped": True} return {"ok": True, "stopped": False} class RegenReq(BaseModel): deep: bool = False @router.post("/chat/{conv_id}/messages/{message_id}/regenerate") async def regenerate(conv_id: str, message_id: str, req: RegenReq, request: Request, auth: AuthUser = Depends(get_current_user), settings: Settings = Depends(get_settings)) -> StreamingResponse: conv = await conv_service.get_conversation(conv_id, auth.id) if not conv: raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Conversation introuvable.") if conv_id in ACTIVE: raise HTTPException(status.HTTP_409_CONFLICT, detail="Une réponse est déjà en cours.") user = await users.get_user(auth.id) if not user: raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail="Compte introuvable.") await conv_service.delete_messages_after(conv_id, message_id) hist = await conv_service.history(conv_id, limit=1) if not hist or hist[-1].role != "user": raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="Aucun message à régénérer.") asyncio.create_task(analytics.record_event(hash_user_id(auth.id), conv.course_code, "regenerate")) turn = Turn(conv, user, emit=None, deep=req.deep, settings=settings) # type: ignore[arg-type] headers = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no"} return StreamingResponse(_stream(turn, hist[-1].content, hist[-1].attachments or [], conv_id, False, request), media_type="text/event-stream", headers=headers)