Spaces:
Sleeping
Sleeping
File size: 2,638 Bytes
f65e025 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """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"})
|