src/ai_agent/api/routers/chat.py

"""Chat router with SSE streaming.

Event protocol (named events; payloads are JSON):

  ``session``        ``{"session_id": "..."}``                            once, first
  ``user_message``   ``{"text": "...", "asset_ids": [...]}``              echo
  ``status``         ``{"phase": "thinking" | "agent_running" | ...}``    occasional
  ``text``           ``{"content": "..."}``                                main reply
  ``recommendation`` ``{"item": {rank, name, accuracy, why, ...}}``       per card
  ``tool_trace``     ``{"trace": {...}}``                                  per call
  ``pending_action`` ``{"action": {type, tool_name, ...}}``                if pause
  ``clarification``  ``{"question": "...", "options": [...]}``             if needed
  ``images``         ``{"paths": [...]}``                                  tool-exec
  ``files``          ``{"items": [{path, label}, ...]}``                   tool-exec
  ``usage``          ``{"total": n, "input": n, "output": n}``             once
  ``error``          ``{"message": "...", "code": "..."}``                 fatal
  ``done``           ``{}``                                                terminal

The underlying agent (``services.chat.process_turn``) is still synchronous;
this router translates the final ``ChatTurnResult`` into a sequence of
events. As soon as pydantic-ai's ``run_stream`` is wired in, the same
endpoint will emit ``text`` deltas live.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, AsyncIterator, Dict, List

from fastapi import APIRouter, Depends, HTTPException, status
from sse_starlette.sse import EventSourceResponse

from ai_agent.api.deps import get_doc_index, require_auth
from ai_agent.api.schemas import ChatStartBody
from ai_agent.retriever.software_doc import SoftwareDoc
from ai_agent.services.chat import (
    ChatRequest,
    ChatTurnResult,
    approve_pending,
    decline_pending,
    process_turn,
)
from ai_agent.services.sessions import Session, get_session_store

log = logging.getLogger("api.routers.chat")

router = APIRouter(prefix="/api/chat", tags=["chat"], dependencies=[Depends(require_auth)])


def _event(name: str, data: Any) -> Dict[str, Any]:
    return {"event": name, "data": json.dumps(data, default=str)}


# Status messages emitted on a heartbeat while the (synchronous) agent runs.
# The list is intentionally aligned with the typical agent timeline:
# pre-flight -> retrieval -> rerank -> VLM call -> shaping the response.
_STATUS_BEATS: List[str] = [
    "analyzing your request…",
    "searching software catalog…",
    "ranking candidates against your image…",
    "consulting the agent — this can take a few seconds…",
    "fetching repo metadata…",
    "preparing recommendations…",
    "still working — large models are slower on first call…",
]
_STATUS_HEARTBEAT_SECONDS = 1.6


def _stream_result(
    session: Session, result: ChatTurnResult
) -> AsyncIterator[Dict[str, Any]]:
    """Translate a finished ChatTurnResult into a sequence of SSE events."""

    async def gen() -> AsyncIterator[Dict[str, Any]]:
        # Always lead with the session id so the client can persist it.
        yield _event("session", {"session_id": session.session_id})

        if result.text:
            yield _event("text", {"content": result.text})

        for rec in result.recommendations:
            yield _event(
                "recommendation",
                {
                    "rank": rec.rank,
                    "name": rec.name,
                    "accuracy": rec.accuracy,
                    "why": rec.why,
                    "doc": rec.doc,
                    "demo_url": rec.demo_url,
                },
            )

        for trace in result.tool_traces:
            yield _event("tool_trace", {"trace": trace})

        if result.clarification:
            yield _event(
                "clarification",
                {
                    "question": result.clarification.question,
                    "context": result.clarification.context,
                    "options": result.clarification.options,
                },
            )

        if result.pending_action:
            pa = result.pending_action
            yield _event(
                "pending_action",
                {
                    "type": pa.type,
                    "tool_name": pa.tool_name,
                    "display_name": pa.display_name,
                    "icon": pa.icon,
                    "image_name": pa.image_name,
                    "demo_url": pa.demo_url,
                    "prompt": pa.prompt,
                },
            )

        if result.images:
            yield _event("images", {"paths": result.images})
        if result.files:
            yield _event(
                "files",
                {"items": [{"path": p, "label": label} for p, label in result.files]},
            )

        if result.usage:
            yield _event("usage", result.usage)

        if result.status == "error" and result.error:
            yield _event("error", {"message": result.text, "code": result.error})

        yield _event("done", {"status": result.status})

    return gen()


def _resolve_or_error(session_id: str) -> Session:
    session = get_session_store().get(session_id)
    if session is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="session_not_found"
        )
    return session


def _with_heartbeat(
    session: Session,
    coro_factory,
) -> AsyncIterator[Dict[str, Any]]:
    """Run a blocking callable in a thread; emit ``status`` heartbeats while
    it executes; then stream the real result.

    ``coro_factory`` is a zero-arg lambda that returns the ``ChatTurnResult``
    when called from a worker thread.
    """

    async def gen() -> AsyncIterator[Dict[str, Any]]:
        yield _event("session", {"session_id": session.session_id})
        yield _event("status", {"phase": "started", "message": _STATUS_BEATS[0]})

        task = asyncio.create_task(asyncio.to_thread(coro_factory))
        i = 1
        while True:
            try:
                result = await asyncio.wait_for(
                    asyncio.shield(task), timeout=_STATUS_HEARTBEAT_SECONDS
                )
                break
            except asyncio.TimeoutError:
                msg = _STATUS_BEATS[min(i, len(_STATUS_BEATS) - 1)]
                yield _event(
                    "status",
                    {"phase": f"beat_{i}", "message": msg},
                )
                i += 1

        # Drain the (now-finished) result through the existing translator,
        # but skip its leading ``session`` event since we already emitted one.
        translator = _stream_result(session, result)
        first = True
        async for ev in translator:
            if first and ev.get("event") == "session":
                first = False
                continue
            first = False
            yield ev

    return gen()


@router.post("")
async def start_chat(
    body: ChatStartBody,
    doc_index: Dict[str, SoftwareDoc] = Depends(get_doc_index),
):
    store = get_session_store()
    session = store.get_or_create(body.session_id)

    # Restore prior conversation transcript into a fresh server-side session.
    if body.seed_history and not session.conversation_history:
        session.conversation_history.extend(body.seed_history)

    request = ChatRequest(
        message=body.message,
        asset_ids=list(body.asset_ids),
        model=body.model,
        top_k=body.top_k,
        num_choices=body.num_choices,
    )
    return EventSourceResponse(
        _with_heartbeat(session, lambda: process_turn(session, request, doc_index))
    )


@router.post("/{session_id}/approve")
async def approve(
    session_id: str,
    doc_index: Dict[str, SoftwareDoc] = Depends(get_doc_index),
):
    session = _resolve_or_error(session_id)
    return EventSourceResponse(
        _with_heartbeat(session, lambda: approve_pending(session))
    )


@router.post("/{session_id}/decline")
async def decline(session_id: str):
    session = _resolve_or_error(session_id)
    result = decline_pending(session)
    return EventSourceResponse(_stream_result(session, result))


@router.post("/{session_id}/confirm-demo")
async def confirm_demo(
    session_id: str,
    doc_index: Dict[str, SoftwareDoc] = Depends(get_doc_index),
):
    """Resume a turn that asked the user to confirm running a demo URL."""
    session = _resolve_or_error(session_id)
    request = ChatRequest(message="yes")
    return EventSourceResponse(
        _with_heartbeat(session, lambda: process_turn(session, request, doc_index))
    )