Spaces:
Paused
Paused
File size: 2,645 Bytes
daea45b | 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 | """Shareable agent traces (Build Small "Sharing is Caring" badge).
Turns a completed smolcode run into an OpenTelemetry-style JSON trace: a root
span minted by LiteForge's `Tracer` plus one child span per agent step, carrying
the step kind, duration, and token counts read from `AgentStep`. Publish a trace
file to the Hub so others can see exactly how the tiny model reasoned.
"""
from __future__ import annotations
import json
import time
from pathlib import Path
import liteforge as lf
def build_trace(agent, task: str, final: str, *, preset: str, model: str) -> dict:
"""Build an OTel-ish trace document from a finished agent run."""
tracer = lf.Tracer("smolcode")
root = tracer.start_span("coding_task")
root.set_attribute("preset", preset)
root.set_attribute("model", model)
root.set_attribute("task", task)
trace_id = root.context.trace_id
root_id = root.context.span_id
spans: list[dict] = []
total_tokens = 0
history = agent.raw_history() if hasattr(agent, "raw_history") else getattr(agent, "history", lambda: [])()
for i, s in enumerate(history):
dur = getattr(s, "duration_ms", None) or 0
tot = getattr(s, "total_tokens", None) or 0
step_no = getattr(s, "step_number", getattr(s, "number", i))
step_type = getattr(s, "step_type", getattr(s, "kind", "step"))
result_text = getattr(s, "result", getattr(s, "detail", ""))
total_tokens += tot or 0
spans.append({
"trace_id": trace_id,
"span_id": f"{root_id[:-len(str(step_no))-1]}{step_no:02d}",
"parent_span_id": root_id,
"name": str(step_type),
"duration_ms": dur,
"attributes": {
"step_number": step_no,
"prompt_tokens": getattr(s, "prompt_tokens", None),
"completion_tokens": getattr(s, "completion_tokens", None),
"total_tokens": tot,
"result": str(result_text)[:200],
},
})
root.end()
return {
"trace_id": trace_id,
"service": "smolcode",
"preset": preset,
"model": model,
"task": task,
"final": final,
"n_steps": len(spans),
"total_tokens": total_tokens,
"root": {"span_id": root_id, "name": "coding_task"},
"spans": spans,
}
def save_trace(trace: dict, out_dir: str | Path = "traces") -> Path:
d = Path(out_dir)
d.mkdir(parents=True, exist_ok=True)
stamp = time.strftime("%Y%m%d-%H%M%S")
path = d / f"trace-{stamp}.json"
path.write_text(json.dumps(trace, indent=2))
return path
|