"""Turn agent run logs into one JSON-lines file of complete trajectories. Purpose: the first stage of Part 27. Reads the logs the course's own agents already write - Part 24's minimal-agent.py transcripts and Part 26's agent-eval.py results - and normalises them into one episode per line: the tool list, the whole message sequence including tool calls and tool results, and the outcome. Everything downstream in this part reads that one shape, so a reader whose agent logs something else has one adapter to write rather than four. Platform: all (standard library only; no model, no accelerator, no network) Minimum memory: 8 GB nominally, and far less in practice: this is text in memory Assumes: Python 3.10 or newer. For --from part-24, a lab notebook containing records written by minimal-agent.py --labbook, whose "transcript" paths still resolve. For --from part-26, the results file agent-eval.py wrote with --out, whose "trajectory_file" paths still resolve. --tools takes an OpenAI-style tool list, which is what Part 24's Toolbox.schemas() returns. Usage: python3 collect-trajectories.py --from part-24 --labbook labbook.md \\ --tools tools.json --out raw/trajectories.jsonl python3 collect-trajectories.py --from part-26 --results agent-results.json \\ --tools tools.json --out raw/trajectories.jsonl python3 collect-trajectories.py --from jsonl --input other-agent.jsonl \\ --out raw/trajectories.jsonl python3 collect-trajectories.py --make-example raw/example-trajectories.jsonl The episode shape, which is the contract for the rest of this part: {"id": "...", "task_id": "...", "task": "...", "source": "part-24/minimal-agent", "model": "local/agent", "scaffold": "minimal-agent", "tools": [ {"type": "function", "function": {"name": ..., "description": ..., "parameters": { JSON schema }}} ], "messages": [ {"role": "system"|"user"|"assistant"|"tool", ...} ], "outcome": {"passed": true, "stopped": "finished", "steps": 4, "turns": 3, "tokens": 2841, "seconds": 22.4}, "lossy": false, "collected_at": "2026-09-09T09:00:00Z"} An assistant message that calls tools carries "tool_calls", each one {"type": "function", "function": {"name": ..., "arguments": { ... }}} with the arguments as an object rather than a string: that is the shape the Transformers chat templates and TRL's SFT trainer expect. trajectories-to-sft.py converts to the string form where a tool chain wants one. """ from __future__ import annotations import argparse import json import sys from datetime import datetime, timezone from pathlib import Path from typing import Any # -------------------------------------------------------------------------------------- # Reading # -------------------------------------------------------------------------------------- def read_jsonl(path: Path) -> list[dict[str, Any]]: """Every JSON object on its own line. Lines that are not objects are skipped, so a notebook that also contains prose is a valid input.""" rows: list[dict[str, Any]] = [] if not path.is_file(): sys.exit(f"{path} does not exist") for line in path.read_text(encoding="utf-8", errors="replace").splitlines(): line = line.strip() if not line.startswith("{"): continue try: value = json.loads(line) except json.JSONDecodeError: continue if isinstance(value, dict): rows.append(value) return rows def parse_arguments(raw: Any) -> dict[str, Any]: """Tool-call arguments as an object, whatever the log stored them as.""" if isinstance(raw, dict): return raw if isinstance(raw, str) and raw.strip(): try: value = json.loads(raw) except json.JSONDecodeError: return {"_unparsed": raw} return value if isinstance(value, dict) else {"_value": value} return {} def now() -> str: return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") # -------------------------------------------------------------------------------------- # Part 24: minimal-agent.py transcripts # -------------------------------------------------------------------------------------- def episode_from_part24(run: dict[str, Any], result: dict[str, Any], tools: list[dict], system_prompt: str | None, keep_reasoning: bool, finish_tool: str) -> dict[str, Any] | None: """One episode from one entry in a Part 24 lab-notebook record. The transcript holds the assistant's text and the arguments of every executed call. The call to the finish tool is the one that never appears as an executed call, because the loop stops on it, so it is rebuilt from the recorded answer. Rebuilding it matters: a model trained on trajectories that end in plain text learns to stop calling finish, and the harness that checks for finish then scores it as a failure. """ transcript_path = Path(str(result.get("transcript", ""))) if not transcript_path.is_file(): print(f" skipped {result.get('id')}: transcript {transcript_path} is gone") return None rows = read_jsonl(transcript_path) if not rows: print(f" skipped {result.get('id')}: transcript is empty") return None messages: list[dict[str, Any]] = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) messages.append({"role": "user", "content": result.get("task", "")}) # Group the transcript by turn, in file order. Each turn is one assistant row # followed by the rows of the tools that turn executed. turns: list[list[dict[str, Any]]] = [] for row in rows: if "assistant" in row or not turns: turns.append([row]) else: turns[-1].append(row) answer = result.get("answer") or "" for group in turns: head = group[0] assistant = head.get("assistant") or {} tool_rows = [r for r in group if r.get("tool")] names = list(assistant.get("tool_calls") or []) if not names: content = assistant.get("content") if content: messages.append({"role": "assistant", "content": content}) continue # Match each announced call to the next executed row with that name. The one # that has no row is the finish call. remaining = list(tool_rows) calls: list[dict[str, Any]] = [] matched: list[dict[str, Any] | None] = [] for name in names: row = next((r for r in remaining if r.get("tool") == name), None) if row is not None: remaining.remove(row) calls.append({"type": "function", "function": { "name": name, "arguments": parse_arguments(row.get("arguments"))}}) matched.append(row) elif name == finish_tool: calls.append({"type": "function", "function": { "name": finish_tool, "arguments": {"answer": answer}}}) matched.append(None) else: # Announced, never executed, and not the finish tool: the loop refused # it. Dropping it silently would train the model on a call it did not # make, so the whole episode is dropped instead. print(f" skipped {result.get('id')}: call to {name!r} has no recorded result") return None message: dict[str, Any] = {"role": "assistant", "content": assistant.get("content"), "tool_calls": calls} if keep_reasoning and assistant.get("reasoning"): message["reasoning_content"] = assistant["reasoning"] messages.append(message) for name, row in zip(names, matched): if row is None: continue messages.append({"role": "tool", "name": name, "content": str(row.get("result", ""))}) return { "id": f"{run.get('model', 'model')}-{result.get('id', 'task')}-{transcript_path.stem}", "task_id": result.get("id"), "task": result.get("task", ""), "source": "part-24/minimal-agent", "model": run.get("model"), "scaffold": "minimal-agent", "tools": tools, "messages": messages, "outcome": { "passed": result.get("passed"), "stopped": result.get("stopped"), "steps": result.get("turns"), "turns": result.get("turns"), "tokens": result.get("tokens"), "seconds": result.get("seconds"), "tool_errors": sum(1 for r in rows if r.get("is_error")), }, "lossy": False, "collected_at": now(), } # -------------------------------------------------------------------------------------- # Part 26: agent-eval.py results # -------------------------------------------------------------------------------------- def episode_from_part26(run: dict[str, Any], row: dict[str, Any], tools: list[dict], system_prompt: str | None) -> dict[str, Any] | None: """One episode from one row of a Part 26 results file. Part 26's trajectory records what a tool was called with and what came back, which is what a reader needs to debug a run. It does not record the assistant's own words, so what is reconstructed here is a plausible trajectory rather than the one that happened, and every episode is marked lossy. Training on these teaches tool sequences and teaches nothing about what the model said between them; the honest fix is to log the assistant messages in the first place, which is what Part 24's agent does. """ path = Path(str(row.get("trajectory_file", ""))) steps = read_jsonl(path) if path.is_file() else [] tool_steps = [s for s in steps if s.get("tool")] if not tool_steps: return None messages: list[dict[str, Any]] = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) meta = next((s.get("meta") for s in steps if isinstance(s.get("meta"), dict)), {}) or {} task_text = meta.get("task") or row.get("task") or "" messages.append({"role": "user", "content": task_text}) for step in tool_steps: name = str(step["tool"]) messages.append({"role": "assistant", "content": None, "tool_calls": [ {"type": "function", "function": {"name": name, "arguments": parse_arguments(step.get("arguments"))}}]}) messages.append({"role": "tool", "name": name, "content": str(step.get("result", ""))}) messages.append({"role": "assistant", "content": row.get("answer", "")}) checks = row.get("checks") or {} return { "id": f"{run.get('run_id', 'run')}-{row.get('id')}-r{row.get('attempt', 1)}", "task_id": row.get("id"), "task": task_text, "source": "part-26/agent-eval", "model": run.get("model"), "scaffold": run.get("scaffold"), "tools": tools, "messages": messages, "outcome": { "passed": checks.get("passed"), "stopped": row.get("stopped"), "steps": row.get("steps"), "turns": None, "tokens": row.get("tokens"), "seconds": row.get("seconds"), "tool_errors": sum(1 for s in tool_steps if s.get("is_error")), }, "lossy": True, "collected_at": now(), } # -------------------------------------------------------------------------------------- # Tools # -------------------------------------------------------------------------------------- def infer_tools(episodes: list[dict[str, Any]]) -> list[dict[str, Any]]: """Build a tool list from the calls that were actually made. This exists so that the pipeline can be run before the real schemas are to hand, and it is not what you should train on. The description of a tool is part of the prompt the model reads at serving time; inventing an empty one here teaches the model to map an empty description onto a call, and at serving time it sees the real description instead. Every inferred entry says so in its own description. """ seen: dict[str, set[str]] = {} for episode in episodes: for message in episode["messages"]: for call in message.get("tool_calls") or []: function = call.get("function", {}) name = function.get("name") if not name: continue keys = seen.setdefault(name, set()) keys.update(k for k in (function.get("arguments") or {}) if isinstance(k, str)) tools = [] for name in sorted(seen): properties = {key: {"type": "string"} for key in sorted(seen[name])} tools.append({"type": "function", "function": { "name": name, "description": "Inferred from observed calls; replace with the tool's real description.", "parameters": {"type": "object", "properties": properties, "required": sorted(properties)}, }}) return tools # -------------------------------------------------------------------------------------- # The worked example # -------------------------------------------------------------------------------------- EXAMPLE_TOOLS = [ {"type": "function", "function": { "name": "list_files", "description": "List the file names in a directory inside the workspace.", "parameters": {"type": "object", "properties": {"path": {"type": "string", "description": "Directory relative to the workspace root."}}, "required": ["path"]}}}, {"type": "function", "function": { "name": "read_file", "description": "Read a text file inside the workspace and return its contents.", "parameters": {"type": "object", "properties": {"path": {"type": "string", "description": "File relative to the workspace root."}}, "required": ["path"]}}}, {"type": "function", "function": { "name": "finish", "description": "Give the final answer and end the run.", "parameters": {"type": "object", "properties": {"answer": {"type": "string", "description": "The answer, in plain prose."}}, "required": ["answer"]}}}, ] EXAMPLE_SYSTEM = ("You are a careful assistant working inside a sandboxed workspace. " "Use the tools to find things out, take one step at a time, and call " "finish exactly once when you have the answer.") # The example transcripts have to contain something worth redacting, and this file is # published, so nothing here may be a real secret or a real address. These four values # are the safe stand-ins: an address in the loopback range, a documentation e-mail # domain, a home path with a placeholder user name, and a token that is a prefix plus a # run of one letter, assembled here so that no line of this file looks like a # credential. Real transcripts contain the real things, which is the whole point of the # scrubber. SAFE_ENDPOINT = "http://127.0.0.1:4000/v1" SAFE_EMAIL = "ops@example.com" SAFE_HOME_PATH = "/home/user/deploy/logs" SAFE_FAKE_TOKEN = "hf_" + "y" * 34 def example_episodes() -> list[dict[str, Any]]: """Four short episodes: two that succeeded, one that succeeded the long way round, and one that failed. They exercise every filter in trajectories-to-sft.py and carry the kind of leak scrub-trajectories.py is looking for.""" def episode(task_id: str, task: str, calls: list[tuple[str, dict, str]], answer: str, passed: bool, stopped: str, tokens: int, seconds: float, tool_errors: int = 0) -> dict[str, Any]: messages: list[dict[str, Any]] = [ {"role": "system", "content": EXAMPLE_SYSTEM}, {"role": "user", "content": task}, ] for name, arguments, result in calls: messages.append({"role": "assistant", "content": None, "tool_calls": [ {"type": "function", "function": {"name": name, "arguments": arguments}}]}) messages.append({"role": "tool", "name": name, "content": result}) messages.append({"role": "assistant", "content": None, "tool_calls": [ {"type": "function", "function": {"name": "finish", "arguments": {"answer": answer}}}]}) return { "id": f"example-{task_id}", "task_id": task_id, "task": task, "source": "part-27/example", "model": "local/agent", "scaffold": "minimal-agent", "tools": EXAMPLE_TOOLS, "messages": messages, "outcome": {"passed": passed, "stopped": stopped, "steps": len(calls) + 1, "turns": len(calls) + 1, "tokens": tokens, "seconds": seconds, "tool_errors": tool_errors}, "lossy": False, "collected_at": now(), } return [ episode("e-index-file", "Which file in the workspace describes the release process?", [("list_files", {"path": "."}, "notes.md\nrelease.md\nrunbook.md"), ("read_file", {"path": "release.md"}, "# Release\nTag the commit, then run the publish job.")], "release.md describes the release process.", True, "finished", 1840, 14.2), episode("e-owner", "Who is listed as the owner of the deploy job, and where do its logs go?", [("read_file", {"path": "runbook.md"}, f"Owner: platform team (contact {SAFE_EMAIL})\n" f"Gateway: {SAFE_ENDPOINT} with token {SAFE_FAKE_TOKEN}\n" f"Logs: {SAFE_HOME_PATH}, kept for 28 days")], "The platform team owns it and the logs are kept for 28 days.", True, "finished", 1210, 9.8), episode("e-long-way", "How many markdown files does the workspace contain?", [("list_files", {"path": "."}, "notes.md\nrelease.md\nrunbook.md"), ("read_file", {"path": "notes.md"}, "Scratch notes."), ("read_file", {"path": "release.md"}, "# Release"), ("read_file", {"path": "runbook.md"}, "# Runbook"), ("list_files", {"path": "."}, "notes.md\nrelease.md\nrunbook.md")], "Three.", True, "finished", 4900, 51.0), episode("e-missing", "What does the deployment checklist say about database migrations?", [("read_file", {"path": "checklist.md"}, "error: no such file inside the workspace: checklist.md")], "There is no deployment checklist in the workspace.", False, "finished", 980, 7.1, tool_errors=1), ] # -------------------------------------------------------------------------------------- # Main # -------------------------------------------------------------------------------------- def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--from", dest="source", choices=["part-24", "part-26", "jsonl"], help="which log shape to read") parser.add_argument("--labbook", default=None, help="part-24: the notebook minimal-agent.py appended to") parser.add_argument("--lab-name", default="part-24/minimal-agent", help="part-24: the lab field to select in the notebook") parser.add_argument("--results", default=None, help="part-26: the results file agent-eval.py wrote with --out") parser.add_argument("--input", default=None, help="jsonl: a file already in this part's episode shape") parser.add_argument("--tools", default=None, help="JSON file holding the OpenAI-style tool list the agent was given") parser.add_argument("--infer-tools", action="store_true", help="build a tool list from the observed calls; read the warning it prints") parser.add_argument("--system-prompt-file", default=None, help="text file holding the agent's system prompt, included in every episode") parser.add_argument("--keep-reasoning", action="store_true", help="keep the assistant's reasoning text as reasoning_content") parser.add_argument("--finish-tool", default="finish", help="name of the tool that ends a run; its call is rebuilt from the answer") parser.add_argument("--out", default="raw/trajectories.jsonl") parser.add_argument("--make-example", default=None, metavar="PATH", help="write four worked episodes to PATH and exit") args = parser.parse_args() if args.make_example: out = Path(args.make_example) out.parent.mkdir(parents=True, exist_ok=True) episodes = example_episodes() out.write_text("".join(json.dumps(e) + "\n" for e in episodes), encoding="utf-8") print(f"wrote {len(episodes)} example episode(s) to {out}") print("Two of them carry the kind of thing scrub-trajectories.py is looking for.") return if not args.source: parser.error("give --from part-24, --from part-26 or --from jsonl, or --make-example") tools: list[dict[str, Any]] = [] if args.tools: loaded = json.loads(Path(args.tools).read_text(encoding="utf-8")) tools = loaded["tools"] if isinstance(loaded, dict) and "tools" in loaded else loaded if not isinstance(tools, list): sys.exit(f"{args.tools} must hold a list of tool definitions") elif not args.infer_tools and args.source != "jsonl": sys.exit("pass --tools with the agent's real tool list, or --infer-tools and read " "the warning. Part 24's Toolbox.schemas() returns exactly this list.") system_prompt = None if args.system_prompt_file: system_prompt = Path(args.system_prompt_file).read_text(encoding="utf-8").strip() episodes: list[dict[str, Any]] = [] if args.source == "part-24": if not args.labbook: parser.error("--from part-24 needs --labbook") runs = [r for r in read_jsonl(Path(args.labbook)) if r.get("lab") == args.lab_name] if not runs: sys.exit(f"no records for lab {args.lab_name!r} in {args.labbook}") for run in runs: for result in run.get("results", []): episode = episode_from_part24(run, result, tools, system_prompt, args.keep_reasoning, args.finish_tool) if episode: episodes.append(episode) elif args.source == "part-26": if not args.results: parser.error("--from part-26 needs --results") payload = json.loads(Path(args.results).read_text(encoding="utf-8")) run = payload.get("run", {}) for row in payload.get("results", []): episode = episode_from_part26(run, row, tools, system_prompt) if episode: episodes.append(episode) if episodes: print("note: Part 26's trajectories record tool calls and results but not the " "assistant's own words, so every episode here is marked lossy.") else: if not args.input: parser.error("--from jsonl needs --input") for row in read_jsonl(Path(args.input)): if "messages" not in row: continue row.setdefault("tools", tools) row.setdefault("lossy", False) row.setdefault("collected_at", now()) episodes.append(row) if not episodes: sys.exit("no episodes were built; check the paths inside the log you passed") if args.infer_tools and not args.tools: tools = infer_tools(episodes) for episode in episodes: episode["tools"] = tools print(f"WARNING: inferred {len(tools)} tool definition(s) from the calls that were " "made. The descriptions are placeholders and they will be part of the " "training prompt. Replace them with the real schemas before you train " "anything you intend to serve.") out = Path(args.out) out.parent.mkdir(parents=True, exist_ok=True) out.write_text("".join(json.dumps(e) + "\n" for e in episodes), encoding="utf-8") passed = sum(1 for e in episodes if e["outcome"].get("passed") is True) lossy = sum(1 for e in episodes if e.get("lossy")) tool_calls = sum(len(m.get("tool_calls") or []) for e in episodes for m in e["messages"]) print(f"episodes: {len(episodes)}") print(f"outcome passed:{passed:>4}") print(f"lossy: {lossy:>4}") print(f"tool calls: {tool_calls:>4}") print(f"tools: {len(tools)}") print(f"written to {out}") print("\nNext: scrub-trajectories.py, before you read them and before you train on them.") if __name__ == "__main__": main()