backend/routers/chat.py

"""SSE streaming chat endpoint."""

import asyncio
import json
from typing import List, Literal, Optional

from fastapi import APIRouter, Depends
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from utils.response_cleaner import clean_model_response

router = APIRouter(prefix="/api")


class ChatRequest(BaseModel):
    query: str
    history: Optional[List[dict]] = []
    session_id: Optional[str] = None
    document_ids: Optional[List[str]] = None
    mode: Literal["all", "single"] = "all"


def get_qa_service():
    from core.bootstrap import get_qa_service as _get
    return _get()


def get_chat_store():
    from core.bootstrap import get_chat_store as _get
    return _get()


@router.post("/chat")
async def chat(request: ChatRequest, qa=Depends(get_qa_service), store=Depends(get_chat_store)):
    async def event_generator():
        full_answer = []
        sources = []
        confidence = None

        def run_qa():
            return list(
                qa.stream_answer(
                    request.query,
                    request.history,
                    document_ids=request.document_ids,
                    mode=request.mode,
                )
            )

        items = await asyncio.to_thread(run_qa)

        for item in items:
            if isinstance(item, dict):
                if item.get("type") == "sources":
                    sources = item["data"]
                    yield f"data: {json.dumps({'type': 'sources', 'data': sources})}\n\n"
                elif item.get("type") == "confidence":
                    confidence = item["data"]
                    yield f"data: {json.dumps({'type': 'confidence', 'data': confidence})}\n\n"
            else:
                full_answer.append(item)
                yield f"data: {json.dumps({'type': 'token', 'data': item})}\n\n"

        if request.session_id and full_answer:
            answer = clean_model_response("".join(full_answer))
            store.add_message(request.session_id, "user", request.query)
            store.add_message(
                request.session_id,
                "assistant",
                answer,
                sources=sources,
                confidence=confidence,
            )

        yield "data: [DONE]\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")