163 lines
5.8 KiB
Python
163 lines
5.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Extract a prompt-bearing routing fixture while emitting a prompt-free manifest."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--trace", type=Path, required=True)
|
|
parser.add_argument("--row-ids", type=int, nargs="+", required=True)
|
|
parser.add_argument("--parent-of", type=int, nargs="*", default=[])
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--manifest", type=Path, required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
def common_prefix(left: list[Any], right: list[Any]) -> int:
|
|
count = 0
|
|
for lhs, rhs in zip(left, right):
|
|
if lhs != rhs:
|
|
break
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
target_ids = set(args.row_ids)
|
|
parent_targets = set(args.parent_of)
|
|
if not parent_targets.issubset(target_ids):
|
|
raise SystemExit("--parent-of must be a subset of --row-ids")
|
|
|
|
trace_digest = hashlib.sha256()
|
|
by_row: dict[int, dict[str, Any]] = {}
|
|
by_chat_id: dict[str, dict[str, Any]] = {}
|
|
with args.trace.open("rb") as handle:
|
|
row_id = 0
|
|
while True:
|
|
offset = handle.tell()
|
|
line = handle.readline()
|
|
if not line:
|
|
break
|
|
trace_digest.update(line)
|
|
row = json.loads(line)
|
|
meta = {
|
|
"row_id": row_id,
|
|
"offset": offset,
|
|
"chat_id": str(row.get("chat_id")),
|
|
"parent_chat_id": str(row.get("parent_chat_id")),
|
|
"turn": int(row.get("turn", 1)),
|
|
"input_length": int(row["input_length"]),
|
|
"output_length": int(row["output_length"]),
|
|
"hash_ids": row.get("hash_ids") or [],
|
|
}
|
|
by_chat_id[meta["chat_id"]] = meta
|
|
if row_id in target_ids:
|
|
by_row[row_id] = meta
|
|
row_id += 1
|
|
|
|
missing = target_ids - set(by_row)
|
|
if missing:
|
|
raise SystemExit(f"missing target row IDs: {sorted(missing)}")
|
|
parent_rows: list[dict[str, Any]] = []
|
|
for row_id in args.parent_of:
|
|
parent_id = by_row[row_id]["parent_chat_id"]
|
|
parent = by_chat_id.get(parent_id)
|
|
if parent is None:
|
|
raise SystemExit(f"row {row_id} parent chat {parent_id} is absent")
|
|
parent_rows.append(parent)
|
|
|
|
# Put parents first so online prefix caching can materialize their shared
|
|
# blocks before descendants are admitted.
|
|
ordered_meta = parent_rows + [by_row[row_id] for row_id in args.row_ids]
|
|
if len({row["row_id"] for row in ordered_meta}) != len(ordered_meta):
|
|
raise SystemExit("fixture rows must be unique")
|
|
|
|
output_rows: list[dict[str, Any]] = []
|
|
with args.trace.open("rb") as handle:
|
|
for fixture_index, meta in enumerate(ordered_meta):
|
|
handle.seek(meta["offset"])
|
|
source = json.loads(handle.readline())
|
|
prompt = source.get("prompt")
|
|
if not isinstance(prompt, str) or not prompt:
|
|
raise SystemExit(f"row {meta['row_id']} has no prompt")
|
|
output_rows.append(
|
|
{
|
|
"fixture_index": fixture_index,
|
|
"row_id": str(meta["row_id"]),
|
|
"prompt": prompt,
|
|
"prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
|
|
"input_length": meta["input_length"],
|
|
"output_length": meta["output_length"],
|
|
"turn": meta["turn"],
|
|
"chat_id": meta["chat_id"],
|
|
"parent_chat_id": meta["parent_chat_id"],
|
|
"hash_ids": meta["hash_ids"],
|
|
}
|
|
)
|
|
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
with args.output.open("w") as handle:
|
|
for row in output_rows:
|
|
handle.write(json.dumps(row, sort_keys=True) + "\n")
|
|
|
|
pair_coverage = []
|
|
by_chat = {row["chat_id"]: row for row in output_rows}
|
|
for child in output_rows:
|
|
parent = by_chat.get(child["parent_chat_id"])
|
|
if parent is None:
|
|
continue
|
|
pair_coverage.append(
|
|
{
|
|
"parent_row_id": parent["row_id"],
|
|
"child_row_id": child["row_id"],
|
|
"parent_turn": parent["turn"],
|
|
"child_turn": child["turn"],
|
|
"trace_hash_common_prefix_blocks": common_prefix(
|
|
parent["hash_ids"], child["hash_ids"]
|
|
),
|
|
}
|
|
)
|
|
fixture_digest = hashlib.sha256(args.output.read_bytes()).hexdigest()
|
|
manifest = {
|
|
"schema_version": "qwen30_trace_routing_fixture.v1",
|
|
"source_trace": str(args.trace.resolve()),
|
|
"source_trace_sha256": trace_digest.hexdigest(),
|
|
"fixture_path": str(args.output.resolve()),
|
|
"fixture_sha256": fixture_digest,
|
|
"contains_prompt_text": False,
|
|
"request_count": len(output_rows),
|
|
"rows": [
|
|
{
|
|
key: row[key]
|
|
for key in (
|
|
"fixture_index",
|
|
"row_id",
|
|
"prompt_sha256",
|
|
"input_length",
|
|
"output_length",
|
|
"turn",
|
|
"chat_id",
|
|
"parent_chat_id",
|
|
)
|
|
}
|
|
| {"trace_hash_blocks": len(row["hash_ids"])}
|
|
for row in output_rows
|
|
],
|
|
"prefix_pairs": pair_coverage,
|
|
}
|
|
args.manifest.parent.mkdir(parents=True, exist_ok=True)
|
|
args.manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
|
print(json.dumps({"requests": len(output_rows), "prefix_pairs": pair_coverage}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|