from __future__ import annotations

import json
from pathlib import Path


spans = [
    {"name": "agent.plan", "parent": None, "latency_ms": 180, "tokens": 420, "status": "ok"},
    {"name": "tool.search", "parent": "agent.plan", "latency_ms": 620, "tokens": 0, "status": "ok"},
    {"name": "agent.answer", "parent": None, "latency_ms": 240, "tokens": 610, "status": "ok"},
]

trace = {
    "trace_id": "trace-demo-001",
    "spans": spans,
    "total_latency_ms": sum(s["latency_ms"] for s in spans),
    "total_tokens": sum(s["tokens"] for s in spans),
    "tool_error_count": sum(s["status"] != "ok" and s["name"].startswith("tool.") for s in spans),
    "terminal": "success",
}
assert trace["total_latency_ms"] == 1040
assert trace["total_tokens"] == 1030
path = Path(__file__).with_name("experiment-result.json")
path.write_text(json.dumps(trace, ensure_ascii=False, indent=2) + "\n")
print(json.dumps(trace, ensure_ascii=False, indent=2))