src/api/v2/chat.py

"""Chat endpoint — v2 pilot (pr/5 Phase 2).

`POST /api/v2/chat/stream` is the v2 of the only FE→Python call. It is identical to
`POST /api/v1/chat/stream` except:
  - the request carries an explicit **`analysis_id`** (replacing v1's `room_id`). The
    two are the same session id today (`analysis_id == room_id`), so the warm,
    process-shared `ChatHandler` and the v1 cache/history helpers are reused unchanged.
  - the `done` event carries the assistant **`message_id`** (minted Python-side if Go
    omits it — contract open-Q #1), so the FE can fetch `/api/v1/observability` for the
    turn in parallel with the stream.

Only chat moves to v2; the tools group + observability stay on `/api/v1` (contract:
API_ENDPOINTS_RESTRUCTURE.md §1).

⚠️ Persistence (transitional). This mirrors v1: it still load/saves turn history via the
analysis-keyed message tables so multi-turn context works in the playground. Moving the
read/write to Go-owned `analyses_messages` (and making Python read-only) is DEV_PLAN #25.
Note Sofhia's `/tools/help` is already generative-only — align chat with that under #25.
"""

import json
import uuid
from typing import Any

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession
from sse_starlette.sse import EventSourceResponse

# Reuse the v1 chat machinery verbatim (warm ChatHandler + cache/history helpers) so
# v2 stays a thin field-rename over the same logic. Importing the module-private helpers
# is the established pattern here (handlers/help.py imports `_chat_handler` the same way).
from src.api.v1.chat import (
    _CACHEABLE_INTENTS,
    _chat_cache_key,
    _chat_handler,
    _fast_intent,
    cache_response,
    get_cached_response,
    load_history,
    save_messages,
)
from src.db.postgres.connection import get_db
from src.db.redis.connection import get_redis
from src.middlewares.logging import get_logger, log_execution

logger = get_logger("chat_api_v2")

router = APIRouter(prefix="/api/v2", tags=["Chat"])


def _mint_message_id(provided: str | None) -> str:
    """Use Go's assistant turn id when provided; else mint one (contract open-Q #1)."""
    return provided or f"msg_{uuid.uuid4().hex[:12]}"


class ChatRequest(BaseModel):
    user_id: str
    analysis_id: str
    message: str
    # ⚠️ open-Q #1: Go may mint + pass the assistant turn id; if absent we mint one and
    # echo it on `done` so the FE can correlate /observability with this answer.
    message_id: str | None = None


@router.post("/chat/stream")
@log_execution(logger)
async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
    """Chat endpoint with streaming response (v2 — keyed on `analysis_id`).

    SSE event sequence:
      1. sources  — JSON array of source refs (table for structured; deduped
                    document_id/page_label for unstructured; [] for chat/help/error)
      2. status   — slow-path progress pings (optional)
      3. chunk    — text fragments of the answer
      4. done     — {"message_id": "..."} for the observability lookup
    """
    analysis_id = request.analysis_id
    message_id = _mint_message_id(request.message_id)
    redis = await get_redis()
    cache_key = _chat_cache_key(analysis_id, request.user_id, request.message)

    # v2 `done` always carries the turn id (v1 sent an empty `done`).
    done_event = {"event": "done", "data": json.dumps({"message_id": message_id})}

    # Redis cache hit (stateless `chat` intent only).
    cached = await get_cached_response(redis, cache_key)
    logger.info("cache check", cache_key=cache_key, cache_hit=cached is not None)
    if cached:
        logger.info("Returning cached response")
        cached_text = cached["response"]
        cached_sources = cached["sources"]
        await save_messages(db, analysis_id, request.message, cached_text, sources=cached_sources)

        async def stream_cached():
            yield {"event": "sources", "data": json.dumps(cached_sources)}
            for i in range(0, len(cached_text), 50):
                yield {"event": "chunk", "data": cached_text[i:i + 50]}
            yield done_event

        return EventSourceResponse(stream_cached())

    try:
        # Fast intent: greetings/farewells bypass the LLM entirely.
        direct = _fast_intent(request.message)
        if direct:
            await cache_response(redis, cache_key, direct, sources=[])
            await save_messages(db, analysis_id, request.message, direct, sources=[])

            async def stream_direct():
                yield {"event": "sources", "data": json.dumps([])}
                yield {"event": "chunk", "data": direct}
                yield done_event

            return EventSourceResponse(stream_direct())

        history = await load_history(db, analysis_id, limit=10)
        handler = _chat_handler

        async def stream_response():
            logger.info("stream_response started", analysis_id=analysis_id, user_id=request.user_id)
            full_response = ""
            sources: list[dict[str, Any]] = []
            effective_intent: str | None = None
            async for event in handler.handle(
                request.message, request.user_id, history, analysis_id=analysis_id
            ):
                if event["event"] == "intent":
                    # consumed internally (not forwarded); gates caching below.
                    try:
                        effective_intent = json.loads(event["data"]).get("intent")
                    except (TypeError, ValueError, AttributeError):
                        effective_intent = None
                elif event["event"] == "sources":
                    try:
                        sources = json.loads(event["data"]) or []
                    except (TypeError, ValueError):
                        sources = []
                    yield event
                elif event["event"] == "chunk":
                    full_response += event["data"]
                    yield event
                elif event["event"] == "done":
                    # Only cache stateless `chat` replies (see _CACHEABLE_INTENTS).
                    if effective_intent in _CACHEABLE_INTENTS:
                        await cache_response(redis, cache_key, full_response, sources=sources)
                    try:
                        await save_messages(
                            db, analysis_id, request.message, full_response, sources=sources
                        )
                    except Exception as e:
                        logger.error("save_messages failed", analysis_id=analysis_id, error=str(e))
                    yield done_event
                elif event["event"] == "status":
                    # slow-path progress: forward so the client shows activity.
                    yield event
                elif event["event"] == "error":
                    yield event
                    return

        return EventSourceResponse(stream_response())

    except Exception as e:
        logger.error("Chat failed", error=str(e))
        raise HTTPException(status_code=500, detail=f"Chat failed: {str(e)}") from e