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 {}


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()

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

    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,
                input=_truncate(dict(args)),
                output=out_dict,
                status=status,
                error=out_dict.get("error"),
            )
        )
        if name == "retrieve_data":
            self._record_db_source(output)

    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)
        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(self, analysis_id: str, user_id: str, message_id: str) -> TraceabilityPayload:
        """Freeze the accumulated state into a `TraceabilityPayload`."""
        from datetime import UTC, datetime

        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),
            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