pyrtan / src / pyrtan / trace.py
2,435 bytes · last changed 2026-10-25 in Initial public release
| 1 | """What a run did, step by step. |
| 2 | |
| 3 | A trace is the product that distinguishes a decision tree from a model: every |
| 4 | observation is recorded, every branch names the condition that chose it, and |
| 5 | the whole thing is plain data a person or another service can read. |
| 6 | """ |
| 7 | |
| 8 | from __future__ import annotations |
| 9 | |
| 10 | import json |
| 11 | import time |
| 12 | from dataclasses import dataclass, field, asdict |
| 13 | from typing import Any |
| 14 | |
| 15 | |
| 16 | @dataclass |
| 17 | class Step: |
| 18 | node: str |
| 19 | kind: str |
| 20 | summary: str |
| 21 | observed: Any = None |
| 22 | chosen: str | None = None |
| 23 | next: str | None = None |
| 24 | elapsed_ms: int = 0 |
| 25 | depth: int = 0 |
| 26 | |
| 27 | def to_dict(self) -> dict[str, Any]: |
| 28 | d = asdict(self) |
| 29 | if d["observed"] is None: |
| 30 | del d["observed"] |
| 31 | if d["chosen"] is None: |
| 32 | del d["chosen"] |
| 33 | return d |
| 34 | |
| 35 | |
| 36 | @dataclass |
| 37 | class Trace: |
| 38 | tree: str |
| 39 | steps: list[Step] = field(default_factory=list) |
| 40 | started: float = field(default_factory=time.time) |
| 41 | gaps: list[dict[str, Any]] = field(default_factory=list) |
| 42 | |
| 43 | def add(self, step: Step) -> Step: |
| 44 | self.steps.append(step) |
| 45 | return step |
| 46 | |
| 47 | def gap(self, node: str, on: str, observed: Any) -> None: |
| 48 | """Record a decision whose observed value matched no branch. Gaps are |
| 49 | how a tree grows: each one is a branch an author has not written yet.""" |
| 50 | self.gaps.append({"tree": self.tree, "node": node, "on": on, "observed": _short(observed)}) |
| 51 | |
| 52 | def to_dict(self) -> dict[str, Any]: |
| 53 | return {"tree": self.tree, "started": self.started, "steps": [s.to_dict() for s in self.steps], "gaps": self.gaps} |
| 54 | |
| 55 | def render(self) -> str: |
| 56 | lines = [f"trace: {self.tree}"] |
| 57 | for s in self.steps: |
| 58 | pad = " " * (s.depth + 1) |
| 59 | arrow = f" -> {s.next}" if s.next else "" |
| 60 | chosen = f" [{s.chosen}]" if s.chosen else "" |
| 61 | lines.append(f"{pad}{s.node} ({s.kind}){chosen}: {s.summary}{arrow}") |
| 62 | for g in self.gaps: |
| 63 | lines.append(f" gap: {g['node']} saw {g['on']} = {json.dumps(g['observed'], ensure_ascii=False)} and had no branch for it") |
| 64 | return "\n".join(lines) |
| 65 | |
| 66 | |
| 67 | def _short(value: Any, limit: int = 200) -> Any: |
| 68 | if isinstance(value, str) and len(value) > limit: |
| 69 | return value[:limit] + "…" |
| 70 | if isinstance(value, (dict, list)): |
| 71 | text = json.dumps(value, ensure_ascii=False) |
| 72 | if len(text) > limit: |
| 73 | return text[:limit] + "…" |
| 74 | return value |