File size: 7,307 Bytes
5655a42 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | #!/usr/bin/env python3
"""Simulated agent lookups: how many tool calls / tokens to answer "show me the definition of X"?
Deterministic agent policy (the one a careful model actually follows, given read_file's 2,000-line window):
1. grep -n "def X\\b|class X\\b|^X\\s*=" across the tree -> 1 call, returns (file, line) hits
2. read_file(file, offset=hit_line-20, limit=W) for W in (200,) -> 1 call; if the symbol's end is past the
window, read the next 2,000-line window until it is -> +1 call each
3. done when the whole definition is in context
Costs charged per task:
tool_calls = 1 grep + N reads
tokens = tokens of grep output (all hits, one line each) + tokens of every read window returned
Both trees get the same symbol set: the intersection of symbols that exist (by name) on both, so a symbol moving
to a smaller file counts as a WIN and a symbol that was deleted on one side is excluded, not scored.
Usage: python evals/codebase_navigability/lookup_sim.py <base_tree> <head_tree> [--sample N] [--seed S]
"""
from __future__ import annotations
import argparse
import ast
import json
import os
import random
import statistics
import sys
from pathlib import Path
SKIP_TOP = {".git", "node_modules", "apps", "website", "build", ".venv", "venv", "MagicMock", "__pycache__", ".worktrees", "dist", "evals", "skills", "optional-skills", "docs", "tests"}
WINDOW = 2000 # read_file's default max lines per call
FIRST_READ = 200 # a careful agent's first targeted read around the grep hit (override: --first-read)
MARGIN = 5 # lines of context above the grep hit
def tokenizer():
try:
import tiktoken
enc = tiktoken.get_encoding("o200k_base")
return lambda s: len(enc.encode(s, disallowed_special=()))
except Exception:
return lambda s: len(s.encode()) // 4
def index(tree: Path):
"""name -> list of (relpath, start_line, end_line, lines_of_file)"""
idx: dict[str, list] = {}
files: dict[str, list[str]] = {}
for dp, dns, fns in os.walk(tree):
rel = os.path.relpath(dp, tree)
if rel != "." and rel.split(os.sep)[0] in SKIP_TOP:
dns[:] = []
continue
for f in fns:
if not f.endswith(".py"):
continue
p = Path(dp) / f
src = p.read_text(encoding="utf-8", errors="replace")
lines = src.split("\n")
r = os.path.relpath(p, tree)
files[r] = lines
try:
t = ast.parse(src)
except SyntaxError:
continue
for n in ast.walk(t):
if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
s = min([d.lineno for d in n.decorator_list] + [n.lineno])
idx.setdefault(n.name, []).append((r, s, n.end_lineno, len(lines)))
return idx, files
def grep_hits(name, idx):
"""Simulate `grep -rn` for the symbol: every def/class of that name, one line each."""
return idx.get(name, [])
def simulate(name, idx, files, tok):
hits = grep_hits(name, idx)
if not hits:
return None
grep_out = "\n".join(f"{r}:{s}: def {name}(...)" for r, s, _, _ in hits)
tokens = tok(grep_out)
calls = 1
# The agent opens the FIRST hit (deterministic; ambiguity costs are the same policy on both trees)
r, s, e, n = hits[0]
lines = files[r]
start = max(1, s - MARGIN)
# The agent doesn't know where the definition ends until it reads it: it asks for a window
# sized to the typical definition (FIRST_READ) and pages forward only if the def keeps going.
end = min(n, start + FIRST_READ - 1)
tokens += tok("\n".join(lines[start - 1:end]))
calls += 1
while end < e: # definition continues past the window: page forward
start = end + 1
end = min(n, start + WINDOW - 1)
tokens += tok("\n".join(lines[start - 1:end]))
calls += 1
exact = tok("\n".join(lines[s - 1:e]))
return {"name": name, "file": r, "file_lines": n, "def_lines": e - s + 1, "hits": len(hits), "calls": calls, "tokens": tokens, "exact_tokens": exact}
def main():
global FIRST_READ
ap = argparse.ArgumentParser()
ap.add_argument("base"); ap.add_argument("head")
ap.add_argument("--sample", type=int, default=3000); ap.add_argument("--seed", type=int, default=7)
ap.add_argument("--first-read", type=int, default=FIRST_READ, help="lines in the first targeted read")
ap.add_argument("--out", default=".")
a = ap.parse_args()
FIRST_READ = a.first_read
tok = tokenizer()
bi, bf = index(Path(a.base)); hi, hf = index(Path(a.head))
common = sorted(set(bi) & set(hi))
# weight the sample toward symbols that are actually looked up: public names, exclude dunders/tests
common = [n for n in common if not n.startswith("__") and not n.startswith("test_")]
random.Random(a.seed).shuffle(common)
sample = common[: a.sample]
B = [simulate(n, bi, bf, tok) for n in sample]
H = [simulate(n, hi, hf, tok) for n in sample]
pairs = [(b, h) for b, h in zip(B, H) if b and h]
def agg(rows):
c = [r["calls"] for r in rows]; t = [r["tokens"] for r in rows]
return {"tasks": len(rows), "calls_mean": round(statistics.mean(c), 3), "calls_total": sum(c), "multi_window_tasks": sum(1 for x in c if x > 2),
"tokens_mean": round(statistics.mean(t)), "tokens_p50": int(statistics.median(t)), "tokens_p90": sorted(t)[int(len(t) * .9)], "tokens_total": sum(t),
"file_lines_p50": int(statistics.median(r["file_lines"] for r in rows)), "def_lines_p50": int(statistics.median(r["def_lines"] for r in rows)),
"ambiguous_hits_mean": round(statistics.mean(r["hits"] for r in rows), 2),
"exact_def_tokens_mean": round(statistics.mean(r["exact_tokens"] for r in rows)), "exact_def_tokens_total": sum(r["exact_tokens"] for r in rows)}
res = {"common_symbols": len(common), "sampled": len(pairs), "seed": a.seed, "window": WINDOW, "first_read": FIRST_READ,
"base": agg([b for b, _ in pairs]), "head": agg([h for _, h in pairs]),
"head_cheaper": sum(1 for b, h in pairs if h["tokens"] < b["tokens"]), "head_costlier": sum(1 for b, h in pairs if h["tokens"] > b["tokens"]), "equal": sum(1 for b, h in pairs if h["tokens"] == b["tokens"]),
"worst_regressions": sorted(({"name": h["name"], "base_tok": b["tokens"], "head_tok": h["tokens"], "head_file": h["file"]} for b, h in pairs), key=lambda x: x["base_tok"] - x["head_tok"])[:8],
"best_wins": sorted(({"name": h["name"], "base_tok": b["tokens"], "head_tok": h["tokens"], "base_file": b["file"], "head_file": h["file"]} for b, h in pairs), key=lambda x: x["head_tok"] - x["base_tok"])[:8]}
Path(a.out).mkdir(parents=True, exist_ok=True)
(Path(a.out) / f"lookup_sim_w{FIRST_READ}.json").write_text(json.dumps(res, indent=1))
print(json.dumps({k: v for k, v in res.items() if k not in ("worst_regressions", "best_wins")}, indent=1))
print("best wins:", json.dumps(res["best_wins"][:4], indent=1)); print("worst regressions:", json.dumps(res["worst_regressions"][:4], indent=1))
if __name__ == "__main__":
main()
|