SPB Git forge

spb/uqo-chat

Public
14commits 1branches 0releases
1.4 MBsize
maindefault branch
17 days agolast push
Python 64.6% TypeScript 33.7% CSS 0.8%
6.6 KB · 154 lines python
Raw Blame History
1"""Chat SSE endpoint, stop, regenerate."""23from __future__ import annotations45import asyncio6import json7from collections.abc import AsyncIterator8from typing import Any910from fastapi import APIRouter, Depends, HTTPException, Request, status11from fastapi.responses import StreamingResponse12from pydantic import BaseModel, Field1314from app.core.config import Settings, get_settings15from app.core.logging import get_logger16from app.core.ratelimit import MSG_MESSAGES, limiter17from app.core.security import AuthUser, get_current_user, hash_user_id18from app.llm.agent import Turn, generate_title19from app.services import analytics, users20from app.services import conversations as conv_service2122log = get_logger("api.chat")23router = APIRouter(tags=["chat"])2425ACTIVE: dict[str, Turn] = {}262728class SendReq(BaseModel):29    content: str = Field(..., min_length=1, max_length=12000)30    attachments: list[str] = Field(default_factory=list)31    deep: bool = False323334def _sse(event: str, data: dict[str, Any], event_id: int | None = None) -> str:35    head = f"id: {event_id}\n" if event_id is not None else ""36    return f"{head}event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"373839async def _stream(turn: Turn, user_text: str, attachments: list[str], conv_id: str,40                  first_message: bool, request: Request) -> AsyncIterator[str]:41    queue: asyncio.Queue[tuple[str, dict[str, Any]] | None] = asyncio.Queue()4243    async def emit(event: str, data: dict[str, Any]) -> None:44        await queue.put((event, data))4546    turn.emit = emit47    ACTIVE[conv_id] = turn4849    async def worker() -> None:50        try:51            await asyncio.wait_for(turn.run(user_text, attachments),52                                   timeout=turn.settings.LLM_TURN_TIMEOUT_SECONDS)53        except TimeoutError:54            await emit("error", {"code": "timeout", "message_fr": "Le tour a pris trop de temps."})55            await emit("done", {"message_id": turn.message_id, "finish_reason": "timeout"})56        except Exception as exc:  # noqa: BLE00157            log.exception("turn_failed")58            await emit("error", {"code": "internal", "message_fr": "Erreur interne."})59            await emit("done", {"message_id": turn.message_id, "finish_reason": "error"})60            _ = exc61        finally:62            if first_message:63                await generate_title(conv_id, user_text, emit)64            await queue.put(None)6566    task = asyncio.create_task(worker())67    seq = 068    try:69        while True:70            try:71                item = await asyncio.wait_for(queue.get(), timeout=15)72            except TimeoutError:73                yield ": ping\n\n"74                continue75            if item is None:76                break77            seq += 178            yield _sse(item[0], item[1], seq)79            if await request.is_disconnected():80                turn.cancel.set()81    finally:82        ACTIVE.pop(conv_id, None)83        if not task.done():84            turn.cancel.set()85            try:86                await asyncio.wait_for(task, timeout=5)87            except (TimeoutError, asyncio.CancelledError):88                task.cancel()899091@router.post("/chat/{conv_id}/messages")92async def send_message(conv_id: str, req: SendReq, request: Request,93                       auth: AuthUser = Depends(get_current_user),94                       settings: Settings = Depends(get_settings)) -> StreamingResponse:95    conv = await conv_service.get_conversation(conv_id, auth.id)96    if not conv:97        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Conversation introuvable.")98    if conv_id in ACTIVE:99        raise HTTPException(status.HTTP_409_CONFLICT, detail="Une réponse est déjà en cours.")100    user = await users.get_user(auth.id)101    if not user:102        raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail="Compte introuvable.")103    if user.role == "student":104        limiter.check(f"msg_h:{auth.id}", settings.RATE_MESSAGES_PER_HOUR, 3600, MSG_MESSAGES)105        limiter.check(f"msg_d:{auth.id}", settings.RATE_MESSAGES_PER_DAY, 86400, MSG_MESSAGES)106    prior = await conv_service.history(conv_id, limit=1)107    first = len(prior) == 0108    await conv_service.add_message(conv_id, "user", req.content, req.attachments)109    asyncio.create_task(analytics.classify_and_record(hash_user_id(auth.id), conv.course_code,110                                                      req.content))111    deep = req.deep or bool((user.preferences or {}).get("deep"))112    turn = Turn(conv, user, emit=None, deep=deep, settings=settings)  # type: ignore[arg-type]113    headers = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "keep-alive"}114    return StreamingResponse(_stream(turn, req.content, req.attachments, conv_id, first, request),115                             media_type="text/event-stream", headers=headers)116117118@router.post("/chat/{conv_id}/stop")119async def stop(conv_id: str, auth: AuthUser = Depends(get_current_user)) -> dict:120    turn = ACTIVE.get(conv_id)121    if turn and turn.user.id == auth.id:122        turn.cancel.set()123        return {"ok": True, "stopped": True}124    return {"ok": True, "stopped": False}125126127class RegenReq(BaseModel):128    deep: bool = False129130131@router.post("/chat/{conv_id}/messages/{message_id}/regenerate")132async def regenerate(conv_id: str, message_id: str, req: RegenReq, request: Request,133                     auth: AuthUser = Depends(get_current_user),134                     settings: Settings = Depends(get_settings)) -> StreamingResponse:135    conv = await conv_service.get_conversation(conv_id, auth.id)136    if not conv:137        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Conversation introuvable.")138    if conv_id in ACTIVE:139        raise HTTPException(status.HTTP_409_CONFLICT, detail="Une réponse est déjà en cours.")140    user = await users.get_user(auth.id)141    if not user:142        raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail="Compte introuvable.")143    await conv_service.delete_messages_after(conv_id, message_id)144    hist = await conv_service.history(conv_id, limit=1)145    if not hist or hist[-1].role != "user":146        raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="Aucun message à régénérer.")147    asyncio.create_task(analytics.record_event(hash_user_id(auth.id), conv.course_code,148                                               "regenerate"))149    turn = Turn(conv, user, emit=None, deep=req.deep, settings=settings)  # type: ignore[arg-type]150    headers = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}151    return StreamingResponse(_stream(turn, hist[-1].content, hist[-1].attachments or [], conv_id,152                                     False, request),153                             media_type="text/event-stream", headers=headers)154