File size: 8,541 Bytes
23cc207 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 | """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 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
def _truncate(obj: Any) -> Any:
"""Recursively cap any string to CAP_STR; leave numbers/None untouched."""
if isinstance(obj, str):
return obj[:CAP_STR]
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": _truncate(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 (mirrors `chat_handler._build_sources`), stamped with the query."""
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
|