qwen36_01/worklogs/generate_official_like_dataset.py

import argparse
import json
import random
from pathlib import Path


TOOL_NAMES = [
    "terminal_exec",
    "process_list",
    "read_file",
    "edit_file",
    "write_file",
    "glob_files",
    "grep_code",
    "web_search",
    "web_fetch",
    "browser_open",
    "mcp_codebase_memory_search",
    "mcp_codebase_memory_graph",
    "sessions_spawn",
    "sessions_send",
    "subagents_run",
    "memory_search",
]


SYSTEM_PROMPT = """You are Codex, a coding agent. You help the user modify, test, and reason about code.
Follow the available tools carefully. Prefer precise shell commands, inspect files before editing,
and return concise progress updates. When tool calls are needed, produce valid tool call arguments."""


COMMON_PROJECT_CONTEXT = """
Repository: /home/user/workspace/project
Task: implement features, inspect logs, run tests, patch code, summarize results.
Conventions:
- Use rg for search.
- Avoid destructive commands.
- Keep a worklog.
- Preserve unrelated user changes.
"""


CODE_SNIPPET = """
def fetch_prices(symbols, start_date, end_date, retries=3):
    results = []
    for symbol in symbols:
        payload = {
            "symbol": symbol,
            "start": start_date,
            "end": end_date,
            "adjust": "qfq",
        }
        # retry network request with backoff
        for attempt in range(retries):
            try:
                results.append(client.get("/prices", params=payload))
                break
            except TimeoutError:
                time.sleep(1.5 * (attempt + 1))
    return results
"""


def make_tool(name):
    if "terminal" in name or "process" in name:
        props = {"command": {"type": "string", "description": "Shell command to execute."}}
        required = ["command"]
    elif "read" in name:
        props = {"path": {"type": "string"}}
        required = ["path"]
    elif "edit" in name:
        props = {
            "file_path": {"type": "string"},
            "old_string": {"type": "string"},
            "new_string": {"type": "string"},
        }
        required = ["file_path", "old_string", "new_string"]
    elif "web" in name:
        props = {"query": {"type": "string"}, "count": {"type": "integer"}}
        required = ["query"]
    else:
        props = {"input": {"type": "string"}, "limit": {"type": "integer"}}
        required = ["input"]
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": f"Synthetic benchmark tool: {name}",
            "parameters": {
                "type": "object",
                "properties": props,
                "required": required,
            },
        },
    }


def rand_text(rng, approx_tokens):
    # Mostly ASCII/code-like text keeps the chars/token ratio simple enough for
    # synthetic load generation. This is approximate; server usage is authoritative.
    unit = (
        "tool output line: status=ok path=/home/user/project/src/module.py "
        "grep result includes function names, stack frames, JSON fields, "
        "中文说明：这里包含工具返回、日志片段、代码上下文和转义字符。 "
        + CODE_SNIPPET.replace("\n", "\\n")
        + "\n"
    )
    chars = max(1, approx_tokens * 4)
    return (unit * (chars // len(unit) + 1))[:chars]


def make_tool_call(rng, idx):
    name = rng.choice(TOOL_NAMES)
    if "terminal" in name:
        args = {
            "command": "cd /home/user/project && rg -n \"TODO|FIXME|error\" src tests 2>/dev/null | head -50"
        }
    elif "read" in name:
        args = {"path": "/home/user/workspace/.agents/skills/data-analysis/SKILL.md"}
    elif "edit" in name:
        args = {
            "file_path": "/home/user/scripts/fetch_prices.py",
            "old_string": "# 全局频率限制器实例\n_rate_limiter = RateLimiter(min_interval=2.0, max_per_minute=20)",
            "new_string": "# 全局频率限制器实例 - 根据配额上调\n_rate_limiter = RateLimiter(min_interval=1.5, max_per_minute=40)",
        }
    elif "web" in name:
        args = {"query": "深圳 周末活动 推荐 2026年7月", "count": 10}
    else:
        args = {"input": "inspect project memory and summarize related files", "limit": 20}
    return {
        "id": f"call_{idx}_{rng.randrange(10**8)}",
        "type": "function",
        "function": {"name": name, "arguments": json.dumps(args, ensure_ascii=False)},
    }


def target_prompt_tokens(rng, max_prompt_tokens=None):
    # Approximate official distribution: many 4K-16K, meaningful 32K+ tail,
    # small number of 128K+ requests.
    x = rng.random()
    if x < 0.08:
        value = rng.randint(1000, 4000)
        return min(value, max_prompt_tokens) if max_prompt_tokens else value
    if x < 0.50:
        value = rng.randint(4000, 16000)
        return min(value, max_prompt_tokens) if max_prompt_tokens else value
    if x < 0.74:
        value = rng.randint(16000, 32000)
        return min(value, max_prompt_tokens) if max_prompt_tokens else value
    if x < 0.95:
        high = min(128000, max_prompt_tokens) if max_prompt_tokens else 128000
        value = rng.randint(32000, max(32000, high))
        return value
    high = min(235000, max_prompt_tokens) if max_prompt_tokens else 235000
    return rng.randint(128000, max(128000, high)) if high >= 128000 else high


def target_output_tokens(rng):
    x = rng.random()
    if x < 0.66:
        return rng.randint(64, 256)
    if x < 0.96:
        return rng.randint(256, 2000)
    return rng.randint(2000, 8192)


def build_request(rng, req_id, session_id, shared_prefix_tokens, total_prompt_tokens):
    tool_count = max(1, int(rng.lognormvariate(2.8, 0.8)))
    tool_count = min(tool_count, 92)
    tools = [make_tool(rng.choice(TOOL_NAMES)) for _ in range(tool_count)]

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT + "\n" + COMMON_PROJECT_CONTEXT + rand_text(rng, shared_prefix_tokens),
        },
    ]

    remaining = max(512, total_prompt_tokens - shared_prefix_tokens)
    turns = rng.randint(12, 65)
    per_turn = max(32, remaining // max(1, turns))
    call_idx = 0
    for turn in range(turns):
        messages.append({
            "role": "user",
            "content": (
                f"Session {session_id}, turn {turn}: inspect the codebase and continue the task. "
                "Need exact commands, possible edits, and concise next action."
            ),
        })
        if rng.random() < 0.78:
            calls = [make_tool_call(rng, call_idx)]
            call_idx += 1
            messages.append({
                "role": "assistant",
                # This runtime validates that every message has content or
                # reasoning_content, even when assistant emits tool_calls.
                "content": "",
                "tool_calls": calls,
            })
            tool_len = int(per_turn * rng.uniform(0.5, 2.0))
            if rng.random() > 0.99:
                tool_len = rng.randint(12000, 50000)
            elif rng.random() > 0.90:
                tool_len = rng.randint(2500, 16000)
            messages.append({
                "role": "tool",
                "tool_call_id": calls[0]["id"],
                "name": calls[0]["function"]["name"],
                "content": rand_text(rng, tool_len),
            })
        else:
            messages.append({
                "role": "assistant",
                "content": "我会先检查相关文件和日志，再给出下一步修改建议。",
            })

    messages.append({
        "role": "user",
        "content": "基于以上工具结果，继续完成当前任务。需要时发起下一步 tool_call。",
    })

    return {
        "id": f"synthetic_{req_id:04d}",
        "model": "llm",
        "messages": messages,
        "tools": tools,
        "tool_choice": "auto",
        "stream": True,
        "temperature": 0,
        "max_tokens": target_output_tokens(rng),
        "metadata": {
            "session_id": session_id,
            "target_prompt_tokens_approx": total_prompt_tokens,
            "shared_prefix_tokens_approx": shared_prefix_tokens,
        },
    }


def estimate_tokens(messages, tools):
    chars = 0
    for msg in messages:
        chars += len(msg.get("role") or "") + len(msg.get("content") or "")
        if msg.get("tool_calls"):
            chars += len(json.dumps(msg["tool_calls"], ensure_ascii=False))
        if msg.get("name"):
            chars += len(msg["name"])
    chars += len(json.dumps(tools, ensure_ascii=False))
    return max(1, chars // 4)


def make_history_turn(rng, state, turn_id, approx_tool_tokens):
    user = {
        "role": "user",
        "content": (
            f"Continue session {state['session_id']} turn {turn_id}. "
            "Inspect files, reason about logs, and call the next tool if needed."
        ),
    }
    call = make_tool_call(rng, turn_id)
    assistant = {
        "role": "assistant",
        "content": "",
        "tool_calls": [call],
    }
    tool = {
        "role": "tool",
        "tool_call_id": call["id"],
        "name": call["function"]["name"],
        "content": rand_text(rng, approx_tool_tokens),
    }
    return [user, assistant, tool]


def generate_cumulative_requests(rng, count, session_count, max_prompt_tokens):
    """Generate requests whose prefixes are actually reusable.

    The first version of this generator only stored an intended cache ratio in
    metadata. This mode makes the byte/text prefix identical across requests by
    sharing tools/system content and by growing each session history in place.
    """
    common_tools = [make_tool(name) for name in TOOL_NAMES]
    # Keep a large, fixed global prefix. It imitates common system prompt,
    # agent framework instructions, and mounted tool schemas.
    common_system = {
        "role": "system",
        "content": SYSTEM_PROMPT + "\n" + COMMON_PROJECT_CONTEXT + rand_text(rng, 6000),
    }
    sessions = {}
    session_order = []
    requests = []

    def get_session(sid):
        if sid not in sessions:
            sessions[sid] = {
                "session_id": sid,
                "messages": [common_system],
                "turn_id": 0,
            }
            session_order.append(sid)
        return sessions[sid]

    recent = []
    for req_id in range(count):
        # For small benchmark sets, force meaningful reuse. For larger sets,
        # this still creates short-distance affinity similar to the official
        # median gap of a few requests.
        if recent and rng.random() < 0.65:
            sid = rng.choice(recent[-32:])
        else:
            sid = f"sess_{len(session_order) % max(1, session_count):04d}"
        recent.append(sid)
        state = get_session(sid)
        target = target_prompt_tokens(rng, max_prompt_tokens)

        # Grow the historical context until this request has the target scale.
        while estimate_tokens(state["messages"], common_tools) < max(1024, target - 256):
            remaining = target - estimate_tokens(state["messages"], common_tools)
            chunk = min(max(256, remaining), rng.randint(800, 3500))
            state["messages"].extend(make_history_turn(rng, state, state["turn_id"], chunk))
            state["turn_id"] += 1

        current_user = {
            "role": "user",
            "content": "基于以上工具结果，继续完成当前任务。需要时发起下一步 tool_call。",
        }
        req_messages = list(state["messages"]) + [current_user]
        requests.append({
            "id": f"synthetic_cumulative_{req_id:04d}",
            "model": "llm",
            "messages": req_messages,
            "tools": common_tools,
            "tool_choice": "auto",
            "stream": True,
            "temperature": 0,
            "max_tokens": target_output_tokens(rng),
            "metadata": {
                "session_id": sid,
                "target_prompt_tokens_approx": estimate_tokens(req_messages, common_tools),
            },
        })

        # Pretend the model made a tool call and the environment returned a
        # result, so the next request in this session contains this full prefix.
        state["messages"].append(current_user)
        state["messages"].extend(make_history_turn(rng, state, state["turn_id"], rng.randint(500, 2500))[1:])
        state["turn_id"] += 1

    return requests


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", required=True)
    parser.add_argument("--requests", type=int, default=881)
    parser.add_argument("--sessions", type=int, default=768)
    parser.add_argument("--seed", type=int, default=20260714)
    parser.add_argument("--jsonl", action="store_true")
    parser.add_argument("--cumulative", action="store_true", help="Generate exact reusable session prefixes.")
    parser.add_argument(
        "--max-prompt-tokens",
        type=int,
        default=None,
        help="Cap synthetic prompt-token target. Use 90000 for a server started with --max-model-len 100000.",
    )
    args = parser.parse_args()

    rng = random.Random(args.seed)
    if args.cumulative:
        requests = generate_cumulative_requests(rng, args.requests, args.sessions, args.max_prompt_tokens)
    else:
        session_ids = [f"sess_{i:04d}" for i in range(args.sessions)]
        # Bias toward near-neighbor session reuse by reusing some recent sessions.
        requests = []
        recent = []
        for i in range(args.requests):
            if recent and rng.random() < 0.22:
                session_id = rng.choice(recent[-32:])
            else:
                session_id = session_ids[i % len(session_ids)]
            recent.append(session_id)
            total = target_prompt_tokens(rng, args.max_prompt_tokens)
            # Official weighted cache hit is about 65.6%; request-level varies.
            if total >= 128000:
                cache_ratio = rng.uniform(0.1, 0.35)
            elif total >= 16000:
                cache_ratio = rng.uniform(0.55, 0.72)
            else:
                cache_ratio = rng.uniform(0.75, 0.92)
            shared = int(total * cache_ratio)
            requests.append(build_request(rng, i, session_id, shared, total))

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    if args.jsonl:
        out.write_text("\n".join(json.dumps(x, ensure_ascii=False) for x in requests) + "\n", encoding="utf-8")
    else:
        out.write_text(json.dumps(requests, ensure_ascii=False, indent=2), encoding="utf-8")

    total_prompt = sum(x["metadata"]["target_prompt_tokens_approx"] for x in requests)
    total_shared = sum(x["metadata"].get("shared_prefix_tokens_approx", 0) for x in requests)
    print(json.dumps({
        "out": str(out),
        "requests": len(requests),
        "sessions": len({x["metadata"]["session_id"] for x in requests}),
        "approx_prompt_tokens_avg": total_prompt / len(requests),
        "approx_cache_ratio_weighted": total_shared / total_prompt if total_shared else None,
        "cumulative": args.cumulative,
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()