414 lines
15 KiB
Python
414 lines
15 KiB
Python
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()
|