pablan/backend/app/api/conversations/turns.py

"""One turn: a question in, an answer streamed out, both persisted.

This is the only place that knows about SSE. A mode yields `ModeEvent`s and
knows nothing about HTTP; here they become frames on the wire. Persistence is
deliberately asymmetric: the user message is committed BEFORE streaming starts
so it survives anything the endpoint does, while the assistant message is
written at the end — complete, partial after an abort, or source-list-only when
no model could answer.
"""

import asyncio
import logging
import time
import uuid
from collections.abc import AsyncIterator
from datetime import UTC, datetime
from typing import Annotated, Any

from fastapi import Depends
from fastapi.responses import StreamingResponse
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from app.api.conversations.access import own_conversation
from app.api.conversations.routing import conversations_router
from app.api.conversations.schemas import SendMessage
from app.api.sse import sse
from app.auth.deps import get_current_user
from app.db import get_db
from app.errors import ApiError
from app.llm.errors import LLMError
from app.log import conversation_id as conversation_id_var
from app.models import Conversation, Message, MessageRole, User
from app.modes import get_mode
from app.modes.base import Degraded, Done, Error, Mode, Sources, StateChanged, Token

router = conversations_router()
logger = logging.getLogger("pablan.conversations")


@router.post("/{conversation_id}/messages")
async def send_message(
    conversation_id: uuid.UUID,
    body: SendMessage,
    user: Annotated[User, Depends(get_current_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
) -> StreamingResponse:
    conversation = await own_conversation(db, conversation_id, user, with_messages=True)
    mode = get_mode(conversation.mode.value)
    if mode is None:
        raise ApiError(
            400, f"Mode '{conversation.mode.value}' is not available.", "unknown_mode"
        )

    # The user message is committed before streaming starts — it survives
    # whatever happens to the LLM call.
    db.add(
        Message(
            conversation_id=conversation.id,
            role=MessageRole.user,
            content=body.content,
        )
    )
    conversation.updated_at = datetime.now(UTC)
    await db.commit()

    return StreamingResponse(
        stream_turn(conversation, body.content, mode, db),
        media_type="text/event-stream",
        headers={"cache-control": "no-cache", "x-accel-buffering": "no"},
    )


def _source_payload(chunks: Any) -> list[dict[str, Any]]:
    """The wire shape of a citation — the same dicts are snapshotted onto the
    message, so a reload shows exactly what was streamed."""
    return [
        {
            "document_id": str(chunk.document_id),
            "title": chunk.title,
            "heading_path": chunk.heading_path,
            "excerpt": chunk.excerpt,
            "used": chunk.used,
            "review_pending": chunk.review_pending,
        }
        for chunk in chunks
    ]


async def stream_turn(
    conversation: Conversation, content: str, mode: Mode, db: AsyncSession
) -> AsyncIterator[str]:
    """Convert ModeEvents to SSE frames; persist the assistant reply —
    complete on normal end, partial on client abort or endpoint failure."""
    context_token = conversation_id_var.set(str(conversation.id))
    started = time.monotonic()
    parts: list[str] = []
    sources: list[dict[str, Any]] = []
    # Set when the mode gave up on the model: the turn still has a reply (the
    # retrieved documents), so it is persisted and replayed like any other.
    fallback: str | None = None
    outcome = "ok"
    try:
        try:
            async for event in mode.handle_turn(conversation, content, db):
                match event:
                    case Token(text=text):
                        parts.append(text)
                        yield sse("token", {"text": text})
                    case Sources(chunks=chunks):
                        sources = _source_payload(chunks)
                        yield sse("sources", {"chunks": sources})
                    case StateChanged(phase=phase, count=count):
                        yield sse("state", {"phase": phase, "count": count})
                    case Error(code=code):
                        outcome = "error"
                        yield sse("error", {"code": code})
                    case Degraded(code=code):
                        outcome = "degraded"
                        fallback = code
                        yield sse("fallback", {"code": code})
                    case Done():
                        pass  # the router emits the final done after persisting
        except LLMError as exc:
            # Every endpoint failure inside a mode ends the turn the same way,
            # wherever it happened. Retrieval embeds before the model is ever
            # called, so an escaping error would reach the browser as a
            # truncated stream ("connection lost") instead of the reason.
            outcome = "error"
            logger.warning(
                "turn failed",
                extra={
                    "event": "turn_error",
                    "mode": mode.name,
                    "code": exc.code,
                    "cause_type": exc.cause_type,
                    "status_code": exc.status_code,
                },
            )
            yield sse("error", {"code": exc.code})
        except (asyncio.CancelledError, GeneratorExit):
            # Client aborted (stop button): keep what was already streamed.
            outcome = "aborted"
            if parts:
                await asyncio.shield(
                    _persist_partial(db.bind, conversation.id, "".join(parts), sources)
                )
            raise
        # Whatever arrived before the end is the reply, complete or not — for a
        # fallback turn that is the source list alone.
        if parts or fallback:
            message_id = await _persist_assistant(
                db, conversation, "".join(parts), sources, fallback=fallback
            )
            yield sse("done", {"message_id": str(message_id)})
    finally:
        conversation_id_var.reset(context_token)
        logger.info(
            "turn finished",
            extra={
                "event": "turn",
                "mode": mode.name,
                "outcome": outcome,
                "duration_ms": round((time.monotonic() - started) * 1000),
                "token_events": len(parts),
            },
        )


def _assistant_meta(
    sources: list[dict[str, Any]], fallback: str | None
) -> dict[str, Any]:
    meta: dict[str, Any] = {}
    if sources:
        meta["sources"] = sources
    if fallback:
        # Why there is no generated text, kept so a reload replays the turn as
        # what it was rather than as an empty reply.
        meta["fallback"] = fallback
    return meta


async def _persist_assistant(
    db: AsyncSession,
    conversation: Conversation,
    content: str,
    sources: list[dict[str, Any]],
    *,
    fallback: str | None = None,
) -> uuid.UUID:
    message = Message(
        conversation_id=conversation.id,
        role=MessageRole.assistant,
        content=content,
        meta=_assistant_meta(sources, fallback),
    )
    db.add(message)
    conversation.updated_at = datetime.now(UTC)
    await db.commit()
    return message.id


async def _persist_partial(
    bind: Any,
    conversation_id: uuid.UUID,
    content: str,
    sources: list[dict[str, Any]],
) -> None:
    """Write what was streamed before the client hung up.

    On a FRESH session on the same engine as the request session: the request
    session is being torn down mid-cancel, so it cannot be used to commit, and
    binding to the same engine keeps this working under the test overrides.
    """
    async with async_sessionmaker(bind, expire_on_commit=False)() as db:
        db.add(
            Message(
                conversation_id=conversation_id,
                role=MessageRole.assistant,
                content=content,
                meta=_assistant_meta(sources, None),
            )
        )
        conversation = await db.get(Conversation, conversation_id)
        if conversation is not None:
            conversation.updated_at = datetime.now(UTC)
        await db.commit()