|
| 1 | +"""Measure model-visible tool-catalog cost in Codex ATIF trajectories.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import argparse |
| 6 | +import hashlib |
| 7 | +import json |
| 8 | +import statistics |
| 9 | +from pathlib import Path |
| 10 | +from typing import Any |
| 11 | + |
| 12 | +from benchmarks.tooling.errors import HarborSuiteError |
| 13 | + |
| 14 | +_JACOBIAN_FIND = "tools.mcp__jacobian__math_find(" |
| 15 | +_JACOBIAN_RUN = "tools.mcp__jacobian__math_run(" |
| 16 | + |
| 17 | + |
| 18 | +def _read_trajectory(path: Path) -> dict[str, Any]: |
| 19 | + try: |
| 20 | + value = json.loads(path.read_text(encoding="utf-8")) |
| 21 | + except (OSError, UnicodeError, json.JSONDecodeError) as exc: |
| 22 | + raise HarborSuiteError(f"unable to read ATIF trajectory {path}: {exc}") from exc |
| 23 | + if not isinstance(value, dict) or not isinstance(value.get("steps"), list): |
| 24 | + raise HarborSuiteError( |
| 25 | + f"invalid ATIF trajectory {path}: steps must be an array" |
| 26 | + ) |
| 27 | + return value |
| 28 | + |
| 29 | + |
| 30 | +def _visible_bytes(value: object) -> int: |
| 31 | + if isinstance(value, str): |
| 32 | + encoded = value.encode("utf-8") |
| 33 | + else: |
| 34 | + encoded = json.dumps( |
| 35 | + value, ensure_ascii=False, separators=(",", ":"), sort_keys=True |
| 36 | + ).encode("utf-8") |
| 37 | + return len(encoded) |
| 38 | + |
| 39 | + |
| 40 | +def _integer(value: object) -> int | None: |
| 41 | + return value if isinstance(value, int) and not isinstance(value, bool) else None |
| 42 | + |
| 43 | + |
| 44 | +def _number(value: object) -> int | float | None: |
| 45 | + return ( |
| 46 | + value |
| 47 | + if isinstance(value, (int, float)) and not isinstance(value, bool) |
| 48 | + else None |
| 49 | + ) |
| 50 | + |
| 51 | + |
| 52 | +def _visible_results(step: dict[str, Any]) -> dict[str, int]: |
| 53 | + observation = step.get("observation") |
| 54 | + results = observation.get("results", []) if isinstance(observation, dict) else [] |
| 55 | + visible_by_call: dict[str, int] = {} |
| 56 | + if not isinstance(results, list): |
| 57 | + return visible_by_call |
| 58 | + for result in results: |
| 59 | + if not isinstance(result, dict): |
| 60 | + continue |
| 61 | + source_call_id = result.get("source_call_id") |
| 62 | + if isinstance(source_call_id, str) and "content" in result: |
| 63 | + visible_by_call[source_call_id] = _visible_bytes(result["content"]) |
| 64 | + return visible_by_call |
| 65 | + |
| 66 | + |
| 67 | +def _analyze_step(step: object) -> tuple[int, int, int, int, int, int]: |
| 68 | + if not isinstance(step, dict): |
| 69 | + return (0, 0, 0, 0, 0, 0) |
| 70 | + visible_by_call = _visible_results(step) |
| 71 | + tool_calls = step.get("tool_calls", []) |
| 72 | + if not isinstance(tool_calls, list): |
| 73 | + return (0, 0, 0, 0, 0, 0) |
| 74 | + scan_count = scan_bytes = unbound_scan_count = tool_output_bytes = 0 |
| 75 | + direct_find_references = direct_run_references = 0 |
| 76 | + for call in tool_calls: |
| 77 | + if not isinstance(call, dict): |
| 78 | + continue |
| 79 | + call_id = call.get("tool_call_id") |
| 80 | + visible = visible_by_call.get(call_id, 0) if isinstance(call_id, str) else 0 |
| 81 | + tool_output_bytes += visible |
| 82 | + arguments = call.get("arguments") |
| 83 | + source = arguments.get("input") if isinstance(arguments, dict) else None |
| 84 | + if call.get("function_name") != "exec" or not isinstance(source, str): |
| 85 | + continue |
| 86 | + direct_find_references += source.count(_JACOBIAN_FIND) |
| 87 | + direct_run_references += source.count(_JACOBIAN_RUN) |
| 88 | + if "ALL_TOOLS" in source: |
| 89 | + scan_count += 1 |
| 90 | + scan_bytes += visible |
| 91 | + if not isinstance(call_id, str) or call_id not in visible_by_call: |
| 92 | + unbound_scan_count += 1 |
| 93 | + return ( |
| 94 | + scan_count, |
| 95 | + scan_bytes, |
| 96 | + unbound_scan_count, |
| 97 | + tool_output_bytes, |
| 98 | + direct_find_references, |
| 99 | + direct_run_references, |
| 100 | + ) |
| 101 | + |
| 102 | + |
| 103 | +def analyze_trajectory(path: Path) -> dict[str, Any]: |
| 104 | + """Extract directory projection and token-cost facts from one ATIF trace.""" |
| 105 | + |
| 106 | + trajectory = _read_trajectory(path) |
| 107 | + scan_count = 0 |
| 108 | + scan_bytes = 0 |
| 109 | + unbound_scan_count = 0 |
| 110 | + tool_output_bytes = 0 |
| 111 | + direct_find_references = 0 |
| 112 | + direct_run_references = 0 |
| 113 | + |
| 114 | + for step in trajectory["steps"]: |
| 115 | + step_counts = _analyze_step(step) |
| 116 | + scan_count += step_counts[0] |
| 117 | + scan_bytes += step_counts[1] |
| 118 | + unbound_scan_count += step_counts[2] |
| 119 | + tool_output_bytes += step_counts[3] |
| 120 | + direct_find_references += step_counts[4] |
| 121 | + direct_run_references += step_counts[5] |
| 122 | + |
| 123 | + metrics = trajectory.get("final_metrics") |
| 124 | + metrics = metrics if isinstance(metrics, dict) else {} |
| 125 | + prompt_tokens = _integer(metrics.get("total_prompt_tokens")) |
| 126 | + cached_tokens = _integer(metrics.get("total_cached_tokens")) |
| 127 | + uncached_tokens = ( |
| 128 | + max(0, prompt_tokens - cached_tokens) |
| 129 | + if prompt_tokens is not None and cached_tokens is not None |
| 130 | + else None |
| 131 | + ) |
| 132 | + return { |
| 133 | + "trajectory": str(path), |
| 134 | + "trajectory_sha256": "sha256:" + hashlib.sha256(path.read_bytes()).hexdigest(), |
| 135 | + "agent": trajectory.get("agent"), |
| 136 | + "all_tools_scan_count": scan_count, |
| 137 | + "all_tools_model_visible_bytes": scan_bytes, |
| 138 | + "all_tools_unbound_observation_count": unbound_scan_count, |
| 139 | + "tool_model_visible_bytes": tool_output_bytes, |
| 140 | + "direct_jacobian_find_references": direct_find_references, |
| 141 | + "direct_jacobian_run_references": direct_run_references, |
| 142 | + "prompt_tokens": prompt_tokens, |
| 143 | + "cached_prompt_tokens": cached_tokens, |
| 144 | + "uncached_prompt_tokens": uncached_tokens, |
| 145 | + "completion_tokens": _integer(metrics.get("total_completion_tokens")), |
| 146 | + "cost_usd": _number(metrics.get("total_cost_usd")), |
| 147 | + } |
| 148 | + |
| 149 | + |
| 150 | +def _median(trials: list[dict[str, Any]], field: str) -> int | float | None: |
| 151 | + values = [ |
| 152 | + trial[field] for trial in trials if isinstance(trial.get(field), (int, float)) |
| 153 | + ] |
| 154 | + return statistics.median(values) if values else None |
| 155 | + |
| 156 | + |
| 157 | +def build_report(paths: list[Path], *, label: str) -> dict[str, Any]: |
| 158 | + """Build one digest-bound observation report for a set of trajectories.""" |
| 159 | + |
| 160 | + if not paths: |
| 161 | + raise HarborSuiteError("at least one ATIF trajectory is required") |
| 162 | + trials = [analyze_trajectory(path) for path in paths] |
| 163 | + return { |
| 164 | + "schema_version": "1", |
| 165 | + "label": label, |
| 166 | + "trial_count": len(trials), |
| 167 | + "summary": { |
| 168 | + "all_tools_scan_trials": sum( |
| 169 | + int(trial["all_tools_scan_count"] > 0) for trial in trials |
| 170 | + ), |
| 171 | + "all_tools_scan_count": sum( |
| 172 | + int(trial["all_tools_scan_count"]) for trial in trials |
| 173 | + ), |
| 174 | + "all_tools_model_visible_bytes": sum( |
| 175 | + int(trial["all_tools_model_visible_bytes"]) for trial in trials |
| 176 | + ), |
| 177 | + "all_tools_unbound_observation_count": sum( |
| 178 | + int(trial["all_tools_unbound_observation_count"]) for trial in trials |
| 179 | + ), |
| 180 | + "median_prompt_tokens": _median(trials, "prompt_tokens"), |
| 181 | + "median_cached_prompt_tokens": _median(trials, "cached_prompt_tokens"), |
| 182 | + "median_uncached_prompt_tokens": _median(trials, "uncached_prompt_tokens"), |
| 183 | + "median_completion_tokens": _median(trials, "completion_tokens"), |
| 184 | + "median_cost_usd": _median(trials, "cost_usd"), |
| 185 | + }, |
| 186 | + "trials": trials, |
| 187 | + } |
| 188 | + |
| 189 | + |
| 190 | +def _parser() -> argparse.ArgumentParser: |
| 191 | + parser = argparse.ArgumentParser(description=__doc__) |
| 192 | + parser.add_argument("trajectories", nargs="+", type=Path) |
| 193 | + parser.add_argument("--label", default="observation") |
| 194 | + parser.add_argument("--output", type=Path) |
| 195 | + return parser |
| 196 | + |
| 197 | + |
| 198 | +def main() -> int: |
| 199 | + args = _parser().parse_args() |
| 200 | + try: |
| 201 | + report = build_report(args.trajectories, label=args.label) |
| 202 | + except HarborSuiteError as exc: |
| 203 | + raise SystemExit(str(exc)) from exc |
| 204 | + rendered = json.dumps(report, indent=2, sort_keys=True) + "\n" |
| 205 | + if args.output is None: |
| 206 | + print(rendered, end="") |
| 207 | + else: |
| 208 | + args.output.parent.mkdir(parents=True, exist_ok=True) |
| 209 | + args.output.write_text(rendered, encoding="utf-8") |
| 210 | + return 0 |
| 211 | + |
| 212 | + |
| 213 | +if __name__ == "__main__": |
| 214 | + raise SystemExit(main()) |
| 215 | + |
| 216 | + |
| 217 | +__all__ = ["analyze_trajectory", "build_report"] |
0 commit comments