backend/app/api/chat.py

"""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"})