src/agents/chat_handler.py

"""ChatHandler — top-level Phase 2 chat orchestrator.

End-to-end flow per user message:

  1. `IntentRouter.classify` → `chat` / `unstructured` / `structured`.
  2. Route:
       - `chat`         → no context. Pass straight to ChatbotAgent.
       - `structured`   → CatalogReader → QueryService → QueryResult.
       - `unstructured` → DocumentRetriever (placeholder, raises until TAB
                          ships) → list[DocumentChunk].
  3. `ChatbotAgent.astream` → yield text tokens.
  4. Wrap each step into an SSE-style event dict so the API endpoint can
     stream them as Server-Sent Events.

Phase 1's chat endpoint (`src/api/v1/chat.py`) is intentionally NOT touched
in this PR. PR7 cleanup will rewire it to call `ChatHandler.handle(...)`.

All dependencies are injectable for tests. Default constructors lazy-build
production deps (no `Settings()` triggered at import time as long as you
inject mocks).
"""

from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator, Callable
from typing import TYPE_CHECKING, Any

from langchain_core.messages import BaseMessage

from src.middlewares.logging import get_logger
from src.retrieval.base import RetrievalResult

from .chatbot import ChatbotAgent, DocumentChunk
from .orchestration import OrchestratorAgent

if TYPE_CHECKING:
    from ..catalog.reader import CatalogReader
    from ..query.service import QueryService
    from ..retrieval.router import RetrievalRouter
    from .slow_path.coordinator import SlowPathCoordinator
    from .slow_path.store import AnalysisStore

logger = get_logger("chat_handler")


class ChatHandler:
    """Top-level chat orchestrator.

    Returns an `AsyncIterator[dict]` of SSE-style events with shape
    `{"event": <name>, "data": <str>}`. Event types:
      - `intent`  — emitted once after classification (JSON-encoded decision)
      - `sources` — JSON array of source refs (one per structured table, or
                    per (document_id, page_label) for unstructured)
      - `chunk`   — text fragment of the streaming answer (one per token)
      - `done`    — end of stream (data is empty string)
      - `error`   — failure; data is a user-facing message
    """

    def __init__(
        self,
        intent_router: OrchestratorAgent | None = None,
        answer_agent: ChatbotAgent | None = None,
        catalog_reader: CatalogReader | None = None,
        query_service: QueryService | None = None,
        document_retriever: RetrievalRouter | None = None,
        *,
        enable_slow_path: bool = False,
        slow_path_coordinator_factory: (
            Callable[[str], SlowPathCoordinator] | None
        ) = None,
        analysis_store: AnalysisStore | None = None,
        enable_tracing: bool = False,
    ) -> None:
        self._intent_router = intent_router
        self._answer_agent = answer_agent
        self._catalog_reader = catalog_reader
        self._query_service = query_service
        self._document_retriever = document_retriever
        # Langfuse tracing (tokens + latency). OFF by default so tests never hit
        # Langfuse; the live endpoint opts in with ChatHandler(enable_tracing=True).
        self._enable_tracing = enable_tracing
        # Slow analytical path (Planner -> TaskRunner -> Assembler). OFF by default:
        # gated until the lead's real BusinessContext lands. When True, `structured`
        # intents route here instead of the single-query QueryService path. The
        # factory + store are injectable for tests.
        self._enable_slow_path = enable_slow_path
        self._slow_path_factory = slow_path_coordinator_factory
        self._analysis_store = analysis_store

    # ------------------------------------------------------------------
    # Lazy default-dep builders
    # ------------------------------------------------------------------

    def _get_intent_router(self) -> OrchestratorAgent:
        if self._intent_router is None:
            self._intent_router = OrchestratorAgent()
        return self._intent_router

    def _get_answer_agent(self) -> ChatbotAgent:
        if self._answer_agent is None:
            self._answer_agent = ChatbotAgent()
        return self._answer_agent

    def _get_catalog_reader(self) -> CatalogReader:
        if self._catalog_reader is None:
            from ..catalog.reader import CatalogReader
            from ..catalog.store import CatalogStore

            self._catalog_reader = CatalogReader(CatalogStore())
        return self._catalog_reader

    def _get_query_service(self) -> QueryService:
        if self._query_service is None:
            from ..query.service import QueryService

            self._query_service = QueryService()
        return self._query_service

    def _get_document_retriever(self) -> RetrievalRouter:
        if self._document_retriever is None:
            from ..retrieval.router import RetrievalRouter

            self._document_retriever = RetrievalRouter()
        return self._document_retriever

    # ------------------------------------------------------------------
    # Public entry
    # ------------------------------------------------------------------

    async def handle(
        self,
        message: str,
        user_id: str,
        history: list[BaseMessage] | None = None,
    ) -> AsyncIterator[dict[str, Any]]:
        tracer = self._make_tracer(user_id, message)

        # ---- 1. Classify intent --------------------------------------
        try:
            oc = tracer.callbacks()  # orchestrator: PII-safe, full capture
            ckw = {"callbacks": oc} if oc else {}
            decision = await self._get_intent_router().classify(message, history, **ckw)
        except Exception as e:
            logger.error("intent classification failed", error=str(e))
            yield {"event": "error", "data": f"Could not classify message: {e}"}
            return

        yield {"event": "intent", "data": decision.model_dump_json()}

        rewritten = decision.rewritten_query or message
        query_result = None
        chunks: list[DocumentChunk] | None = None
        raw_chunks: Any = None

        # ---- 2. Route ------------------------------------------------
        if decision.source_hint == "structured":
            try:
                # One memoizing reader per request: the same catalog is otherwise
                # re-fetched from the catalog DB 4-5x across the slow-path run. This
                # collapses those to one round-trip per source_hint and pins a single
                # consistent snapshot for plan + execution.
                from ..catalog.reader import MemoizingCatalogReader

                req_reader = MemoizingCatalogReader(self._get_catalog_reader())
                catalog = await req_reader.read(user_id, "structured")
                if self._enable_slow_path:
                    async for event in self._run_slow_path(
                        user_id, rewritten, catalog, tracer, req_reader
                    ):
                        yield event
                    return
                query_result = await self._get_query_service().run(
                    user_id, rewritten, catalog
                )
            except Exception as e:
                logger.error(
                    "structured route failed",
                    user_id=user_id,
                    error=str(e),
                )
                yield {"event": "error", "data": f"Structured query failed: {e}"}
                return
        elif decision.source_hint == "unstructured":
            try:
                raw_chunks = await self._get_document_retriever().retrieve(
                    rewritten, user_id
                )
                chunks = _normalize_chunks(raw_chunks)
            except NotImplementedError:
                logger.warning("DocumentRetriever placeholder hit", user_id=user_id)
                yield {
                    "event": "error",
                    "data": "Document retrieval is not yet available — pending implementation.",
                }
                return
            except Exception as e:
                logger.error(
                    "unstructured route failed", user_id=user_id, error=str(e)
                )
                yield {"event": "error", "data": f"Document retrieval failed: {e}"}
                return
        # else: chat path — no context

        # ---- 2b. Emit sources ---------------------------------------
        sources = _build_sources(
            decision.source_hint, user_id, query_result, raw_chunks
        )
        logger.info("built sources", source_hint=decision.source_hint, sources_count=len(sources), raw_chunks_count=len(raw_chunks) if raw_chunks else 0)
        yield {"event": "sources", "data": json.dumps(sources)}

        # ---- 3. Stream answer ----------------------------------------
        # masked: the answer call sees real query rows / doc chunks (possible PII).
        mc = tracer.callbacks(masked=True)
        akw = {"callbacks": mc} if mc else {}
        try:
            async for token in self._get_answer_agent().astream(
                message,
                history=history,
                query_result=query_result,
                chunks=chunks,
                **akw,
            ):
                yield {"event": "chunk", "data": token}
        except Exception as e:
            logger.error("answer streaming failed", user_id=user_id, error=str(e))
            yield {"event": "error", "data": f"Answer generation failed: {e}"}
            return

        tracer.end()
        yield {"event": "done", "data": ""}

    # ------------------------------------------------------------------
    # Slow analytical path (gated, off by default)
    # ------------------------------------------------------------------

    def _make_tracer(self, user_id: str, question: str) -> Any:
        """One Langfuse trace per request (or a NullTracer when disabled)."""
        if not self._enable_tracing:
            from ..observability.langfuse.tracing import NullTracer

            return NullTracer()
        from ..observability.langfuse.tracing import RequestTracer

        return RequestTracer.start(user_id=user_id, question=question)

    def _get_slow_path_coordinator(
        self, user_id: str, tracer: Any = None, catalog_reader: CatalogReader | None = None
    ) -> SlowPathCoordinator:
        """Build the per-request slow-path coordinator (composition root).

        The data-access tools need the authenticated `user_id` + `CatalogReader`,
        so the `CompositeToolInvoker` is constructed per request. The slow-path
        agent code stays tool-agnostic (INV-7) — only here, the composition root,
        do we name concrete tool implementations. When tracing is active the invoker
        is wrapped so each tool call records a metadata-only span.
        """
        if self._slow_path_factory is not None:
            return self._slow_path_factory(user_id)

        from ..tools.data_access import DataAccessToolInvoker
        from ..tools.invoker import AnalyticsToolInvoker, CompositeToolInvoker
        from .planner.registry import default_registry
        from .planner.service import PlannerService
        from .slow_path.assembler import Assembler
        from .slow_path.coordinator import SlowPathCoordinator
        from .slow_path.task_runner import TaskRunner

        invoker: Any = CompositeToolInvoker(
            DataAccessToolInvoker(user_id, catalog_reader or self._get_catalog_reader()),
            AnalyticsToolInvoker(),
        )
        if tracer is not None and getattr(tracer, "active", False):
            from ..observability.langfuse.tracing import TracingToolInvoker

            invoker = TracingToolInvoker(invoker, tracer)
        registry = default_registry()
        return SlowPathCoordinator(
            PlannerService(), TaskRunner(invoker, registry), Assembler(), registry
        )

    def _get_analysis_store(self) -> AnalysisStore:
        if self._analysis_store is None:
            from .slow_path.store import NullAnalysisStore

            self._analysis_store = NullAnalysisStore()
        return self._analysis_store

    async def _run_slow_path(
        self,
        user_id: str,
        query: str,
        catalog: Any,
        tracer: Any = None,
        catalog_reader: CatalogReader | None = None,
    ) -> AsyncIterator[dict[str, Any]]:
        """Run the slow path and stream its assembled answer as SSE events.

        Context comes from the `get_business_context` seam (a stub today); the
        `analysis_record` is persisted via the `AnalysisStore` seam (a no-op today).
        `chat_answer` is emitted as a single `chunk` (the Assembler returns the whole
        object — true token streaming is a later step).
        """
        from .planner.business_context import get_business_context
        from .planner.inputs import Constraints

        if tracer is None:
            from ..observability.langfuse.tracing import NullTracer

            tracer = NullTracer()

        coordinator = self._get_slow_path_coordinator(user_id, tracer, catalog_reader)
        context = await get_business_context(user_id)

        # DB3: warm the user's DB connection in parallel with planning so the
        # handshake overlaps the ~4s Planner call. Default path only — an injected
        # coordinator factory (tests / custom) may not use the real DbExecutor.
        if self._slow_path_factory is None:
            from ..query.executor.db import DbExecutor

            asyncio.create_task(DbExecutor.prewarm(catalog, user_id))  # noqa: RUF006

        pc = tracer.callbacks()              # planner: PII-safe, full capture
        ac = tracer.callbacks(masked=True)   # assembler: sees real rows -> masked
        run_kw: dict[str, Any] = {}
        if pc:
            run_kw["planner_callbacks"] = pc
        if ac:
            run_kw["assembler_callbacks"] = ac

        # R4: bridge the coordinator's per-stage progress callback to SSE `status`
        # events so the stream isn't silent for ~12s (and proxies don't drop the
        # idle connection). Status events only appear if the coordinator calls back.
        progress_q: asyncio.Queue[str] = asyncio.Queue()

        async def _progress(stage: str) -> None:
            await progress_q.put(stage)

        run_task = asyncio.create_task(
            coordinator.run(
                context, catalog, query, Constraints(), progress=_progress, **run_kw
            )
        )
        getter: asyncio.Task = asyncio.create_task(progress_q.get())
        pending: set[asyncio.Task] = {run_task, getter}
        while True:
            done, pending = await asyncio.wait(
                pending, return_when=asyncio.FIRST_COMPLETED
            )
            if getter in done:
                yield {"event": "status", "data": getter.result()}
                getter = asyncio.create_task(progress_q.get())
                pending = pending | {getter}
            if run_task in done:
                getter.cancel()
                while not progress_q.empty():
                    yield {"event": "status", "data": progress_q.get_nowait()}
                break

        try:
            result = run_task.result()
        except Exception as e:
            logger.error("slow path failed", user_id=user_id, error=str(e))
            yield {"event": "error", "data": f"Analysis failed: {e}"}
            return

        yield {"event": "sources", "data": json.dumps([])}  # TODO: derive from record
        yield {"event": "chunk", "data": result.chat_answer}
        try:
            await self._get_analysis_store().save(result.analysis_record)
        except Exception as e:  # persistence must never break the user's answer
            logger.error("analysis_record persist failed", user_id=user_id, error=str(e))
        tracer.end()  # output omitted (chat_answer may contain PII on Cloud)
        yield {"event": "done", "data": ""}


def _build_sources(
    source_hint: str,
    user_id: str,
    query_result: Any,
    raw_chunks: Any,
) -> list[dict[str, Any]]:
    """Build the sources payload for the SSE `sources` event.

    - structured: one entry per executed table (table_name only).
    - unstructured: deduped by (document_id, page_label), Phase 1 shape.
    - chat or error: empty list.
    """
    if source_hint == "structured":
        if query_result is None or getattr(query_result, "error", None):
            return []
        table_name = getattr(query_result, "table_name", "") or ""
        if not table_name:
            return []
        return [{
            "document_id": f"{user_id}_{table_name}",
            "filename": table_name,
            "page_label": None,
        }]

    if source_hint == "unstructured" and raw_chunks:
        seen: set[tuple[Any, Any]] = set()
        sources: list[dict[str, Any]] = []
        for item in raw_chunks:
            if isinstance(item, RetrievalResult):
                data = item.metadata.get("data", {})
            elif isinstance(item, dict):
                data = item
            else:
                continue
            key = (data.get("document_id"), data.get("page_label"))
            if key in seen or key == (None, None):
                continue
            seen.add(key)
            sources.append({
                "document_id": data.get("document_id"),
                "filename": data.get("filename", "Unknown"),
                "page_label": data.get("page_label", "Unknown"),
            })
        return sources

    return []


def _normalize_chunks(raw: Any) -> list[DocumentChunk]:
    """Convert whatever the retriever returns into list[DocumentChunk].

    The Phase 2 `DocumentRetriever.retrieve` interface is a stub today;
    when TAB owner ships it, it should return `list[DocumentChunk]`
    directly so this normalizer becomes a no-op. Until then we coerce
    common shapes (dict-with-content, plain string) defensively.
    """
    if not raw:
        return []
    if isinstance(raw, list) and all(isinstance(c, DocumentChunk) for c in raw):
        return raw
    chunks: list[DocumentChunk] = []
    for item in raw:
        if isinstance(item, DocumentChunk):
            chunks.append(item)
        elif isinstance(item, dict):
            chunks.append(
                DocumentChunk(
                    content=str(item.get("content", "")),
                    filename=item.get("filename"),
                    page_label=item.get("page_label"),
                )
            )
        elif isinstance(item, RetrievalResult):
            data = item.metadata.get("data", {})
            page = data.get("page_label")
            chunks.append(DocumentChunk(
                content=item.content,
                filename=data.get("filename"),
                page_label=str(page) if page is not None else None,
            ))
        elif isinstance(item, str):
            chunks.append(DocumentChunk(content=item))
    return chunks