#!/usr/bin/env python3 """Convert public mini-swe-agent trajectories to Pi/Qwen structured tool chats.""" from __future__ import annotations import argparse import json import re from pathlib import Path from datasets import Dataset, load_dataset SOURCE = "Kwai-Klear/SWE-smith-mini_swe_agent_plus-trajectories-66k" BASH_BLOCK = re.compile(r"```(?:bash|sh)\s*\n(.*?)\n```", re.DOTALL | re.IGNORECASE) FINAL_MARKER = "MINI_SWE_AGENT_FINAL_OUTPUT" PI_SYSTEM_PROMPT = """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. Available tools: - read: Read file contents - bash: Execute bash commands - edit: Make precise file edits with exact text replacement - write: Create or overwrite files Guidelines: - Inspect the repository and relevant instructions before changing code. - Use tools to make the requested change; do not merely describe commands. - Prefer focused searches and reads over dumping large files. - Make minimal, targeted edits and preserve unrelated behavior. - Run relevant focused tests after editing. - When the task is complete, stop using tools and give a concise final summary. - Be concise in your responses. - Show file paths clearly when working with files.""" PI_TOOLS = [ { "type": "function", "function": { "name": "read", "description": "Read the contents of a file.", "parameters": { "type": "object", "properties": { "path": {"type": "string"}, "offset": {"type": "number"}, "limit": {"type": "number"}, }, "required": ["path"], }, }, }, { "type": "function", "function": { "name": "bash", "description": "Execute a bash command in the current working directory.", "parameters": { "type": "object", "properties": { "command": {"type": "string"}, "timeout": {"type": "number"}, }, "required": ["command"], }, }, }, { "type": "function", "function": { "name": "edit", "description": "Edit one file using exact text replacement.", "parameters": { "type": "object", "properties": { "path": {"type": "string"}, "edits": { "type": "array", "items": { "type": "object", "properties": { "oldText": {"type": "string"}, "newText": {"type": "string"}, }, "required": ["oldText", "newText"], }, }, }, "required": ["path", "edits"], }, }, }, { "type": "function", "function": { "name": "write", "description": "Create or overwrite a file.", "parameters": { "type": "object", "properties": { "path": {"type": "string"}, "content": {"type": "string"}, }, "required": ["path", "content"], }, }, }, ] TOOLS_JSON = json.dumps(PI_TOOLS, separators=(",", ":")) def message( role: str, content: str, *, reasoning_content: str | None = None, tool_calls: list[dict] | None = None, tool_call_id: str | None = None, name: str | None = None, ) -> dict: return { "role": role, "content": content, "reasoning_content": reasoning_content, "tool_calls": tool_calls, "tool_call_id": tool_call_id, "name": name, } def clean_reasoning(text: str) -> str: text = text.strip() text = re.sub(r"^THOUGHT:\s*", "", text, flags=re.IGNORECASE) return text.strip() def convert(example: dict) -> dict: source_messages = example["messages"] converted = [message("system", PI_SYSTEM_PROMPT)] pending_call_id: str | None = None call_index = 0 initial_user_seen = False malformed_actions = 0 for source in source_messages: role = source["role"] content = source.get("content") or "" if role == "system": continue if role == "user": if pending_call_id is not None: converted.append( message( "tool", content, tool_call_id=pending_call_id, name="bash", ) ) pending_call_id = None else: converted.append(message("user", content)) initial_user_seen = True continue if role != "assistant": continue match = BASH_BLOCK.search(content) is_final = FINAL_MARKER in content if is_final: converted.append( message( "assistant", "Implemented the requested change and verified it with focused checks.", ) ) pending_call_id = None continue if match is None: malformed_actions += 1 converted.append(message("assistant", content.strip())) pending_call_id = None continue reasoning = clean_reasoning(content[: match.start()]) command = match.group(1).strip() call_index += 1 pending_call_id = f"call_{call_index:04d}" converted.append( message( "assistant", "", reasoning_content=reasoning, tool_calls=[ { "id": pending_call_id, "type": "function", "function": { "name": "bash", "arguments": json.dumps( {"command": command}, separators=(",", ":") ), }, } ], ) ) if not initial_user_seen: raise ValueError(f"trajectory {example['instance_id']} has no user request") if converted[-1]["role"] != "assistant" or converted[-1]["tool_calls"]: converted.append( message( "assistant", "Implemented the requested change and verified it with focused checks.", ) ) return { "instance_id": example["instance_id"], "messages": converted, "tools": TOOLS_JSON, "source": SOURCE, "num_tool_calls": call_index, "malformed_actions": malformed_actions, } def evaluation_instance_ids(path: Path) -> set[str]: """Exact identifiers only, used solely as a decontamination blacklist.""" if not path.exists(): return set() return {entry.name for entry in path.iterdir() if entry.is_dir()} def write_shards(dataset: Dataset, output_dir: Path, shards: int) -> None: output_dir.mkdir(parents=True, exist_ok=True) for index in range(shards): shard = dataset.shard(num_shards=shards, index=index, contiguous=True) path = output_dir / f"train-{index:05d}-of-{shards:05d}.parquet" shard.to_parquet(path) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--output", type=Path, required=True) parser.add_argument("--limit", type=int) parser.add_argument("--num-proc", type=int, default=32) parser.add_argument("--shards", type=int, default=16) parser.add_argument( "--eval-id-dir", type=Path, default=Path("/root/work/shared/tasksets/swe-bench-verified"), ) args = parser.parse_args() split = "train" if args.limit is None else f"train[:{args.limit}]" raw = load_dataset(SOURCE, split=split) blacklist = evaluation_instance_ids(args.eval_id_dir) before = len(raw) raw = raw.filter( lambda instance_id: instance_id not in blacklist, input_columns=["instance_id"], num_proc=args.num_proc, desc="Removing exact evaluation identifiers", ) removed = before - len(raw) converted = raw.map( convert, remove_columns=raw.column_names, num_proc=args.num_proc, desc="Converting trajectories to Pi structured calls", ) write_shards(converted, args.output, min(args.shards, len(converted))) report = { "source": SOURCE, "source_revision": "750b2c11239fd5e32f97e6cfb9bf80fb9a9a2983", "source_rows": before, "rows_removed_as_exact_eval_ids": removed, "output_rows": len(converted), "tool_calls": sum(converted["num_tool_calls"]), "malformed_actions": sum(converted["malformed_actions"]), "transformation": "fenced bash -> Pi structured bash calls; terminal marker -> concise final response", } (args.output / "provenance.json").write_text(json.dumps(report, indent=2) + "\n") print(json.dumps(report, indent=2)) if __name__ == "__main__": main()