engine/trace.py

"""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