Extract Span

1"""Extract a specific span's full input/output state by name.
2
3Usage:
4  SPAN="upsert_dashboard" python3 scripts/extract_span.py FILE
5  SPAN="router" python3 scripts/extract_span.py FILE
6
7Env vars:
8  SPAN     — span name to match (case-insensitive substring match)
9  MAX_LEN  — truncation limit (default 0 = unlimited)
10"""
11
12import json
13import os
14import sys
15
16
17def load_trace_file(path):
18    with open(path) as f:
19        raw = json.load(f)
20    if isinstance(raw, list) and raw and raw[0].get("type") == "text":
21        raw = json.loads(raw[0]["text"])
22    results = raw.get("results", raw)
23    return [results] if isinstance(results, dict) else results
24
25
26def truncate(text, max_len):
27    if max_len <= 0 or len(text) <= max_len:
28        return text
29    return text[:max_len] + f"... [{len(text)} chars total]"
30
31
32span_filter = os.environ.get("SPAN", "").lower()
33if not span_filter:
34    print("Usage: SPAN='span_name' python3 extract_span.py file.json", file=sys.stderr)
35    sys.exit(1)
36
37max_len = int(os.environ.get("MAX_LEN", "0"))
38
39traces = load_trace_file(sys.argv[1])
40found = 0
41for trace in traces:
42    events = sorted(trace.get("events", []), key=lambda e: e.get("createdAt", ""))
43    for ev in events:
44        if ev.get("event") != "$ai_span":
45            continue
46        p = ev.get("properties", {})
47        name = p.get("$ai_span_name", "")
48        if span_filter not in name.lower():
49            continue
50        found += 1
51        error = " [ERROR]" if p.get("$ai_is_error") else ""
52        print(f"\n{'='*80}")
53        print(f"SPAN: {name}  ({p.get('$ai_latency', '?')}s){error}")
54        print(f"Created: {ev.get('createdAt', '?')}")
55        print(f"Parent: {p.get('$ai_parent_id', '(root)')}")
56        print(f"{'='*80}")
57
58        inp = p.get("$ai_input_state")
59        out = p.get("$ai_output_state")
60
61        if inp is not None:
62            formatted = json.dumps(inp, indent=2, default=str) if not isinstance(inp, str) else inp
63            print(f"\n--- INPUT STATE ---")
64            print(truncate(formatted, max_len))
65
66        if out is not None:
67            formatted = json.dumps(out, indent=2, default=str) if not isinstance(out, str) else out
68            print(f"\n--- OUTPUT STATE ---")
69            print(truncate(formatted, max_len))
70
71        if not inp and not out:
72            print("\n  (no input_state or output_state)")
73
74if found == 0:
75    print(f"No spans matching '{span_filter}' found.", file=sys.stderr)
76    sys.exit(1)