| |
| """Synthetic finish-the-task + no-loop trajectories in pi tool format.""" |
| from __future__ import annotations |
|
|
| import json |
| import random |
| from pathlib import Path |
|
|
| from datasets import Dataset, load_dataset |
|
|
| SYS = """You are an expert coding assistant operating inside pi, a coding agent harness. You help users by reading files, executing commands, editing code, and writing new files. |
| |
| Guidelines: |
| - Be concise. Call a tool or give a short final answer and stop. |
| - Chain shell steps in one bash call. Do not loop ls/cd. |
| - When the required output exists, stop. |
| """ |
|
|
| TOOLS = json.dumps( |
| [ |
| { |
| "type": "function", |
| "function": { |
| "name": "bash", |
| "description": "Execute a bash command", |
| "parameters": { |
| "type": "object", |
| "properties": {"command": {"type": "string"}}, |
| "required": ["command"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "write", |
| "description": "Write a file", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "path": {"type": "string"}, |
| "content": {"type": "string"}, |
| }, |
| "required": ["path", "content"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "read", |
| "description": "Read a file", |
| "parameters": { |
| "type": "object", |
| "properties": {"path": {"type": "string"}}, |
| "required": ["path"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "ls", |
| "description": "List a directory", |
| "parameters": { |
| "type": "object", |
| "properties": {"path": {"type": "string"}}, |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "edit", |
| "description": "Edit a file", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "path": {"type": "string"}, |
| "edits": {"type": "array"}, |
| }, |
| "required": ["path", "edits"], |
| }, |
| }, |
| }, |
| ] |
| ) |
|
|
|
|
| def tc(i, name, args): |
| return { |
| "id": f"c{i}", |
| "type": "function", |
| "function": {"name": name, "arguments": json.dumps(args)}, |
| } |
|
|
|
|
| def asst(text, calls=None): |
| m = {"role": "assistant", "content": text} |
| if calls: |
| m["tool_calls"] = calls |
| return m |
|
|
|
|
| def tool(i, content): |
| return {"role": "tool", "tool_call_id": f"c{i}", "content": content} |
|
|
|
|
| def pack(msgs, src): |
| return {"messages": [{"role": "system", "content": SYS}] + msgs, "tools": TOOLS, "source": src} |
|
|
|
|
| def main(): |
| rng = random.Random(2) |
| rows = [] |
|
|
| for i in range(400): |
| fname = f"out_{i}.txt" |
| body = rng.choice(["ok", "done", "42", "hello", "pass"]) |
| rows.append( |
| pack( |
| [ |
| {"role": "user", "content": f"Write {body} to /output/{fname} and stop."}, |
| asst("Writing the file.", [tc(0, "bash", {"command": f"mkdir -p /output && printf '%s\\n' '{body}' > /output/{fname} && wc -c /output/{fname}"})]), |
| tool(0, f"4 /output/{fname}\n"), |
| asst(f"Wrote /output/{fname}."), |
| ], |
| "fin:write", |
| ) |
| ) |
|
|
| for i in range(300): |
| rows.append( |
| pack( |
| [ |
| {"role": "user", "content": "Explore the repo and set the version in pkg/__init__.py to 1.2.3."}, |
| asst("Listing once, then editing.", [tc(0, "bash", {"command": "ls && sed -n '1,20p' pkg/__init__.py"})]), |
| tool(0, "__version__ = '0.1.0'\n"), |
| asst( |
| "Updating the version.", |
| [ |
| tc( |
| 1, |
| "edit", |
| { |
| "path": "pkg/__init__.py", |
| "edits": [{"oldText": "__version__ = '0.1.0'", "newText": "__version__ = '1.2.3'"}], |
| }, |
| ) |
| ], |
| ), |
| tool(1, "Edited pkg/__init__.py"), |
| asst("Set __version__ to 1.2.3."), |
| ], |
| "fin:edit-once", |
| ) |
| ) |
|
|
| for i in range(250): |
| rows.append( |
| pack( |
| [ |
| {"role": "user", "content": "Run the tests. If they pass, stop. Do not keep listing files."}, |
| asst("Running tests.", [tc(0, "bash", {"command": "python -m pytest -q"})]), |
| tool(0, ".\n1 passed\n"), |
| asst("Tests passed."), |
| ], |
| "fin:tests-pass", |
| ) |
| ) |
|
|
| for i in range(250): |
| rows.append( |
| pack( |
| [ |
| {"role": "user", "content": "Fix add() so 2+2 is 4, then retest."}, |
| asst("Reading and testing.", [tc(0, "bash", {"command": "cat mathlib.py && python -m pytest -q"})]), |
| tool(0, "def add(a,b):\n return a+b+1\nF\nassert add(2,2)==4\n"), |
| asst( |
| "Fixing the off-by-one.", |
| [ |
| tc( |
| 1, |
| "edit", |
| { |
| "path": "mathlib.py", |
| "edits": [{"oldText": " return a+b+1", "newText": " return a+b"}], |
| }, |
| ) |
| ], |
| ), |
| tool(1, "Edited"), |
| asst("Re-running tests.", [tc(2, "bash", {"command": "python -m pytest -q"})]), |
| tool(2, ".\n1 passed\n"), |
| asst("Fixed add() and tests pass."), |
| ], |
| "fin:fix-retest", |
| ) |
| ) |
|
|
| |
| existing = load_dataset("/var/lib/agentptb-cache/tb1/tmp/sft-data/pi-sft-short", split="train") |
| idx = list(range(len(existing))) |
| rng.shuffle(idx) |
| for j in idx[:8000]: |
| rec = existing[j] |
| rows.append({"messages": rec["messages"], "tools": rec["tools"], "source": rec["source"]}) |
|
|
| rng.shuffle(rows) |
| out = Path("/var/lib/agentptb-cache/tb1/tmp/sft-data/pi-sft-finish") |
| out.mkdir(parents=True, exist_ok=True) |
| ds = Dataset.from_dict( |
| { |
| "messages": [r["messages"] for r in rows], |
| "tools": [r["tools"] if isinstance(r["tools"], str) else json.dumps(r["tools"]) for r in rows], |
| "source": [r["source"] for r in rows], |
| } |
| ) |
| ds.to_parquet(str(out / "train.parquet")) |
| print("wrote", len(rows), "to", out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
