| 44 | |
| 45 | |
| 46 | class TraceReplayer: |
| 47 | def __init__(self, events: List[Dict[str, Any]]): |
| 48 | self._events = events |
| 49 | self._idx = 0 |
| 50 | |
| 51 | @classmethod |
| 52 | def from_file(cls, path: Path) -> "TraceReplayer": |
| 53 | return cls(list(iter_jsonl(path))) |
| 54 | |
| 55 | def _next(self) -> Dict[str, Any]: |
| 56 | if self._idx >= len(self._events): |
| 57 | raise IndexError("trace exhausted") |
| 58 | ev = self._events[self._idx] |
| 59 | self._idx += 1 |
| 60 | return ev |
| 61 | |
| 62 | def next_llm_message(self, *, step_id: str, iteration: int) -> ReplayMessage: |
| 63 | ev = self._next() |
| 64 | if ev.get("type") != "llm_response": |
| 65 | raise ValueError( |
| 66 | f"trace mismatch: expected llm_response, got {ev.get('type')}" |
| 67 | ) |
| 68 | # Validate alignment between expected and actual event |
| 69 | ev_step_id = str(ev.get("step_id", "")) |
| 70 | ev_iteration = ev.get("iteration") |
| 71 | if ev_step_id != str(step_id): |
| 72 | raise ValueError( |
| 73 | f"trace mismatch: expected step={step_id}, got step={ev_step_id}" |
| 74 | ) |
| 75 | try: |
| 76 | ev_iter_int = int(ev_iteration) if ev_iteration is not None else -1 |
| 77 | except (ValueError, TypeError) as e: |
| 78 | raise ValueError( |
| 79 | f"trace mismatch: invalid iteration in trace: {ev_iteration!r}" |
| 80 | ) from e |
| 81 | if ev_iter_int != int(iteration): |
| 82 | raise ValueError( |
| 83 | f"trace mismatch: expected iteration={iteration}, got iteration={ev_iter_int}" |
| 84 | ) |
| 85 | tool_calls = [] |
| 86 | for tc in ev.get("tool_calls") or []: |
| 87 | tool_calls.append( |
| 88 | ReplayToolCall( |
| 89 | id=str(tc.get("id")), |
| 90 | function=ReplayFunction( |
| 91 | name=str((tc.get("function") or {}).get("name") or ""), |
| 92 | arguments=str( |
| 93 | (tc.get("function") or {}).get("arguments") or "{}" |
| 94 | ), |
| 95 | ), |
| 96 | ) |
| 97 | ) |
| 98 | return ReplayMessage( |
| 99 | content=str(ev.get("content") or ""), tool_calls=tool_calls |
| 100 | ) |
| 101 | |
| 102 | def next_tool_result( |
| 103 | self, |
no outgoing calls