src/traceability/scratchpad.py

"""Per-request traceability accumulator + tool-invoker wrapper (KM-691).

`TraceabilityScratchpad` is a mutable, per-request blackboard that `ChatHandler`
fills while answering a turn, then `build()`s into a `TraceabilityPayload` right
before the `done` SSE event. `TraceabilityToolInvoker` wraps the real tool invoker
so every tool call on the slow path / check branch records its full I/O into the
scratchpad (mirrors `TracingToolInvoker` in `src/observability/langfuse/tracing.py`,
whose name is taken — this one records real I/O, not masked metadata).

Everything here is best-effort: recording must never break the user's answer.
"""

from __future__ import annotations

from typing import Any

from pydantic import BaseModel

from src.middlewares.logging import get_logger

from .schemas import PlanningInfo, PlanStep, ToolCallInfo, TraceabilityPayload

logger = get_logger("traceability")

# Truncation caps (bound the JSONB payload) — see plan §3.
CAP_PREVIEW_ROWS = 5
CAP_STR = 300
# Executed queries get a higher cap than free strings: the SQL/query is the point of
# the feature, and 300 chars mangles all but the smallest statements.
CAP_QUERY = 2000


def _truncate(obj: Any) -> Any:
    """Recursively cap strings to CAP_STR and summarize embedded tool results.

    A Pattern-A `analyze_*` input carries its upstream `retrieve_data` result as a
    `ToolOutput` (a BaseModel). Left untouched that would re-embed the FULL upstream
    table (all rows + the untruncated query) into the payload, defeating the caps —
    so a tool-result-shaped model is summarized via `_output_to_dict` (preview ≤ 5
    rows, `row_count` kept), and any other BaseModel is dumped then recursed.
    """
    if isinstance(obj, str):
        return obj[:CAP_STR]
    if isinstance(obj, BaseModel):
        if hasattr(obj, "kind") and hasattr(obj, "rows"):  # a ToolOutput-shaped result
            return _output_to_dict(obj)
        return _truncate(obj.model_dump(mode="json"))
    if isinstance(obj, dict):
        return {k: _truncate(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_truncate(v) for v in obj]
    return obj


def _output_to_dict(output: Any) -> dict[str, Any]:
    """Normalize a tool result (`ToolOutput` or a synth dict) to the wire shape:
    kind/columns/row_count/preview/value/error, all truncation-capped."""
    if isinstance(output, dict):
        return _truncate(output)

    kind = getattr(output, "kind", None)
    result: dict[str, Any] = {"kind": kind}
    rows = getattr(output, "rows", None)
    if rows is not None:
        result["row_count"] = len(rows)
        columns = getattr(output, "columns", None)
        if columns is not None:
            result["columns"] = list(columns)
        result["preview"] = [
            [_truncate(cell) for cell in row] for row in rows[:CAP_PREVIEW_ROWS]
        ]
    value = getattr(output, "value", None)
    if value is not None:
        result["value"] = _truncate(value)
    error = getattr(output, "error", None)
    if error is not None:
        result["error"] = _truncate(error)
    return result


def _meta_of(output: Any) -> dict[str, Any]:
    """Best-effort read of a tool result's `meta` dict (ToolOutput or plain dict)."""
    if isinstance(output, dict):
        meta = output.get("meta")
    else:
        meta = getattr(output, "meta", None)
    return meta if isinstance(meta, dict) else {}


def _summarize(name: str, out_dict: dict[str, Any], meta: dict[str, Any]) -> str:
    """One plain-English line per tool step (fixed templates — never an LLM call)."""
    rc = out_dict.get("row_count")
    cols = out_dict.get("columns")
    ncol = len(cols) if isinstance(cols, list) else None
    if name == "check_data":
        return "Inspected your data source structure"
    if name == "retrieve_data":
        parts = "Retrieved"
        if rc is not None:
            parts += f" {rc} rows"
        if ncol:
            parts += f" across {ncol} columns"
        table = meta.get("table_name")
        if table:
            parts += f" from {table}"
        return parts
    if name == "retrieve_knowledge":
        return f"Searched your documents ({rc if rc is not None else 0} passages found)"
    if name.startswith("analyze_"):
        pretty = name.removeprefix("analyze_").replace("_", " ")
        return f"Ran {pretty} analysis on {ncol} columns" if ncol else f"Ran {pretty} analysis"
    return name.replace("_", " ")


class TraceabilityScratchpad:
    """Mutable per-request accumulator; `build()` freezes it into a payload."""

    def __init__(self) -> None:
        self.message_id: str | None = None  # set at handler entry; None => no flush
        self.intent: str = "chat"  # default until the router classifies
        self._planning: PlanningInfo | None = None
        self._tool_calls: list[ToolCallInfo] = []
        self._db_sources: list[dict[str, Any]] = []
        self._doc_sources: list[dict[str, Any]] = []
        self._doc_seen: set[tuple[Any, Any]] = set()
        self._catalog: Any = None  # set by set_catalog: enables id->name resolution
        self._retrieve_calls: list[dict[str, Any]] = []  # raw retrieve_data IRs + meta

    def set_intent(self, intent: str) -> None:
        self.intent = intent

    def set_catalog(self, catalog: Any) -> None:
        """Provide the catalog used this turn so `build()` can resolve the IR ids in
        each retrieve_data call to real names (the `data_used` layer). No-op-safe:
        without it, `data_used` stays empty and the raw tool_calls still carry the IR."""
        self._catalog = catalog

    def record_tool_call(
        self,
        name: str,
        args: dict[str, Any],
        output: Any,
        task_id: str | None = None,
    ) -> None:
        """Append one tool call (input + normalized output). For `retrieve_data`,
        also derive a database source from the args + executed query in `meta`."""
        out_dict = _output_to_dict(output)
        status = "error" if out_dict.get("kind") == "error" else "success"
        self._tool_calls.append(
            ToolCallInfo(
                order=len(self._tool_calls) + 1,
                task_id=task_id,
                name=name,
                summary=_summarize(name, out_dict, _meta_of(output)),
                input=_truncate(dict(args)),
                output=out_dict,
                status=status,
                error=out_dict.get("error"),
            )
        )
        if name == "retrieve_data" and status == "success":
            self._record_db_source(output)
            self._capture_retrieve(args, output)

    def _capture_retrieve(self, args: Any, output: Any) -> None:
        """Stash the raw retrieve_data IR + provenance meta so `build()` can resolve a
        `DataUsed`. Gated on the SAME `meta.source_id` check as `_record_db_source`, so
        the two stay index-aligned (the db source was just appended)."""
        meta = _meta_of(output)
        if not meta.get("source_id"):
            return
        ir = args.get("ir") if isinstance(args, dict) else None
        if not isinstance(ir, dict):
            return
        query = meta.get("query")
        self._retrieve_calls.append({
            "ir": ir,
            "query": query[:CAP_QUERY] if isinstance(query, str) else None,
            "source_name": meta.get("source_name"),
            "row_count": meta.get("row_count"),
            "db_source_index": len(self._db_sources) - 1,
        })

    def _record_db_source(self, output: Any) -> None:
        # retrieve_data's args are {"ir": ...}; the reliable source_id/table/query
        # live on the tool OUTPUT meta (see tools/data_access.py::_retrieve_data).
        meta = _meta_of(output)
        if not meta.get("source_id"):
            # A failed/aborted retrieval carries no provenance meta — emitting it
            # anyway produced all-null source rows in the payload.
            return
        query = meta.get("query")
        table = meta.get("table_name") or meta.get("table_id")
        self._db_sources.append({
            "type": "database",
            "source_id": meta.get("source_id"),
            "name": table,
            "query": query[:CAP_QUERY] if isinstance(query, str) else None,
            "detail": {"table": table, "row_count": meta.get("row_count")},
        })

    def set_planning_from_record(self, record: Any) -> None:
        """Map an `AnalysisRecord` (goal_restated + tasks_run) to `PlanningInfo`."""
        try:
            steps = [
                PlanStep(
                    step=i + 1,
                    stage=str(getattr(task, "stage", "")),
                    objective=getattr(task, "objective", ""),
                    status=str(getattr(task, "status", "")),
                    tools_used=list(getattr(task, "tools_used", []) or []),
                )
                for i, task in enumerate(getattr(record, "tasks_run", []) or [])
            ]
            self._planning = PlanningInfo(
                goal_restated=getattr(record, "goal_restated", "") or "",
                assumptions=[],  # AnalysisRecord carries no assumptions field (honest: empty)
                steps=steps,
            )
        except Exception as exc:  # never break the answer on a mapping slip
            logger.warning("traceability planning mapping failed", error=str(exc))

    def add_document_sources(self, raw_chunks: Any, query: str) -> None:
        """Dedupe retrieved chunks by (document_id, page_label) into document
        sources, stamped with the query. Sole source of document provenance now
        that the stream no longer emits sources (KM-691)."""
        for item in raw_chunks or []:
            if hasattr(item, "metadata"):
                data = item.metadata.get("data", {})
            elif isinstance(item, dict):
                data = item
            else:
                continue
            key = (data.get("document_id"), data.get("page_label"))
            if key == (None, None) or key in self._doc_seen:
                continue
            self._doc_seen.add(key)
            source: dict[str, Any] = {
                "type": "document",
                "document_id": data.get("document_id"),
                "filename": data.get("filename", "Unknown"),
                "page_label": data.get("page_label"),
                "query": _truncate(query),
            }
            snippet = data.get("snippet") or data.get("content") or data.get("text")
            if isinstance(snippet, str):
                source["snippet"] = snippet[:CAP_STR]
            score = data.get("score")
            if score is not None:
                source["score"] = score
            self._doc_sources.append(source)

    def _build_data_used(self) -> list[Any]:
        """Resolve each captured retrieve_data IR into a `DataUsed` (real names) and
        enrich the matching db source with source_name + all tables touched. Never
        raises — a resolution slip drops that entry, never breaks the answer."""
        if self._catalog is None or not self._retrieve_calls:
            return []
        from ..query.ir.models import QueryIR
        from .resolve import resolve_data_used

        out: list[Any] = []
        for rc in self._retrieve_calls:
            try:
                ir = QueryIR.model_validate(rc["ir"])
                du = resolve_data_used(ir, self._catalog, rc.get("query"), rc.get("row_count"))
                out.append(du)
                idx = rc.get("db_source_index")
                if isinstance(idx, int) and 0 <= idx < len(self._db_sources):
                    self._db_sources[idx]["source_name"] = rc.get("source_name")
                    self._db_sources[idx]["tables"] = [t.name for t in du.tables]
            except Exception as exc:  # never break the answer on a resolve slip
                logger.warning("data_used resolve failed", error=str(exc))
        return out

    def build(self, analysis_id: str, user_id: str, message_id: str) -> TraceabilityPayload:
        """Freeze the accumulated state into a `TraceabilityPayload`."""
        from datetime import UTC, datetime

        data_used = self._build_data_used()  # also enriches self._db_sources in place
        return TraceabilityPayload(
            analysis_id=analysis_id,
            message_id=message_id,
            user_id=user_id,
            intent=self.intent,
            generated_at=datetime.now(UTC),
            planning=self._planning,
            thinking=None,
            tool_calls=list(self._tool_calls),
            data_used=data_used,
            sources=self._doc_sources + self._db_sources,
        )


class TraceabilityToolInvoker:
    """Wraps a ToolInvoker to record each call's full I/O into a scratchpad.

    Implements the ToolInvoker protocol (`async invoke(tool_name, args)`). Recording
    is never-throw so a trace slip can't break the tool run. Distinct from
    `TracingToolInvoker` (Langfuse, masked-metadata-only) — that name is taken.
    """

    def __init__(self, inner: Any, pad: TraceabilityScratchpad) -> None:
        self._inner = inner
        self._pad = pad

    async def invoke(self, tool_name: str, args: dict[str, Any]) -> Any:
        out = await self._inner.invoke(tool_name, args)
        try:
            self._pad.record_tool_call(tool_name, args, out)
        except Exception as exc:  # never break the tool run
            logger.warning("traceability tool record failed", tool=tool_name, error=str(exc))
        return out