| """JSON trace logger for agent runs. One file per run; flushed after every step |
| so a mid-run crash leaves a partial-but-valid log on disk.""" |
| import json |
| import os |
| import uuid |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
|
|
| class Trace: |
| """A single agent-run's trace. Writes JSON to <directory>/<run_id>.json |
| after every step. Use .flush() to force-write after the final answer.""" |
|
|
| def __init__( |
| self, |
| directory: str, |
| question: str, |
| provider: str, |
| model: str, |
| ): |
| os.makedirs(directory, exist_ok=True) |
| self.run_id = uuid.uuid4().hex[:12] |
| self.path = str(Path(directory) / f"{self.run_id}.json") |
| self.started_at = datetime.now(timezone.utc).isoformat() |
| self.question = question |
| self.provider = provider |
| self.model = model |
| self.steps: list[dict] = [] |
| self.final_answer: str | None = None |
| self.total_steps: int = 0 |
| self.error: str | None = None |
| self.flush() |
|
|
| def log_step(self, step: int, call, result: str) -> None: |
| """Record one tool invocation in step N.""" |
| |
| existing = next((s for s in self.steps if s["step"] == step), None) |
| if existing is None: |
| existing = {"step": step, "tool_calls": [], "results": []} |
| self.steps.append(existing) |
| existing["tool_calls"].append({ |
| "id": call.id, |
| "name": call.name, |
| "args": call.arguments, |
| }) |
| existing["results"].append(result) |
| self.flush() |
|
|
| def log_final(self, step: int, final_answer: str) -> None: |
| """Record the model's final answer (returned without tool_calls).""" |
| self.final_answer = final_answer |
| self.total_steps = step + 1 |
| self.flush() |
|
|
| def log_error(self, error: str) -> None: |
| self.error = error |
| self.flush() |
|
|
| def flush(self) -> None: |
| data = { |
| "run_id": self.run_id, |
| "question": self.question, |
| "started_at": self.started_at, |
| "provider": self.provider, |
| "model": self.model, |
| "steps": self.steps, |
| "final_answer": self.final_answer, |
| "total_steps": self.total_steps, |
| "error": self.error, |
| } |
| Path(self.path).write_text(json.dumps(data, indent=2)) |
|
|