| """Resumably extend Pi's tool-routing dataset with harder parser examples.""" |
|
|
| from collections import Counter |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from pathlib import Path |
| from typing import Any |
| import json |
| import os |
| import re |
| import time |
|
|
| from codex_client import CodexClientPool |
|
|
|
|
| ROOT = Path(__file__).parent |
| TOOLS = json.loads((ROOT / "pi_tools.json").read_text(encoding="utf-8")) |
| SEED_OUTPUT = ROOT / "pi_tools_1200.jsonl" |
| OUTPUT = ROOT / "pi_tools_1500.jsonl" |
| PROGRESS = ROOT / "pi_tools_1500.progress.json" |
| TARGETS = {"read": 350, "bash": 450, "edit": 350, "write": 350} |
| TOTAL_TARGET = sum(TARGETS.values()) |
| BATCH_SIZE = 25 |
| MAX_WORKERS = 4 |
| TOOL_MAP = {tool["name"]: tool for tool in TOOLS} |
| TOOLS_TEXT = json.dumps(TOOLS, separators=(",", ":"), ensure_ascii=False) |
|
|
| THEMES = [ |
| "short direct requests", |
| "absolute and relative paths", |
| "spaces, quotes, unicode, and punctuation", |
| "JavaScript and TypeScript projects", |
| "Python projects", |
| "shell pipelines and redirects", |
| "configuration and environment files", |
| "multiline source code", |
| "logs, offsets, and limits", |
| "nested directories and dotfiles", |
| "URLs contrasted with local paths", |
| "realistic debugging and maintenance tasks", |
| ] |
|
|
|
|
| def _strip_fence(text: str) -> str: |
| text = text.strip() |
| if text.startswith("```"): |
| text = text.split("\n", 1)[1] if "\n" in text else text[3:] |
| if text.endswith("```"): |
| text = text[:-3] |
| return text.strip() |
|
|
|
|
| def _query_key(query: str) -> str: |
| return re.sub(r"\s+", " ", query.strip().casefold()) |
|
|
|
|
| def _load_existing() -> list[dict[str, Any]]: |
| source = OUTPUT if OUTPUT.exists() else SEED_OUTPUT |
| if not source.exists(): |
| return [] |
| records = [] |
| for line_number, line in enumerate(source.read_text(encoding="utf-8").splitlines(), 1): |
| if not line.strip(): |
| continue |
| try: |
| record = json.loads(line) |
| answers = json.loads(record["answers"]) |
| name = answers[0]["name"] |
| if name not in TOOL_MAP: |
| raise ValueError(f"unknown tool {name}") |
| records.append(record) |
| except Exception as error: |
| raise ValueError(f"invalid existing dataset line {line_number}: {error}") from error |
| return records |
|
|
|
|
| def _leaf_values(value: Any): |
| if isinstance(value, dict): |
| for child in value.values(): |
| yield from _leaf_values(child) |
| elif isinstance(value, list): |
| for child in value: |
| yield from _leaf_values(child) |
| elif value is not None: |
| yield value |
|
|
|
|
| def _valid_example(example: Any, expected_tool: str) -> tuple[bool, str]: |
| if not isinstance(example, dict): |
| return False, "not an object" |
| query = example.get("query") |
| answers = example.get("answers") |
| if not isinstance(query, str) or not query.strip(): |
| return False, "missing query" |
| if not isinstance(answers, list) or len(answers) != 1: |
| return False, "answer count" |
| call = answers[0] |
| if not isinstance(call, dict) or call.get("name") != expected_tool: |
| return False, "wrong tool" |
| arguments = call.get("arguments") |
| if not isinstance(arguments, dict): |
| return False, "arguments" |
|
|
| schema = TOOL_MAP[expected_tool]["parameters"] |
| if set(arguments) - set(schema): |
| return False, "unknown argument" |
| required = {name for name, info in schema.items() if info.get("required")} |
| if required - set(arguments): |
| return False, "missing argument" |
|
|
| for name, value in arguments.items(): |
| expected_type = schema[name]["type"] |
| if expected_type == "string" and not isinstance(value, str): |
| return False, "string type" |
| if expected_type == "number" and (not isinstance(value, (int, float)) or isinstance(value, bool)): |
| return False, "number type" |
| if expected_type == "array" and not isinstance(value, list): |
| return False, "array type" |
|
|
| if expected_tool == "edit": |
| edits = arguments.get("edits") |
| if not edits or not all( |
| isinstance(edit, dict) |
| and set(edit) == {"oldText", "newText"} |
| and isinstance(edit["oldText"], str) |
| and isinstance(edit["newText"], str) |
| for edit in edits |
| ): |
| return False, "edit shape" |
|
|
| |
| |
| for value in _leaf_values(arguments): |
| rendered = str(value) |
| if rendered and rendered not in query: |
| return False, f"ungrounded value: {rendered[:40]}" |
| return True, "" |
|
|
|
|
| def _prompt(tool_name: str, count: int, batch_number: int, avoid: list[str], current_count: int) -> str: |
| theme = THEMES[(batch_number - 1) % len(THEMES)] |
| if tool_name == "edit": |
| theme = "short single-line exact replacements with diverse file types and syntax" |
| if tool_name == "bash" and current_count < 400: |
| theme = """high-complexity shell-command parsing. Invent realistic setup or maintenance scenarios such as installing X, configuring Y, launching a service, building a project, transforming files, inspecting a system, or chaining arbitrary shell steps. The query must already contain the exact complete command; the router only extracts it. Use varied quoting, flags, environment variables, pipes, redirects, subshells, and multiline commands. Do not require the router to synthesize specialist commands from vague goals""" |
| elif current_count >= 300: |
| theme = """noisy or mildly ambiguous planner messages that still contain one unambiguous tool action and every exact argument. Add natural hesitations, corrections, irrelevant context, or conversational clutter, but keep the intended call objectively recoverable""" |
| avoid_text = "\n".join(f"- {query}" for query in avoid[-30:]) or "- none" |
| tool = TOOL_MAP[tool_name] |
| return f"""Generate exactly {count} training examples for a coding-agent tool router. |
| |
| ALL AVAILABLE TOOLS: |
| {TOOLS_TEXT} |
| |
| This batch must call only `{tool_name}`: |
| {json.dumps(tool, separators=(',', ':'), ensure_ascii=False)} |
| |
| Batch theme: {theme} |
| |
| Rules: |
| - Return one valid JSON array and nothing else. |
| - Every item is exactly {{"query":"...","answers":[{{"name":"{tool_name}","arguments":{{...}}}}]}}. |
| - Exactly one call per example and exactly {count} examples. |
| - The query is a short, explicit planner intent, not a casual end-user request. |
| - Include every argument value literally and exactly in the query. Never infer, shorten, repair, normalize, or invent a value. |
| - Preserve shell commands, paths, file content, oldText, newText, quoting, escapes, whitespace, and newlines exactly. |
| - Use only declared parameters and include every required parameter. |
| - Make examples meaningfully different, not template-like substitutions. |
| - For `bash`, `command` must contain the complete executable command, not merely a URL, filename, or description. |
| - For `read`, paths are filesystem paths; never use a web URL as a path. |
| - For `edit`, `edits` is a non-empty array of objects with only `oldText` and `newText` strings. |
| - For `edit`, keep every oldText and newText on one line, and place each exact value directly in the query without Markdown escaping or a code fence. |
| - For `write`, `content` is the complete file content and appears verbatim in the query. |
| - Noise may surround an action, but never make the expected arguments unknowable or require specialist domain knowledge. |
| |
| Do not repeat these earlier queries: |
| {avoid_text} |
| """ |
|
|
|
|
| def _generate_batch(tool_name: str, count: int, batch_number: int, avoid: list[str], current_count: int): |
| response = CodexClientPool().get().models.generate_content( |
| model="codex-cli", |
| contents=_prompt(tool_name, count, batch_number, avoid, current_count), |
| config={"temperature": 0.8, "max_output_tokens": 16384}, |
| ) |
| parsed = json.loads(_strip_fence(response.text)) |
| if not isinstance(parsed, list): |
| raise ValueError("response is not an array") |
| return tool_name, parsed |
|
|
|
|
| def _persist(records: list[dict[str, Any]], failures: dict[str, int]) -> Counter: |
| counts = Counter() |
| lines = [] |
| for record in records: |
| name = json.loads(record["answers"])[0]["name"] |
| counts[name] += 1 |
| lines.append(json.dumps(record, ensure_ascii=False)) |
|
|
| output_tmp = OUTPUT.with_suffix(OUTPUT.suffix + ".tmp") |
| output_tmp.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8") |
| os.replace(output_tmp, OUTPUT) |
|
|
| state = { |
| "targets": TARGETS, |
| "total": len(records), |
| "per_tool": {name: counts[name] for name in TOOL_MAP}, |
| "failures": failures, |
| "updated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), |
| "complete": all(counts[name] >= TARGETS[name] for name in TOOL_MAP), |
| } |
| progress_tmp = PROGRESS.with_suffix(PROGRESS.suffix + ".tmp") |
| progress_tmp.write_text(json.dumps(state, indent=2) + "\n", encoding="utf-8") |
| os.replace(progress_tmp, PROGRESS) |
| return counts |
|
|
|
|
| def main() -> None: |
| records = _load_existing() |
| seen = {_query_key(record["query"]) for record in records} |
| failures = {name: 0 for name in TOOL_MAP} |
| counts = _persist(records, failures) |
| print(f"Resuming at {len(records)}/{TOTAL_TARGET}: {dict(counts)}", flush=True) |
|
|
| round_number = 0 |
| while any(counts[name] < TARGETS[name] for name in TOOL_MAP): |
| round_number += 1 |
| jobs = {} |
| with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool: |
| for name in TOOL_MAP: |
| remaining = TARGETS[name] - counts[name] |
| if remaining <= 0: |
| continue |
| count = min(BATCH_SIZE, remaining) |
| prior = [ |
| record["query"] for record in records |
| if json.loads(record["answers"])[0]["name"] == name |
| ] |
| batch_number = (counts[name] // BATCH_SIZE) + 1 |
| future = pool.submit(_generate_batch, name, count, batch_number, prior, counts[name]) |
| jobs[future] = name |
|
|
| for future in as_completed(jobs): |
| scheduled_name = jobs[future] |
| try: |
| name, examples = future.result() |
| accepted = 0 |
| rejected = Counter() |
| remaining = TARGETS[name] - counts[name] |
| for example in examples: |
| valid, reason = _valid_example(example, name) |
| key = _query_key(example.get("query", "")) if isinstance(example, dict) else "" |
| if not valid: |
| rejected[reason] += 1 |
| continue |
| if key in seen: |
| rejected["duplicate"] += 1 |
| continue |
| if accepted >= remaining: |
| break |
| answers = example["answers"] |
| records.append({ |
| "query": example["query"], |
| "tools": TOOLS_TEXT, |
| "answers": json.dumps(answers, separators=(",", ":"), ensure_ascii=False), |
| "source": "synth-codex-cli", |
| "model": "codex-cli", |
| }) |
| seen.add(key) |
| accepted += 1 |
| failures[name] = 0 if accepted else failures[name] + 1 |
| counts = _persist(records, failures) |
| print( |
| f"{name}: +{accepted}, rejected={dict(rejected)}, " |
| f"progress={counts[name]}/{TARGETS[name]}, total={len(records)}/{TOTAL_TARGET}", |
| flush=True, |
| ) |
| except Exception as error: |
| failures[scheduled_name] += 1 |
| _persist(records, failures) |
| print(f"{scheduled_name} batch failed (progress preserved): {error}", flush=True) |
|
|
| if max(failures.values()) >= 5: |
| raise RuntimeError(f"five consecutive failures; rerun to resume: {failures}") |
|
|
| counts = _persist(records, failures) |
| print(json.dumps({"output": str(OUTPUT), "total": len(records), "per_tool": dict(counts)}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
