"""Chat endpoint with Server-Sent Events streaming. Streams the agent's answer token-by-token, then a final `citations` event so the frontend can highlight grounded passages on the document. """ from __future__ import annotations import json from fastapi import APIRouter, HTTPException from fastapi.responses import StreamingResponse from app.agent.orchestrator import run_chat from app.core.logging import get_logger from app.schemas.documents import ChatRequest from app.services import storage log = get_logger(__name__) router = APIRouter(prefix="/api/chat", tags=["chat"]) def _sse(event: str, data) -> str: return f"event: {event}\ndata: {json.dumps(data, default=str)}\n\n" def _doc_context(document_id: str | None) -> str: if not document_id: return "" d = storage.get(document_id) if not d: return "" parts = [f"Active document: '{d.filename}' (id={d.id}), status={d.status.value}, " f"{d.num_pages} pages."] if d.classification: parts.append(f"Detected type: {d.classification.doc_type} " f"({d.classification.confidence:.0%}).") if d.extraction and d.extraction.fields: kv = ", ".join(f"{f.name}={f.value}" for f in d.extraction.fields[:12]) parts.append(f"Already-extracted fields: {kv}") return " ".join(parts) @router.post("") async def chat(req: ChatRequest): if not req.messages: raise HTTPException(400, "messages cannot be empty") async def gen(): try: stream, ctx = await run_chat( history=req.messages, document_id=req.document_id, doc_context=_doc_context(req.document_id), provider=req.provider, ) yield _sse("start", {"ok": True}) async for token in stream: yield _sse("token", {"text": token}) # de-dup citations by chunk_id seen, cites = set(), [] for c in ctx.citations: key = c.chunk_id or c.text[:40] if key in seen: continue seen.add(key) cites.append(c.model_dump()) yield _sse("citations", {"citations": cites}) yield _sse("done", {"ok": True}) except Exception as e: # pragma: no cover log.exception("chat failed") yield _sse("error", {"message": str(e)}) return StreamingResponse(gen(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})