File size: 2,270 Bytes
4948993
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""

Structured per-agent I/O tracing for FormScout.

Records every agent's input/output as JSON-serializable dicts.

Used for the Sharing-is-Caring badge (publish full trace to Hub).

"""
from __future__ import annotations

import json
import time
from dataclasses import asdict, is_dataclass
from pathlib import Path
from typing import Any


class TraceRecord:
    """A single agent execution record."""

    def __init__(self, agent_name: str, input_data: Any, output_data: Any, duration_ms: float):
        self.agent_name = agent_name
        self.input_summary = self._summarize(input_data)
        self.output_summary = self._summarize(output_data)
        self.duration_ms = duration_ms
        self.timestamp = time.time()

    def _summarize(self, data: Any) -> dict:
        """Convert dataclass or dict to JSON-safe summary."""
        if is_dataclass(data) and not isinstance(data, type):
            d = asdict(data)
            # Don't serialize raw frames (numpy arrays)
            if "frames" in d:
                d["frames"] = f"[{len(d['frames'])} frames]"
            return d
        if isinstance(data, dict):
            return data
        return {"value": str(data)}

    def to_dict(self) -> dict:
        return {
            "agent": self.agent_name,
            "timestamp": self.timestamp,
            "duration_ms": self.duration_ms,
            "input": self.input_summary,
            "output": self.output_summary,
        }


class PipelineTrace:
    """Collects trace records for a full pipeline run."""

    def __init__(self):
        self.records: list[TraceRecord] = []
        self.start_time = time.time()

    def add(self, record: TraceRecord):
        self.records.append(record)

    def to_dict(self) -> dict:
        return {
            "total_duration_ms": (time.time() - self.start_time) * 1000,
            "n_agents": len(self.records),
            "agents": [r.to_dict() for r in self.records],
        }

    def save(self, path: str | Path):
        """Save trace as JSON."""
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        with open(p, "w") as f:
            json.dump(self.to_dict(), f, indent=2, default=str)