kyle
eb4e46d036
Exempt title from the vendor-token check. Truncate snippets to 160 chars and clip all four packs to the same per-query k. run.py fails cached replies with a rerun message instead of a config mismatch.
350 lines
11 KiB
Python
350 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""Shuffle four named result dumps into A–D and write a sealed key file.
|
||
|
||
Blinded packs are a list of hit-lists (one per query, same order as
|
||
search_queries.txt). Each hit is only {rank, url, title, snippet}.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import random
|
||
import re
|
||
import sys
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
LABELS = ("A", "B", "C", "D")
|
||
CONFIGS = (
|
||
"searxng",
|
||
"searxng_brave",
|
||
"searxng_brave_tavily",
|
||
"hermes_tavily",
|
||
)
|
||
HIT_KEYS = ("rank", "url", "title", "snippet")
|
||
LEAK_RE = re.compile(r"searxng|brave|tavily|hermes", re.I)
|
||
EXEMPT_FIELDS = frozenset({"url", "title", "snippet"})
|
||
SNIPPET_MAX = 160
|
||
DEFAULT_K = 8
|
||
|
||
|
||
def load_queries(path: Path) -> list[str]:
|
||
return [ln.strip() for ln in path.read_text().splitlines() if ln.strip()]
|
||
|
||
|
||
def load_targets(path: Path) -> list[str]:
|
||
out: list[str] = []
|
||
for ln in path.read_text().splitlines():
|
||
s = ln.strip()
|
||
if not s or s.startswith("#"):
|
||
continue
|
||
out.append(s)
|
||
return out
|
||
|
||
|
||
def _hit_list(block: object) -> list:
|
||
if isinstance(block, list):
|
||
if block and all(isinstance(x, dict) and ("url" in x or "title" in x) for x in block):
|
||
return block
|
||
return []
|
||
if not isinstance(block, dict):
|
||
return []
|
||
for key in ("results", "web", "hits"):
|
||
v = block.get(key)
|
||
if isinstance(v, list):
|
||
return v
|
||
data = block.get("data")
|
||
if isinstance(data, list):
|
||
return data
|
||
if isinstance(data, dict):
|
||
for key in ("web", "results", "hits"):
|
||
v = data.get(key)
|
||
if isinstance(v, list):
|
||
return v
|
||
return []
|
||
|
||
|
||
def _items(raw: object) -> list:
|
||
if isinstance(raw, list):
|
||
return raw
|
||
if isinstance(raw, dict):
|
||
for key in ("queries", "results"):
|
||
v = raw.get(key)
|
||
if isinstance(v, list) and v and isinstance(v[0], (dict, list)):
|
||
if key == "results" and v and isinstance(v[0], dict) and "url" in v[0]:
|
||
return [raw]
|
||
return v
|
||
return [raw]
|
||
raise ValueError(f"raw dump must be a list or object, got {type(raw).__name__}")
|
||
|
||
|
||
def normalize(raw: object, n: int) -> list[list[dict]]:
|
||
items = _items(raw)
|
||
if len(items) != n:
|
||
raise ValueError(f"expected {n} queries in dump, got {len(items)}")
|
||
out: list[list[dict]] = []
|
||
for item in items:
|
||
hits = []
|
||
if isinstance(item, list):
|
||
raw_hits = item
|
||
else:
|
||
raw_hits = _hit_list(item)
|
||
for i, h in enumerate(raw_hits, 1):
|
||
if not isinstance(h, dict):
|
||
continue
|
||
rank = h.get("rank", h.get("position", i))
|
||
try:
|
||
rank_i = int(rank)
|
||
except (TypeError, ValueError):
|
||
rank_i = i
|
||
snip = str(h.get("snippet") or h.get("description") or h.get("content") or "")
|
||
if len(snip) > SNIPPET_MAX:
|
||
snip = snip[:SNIPPET_MAX]
|
||
hits.append(
|
||
{
|
||
"rank": rank_i,
|
||
"url": str(h.get("url") or ""),
|
||
"title": str(h.get("title") or ""),
|
||
"snippet": snip,
|
||
}
|
||
)
|
||
out.append(hits)
|
||
return out
|
||
|
||
|
||
def equalise(packs: dict[str, list], k: int) -> None:
|
||
n = None
|
||
for cfg in CONFIGS:
|
||
if cfg not in packs:
|
||
raise ValueError(f"missing pack {cfg}")
|
||
got = len(packs[cfg])
|
||
if n is None:
|
||
n = got
|
||
elif got != n:
|
||
raise ValueError(f"{cfg}: {got} queries, want {n}")
|
||
assert n is not None
|
||
for cfg in CONFIGS:
|
||
for i in range(n):
|
||
packs[cfg][i] = packs[cfg][i][:k]
|
||
for i in range(n):
|
||
m = min(len(packs[cfg][i]) for cfg in CONFIGS)
|
||
for cfg in CONFIGS:
|
||
hits = packs[cfg][i][:m]
|
||
for j, h in enumerate(hits, 1):
|
||
h["rank"] = j
|
||
packs[cfg][i] = hits
|
||
|
||
|
||
def leak_hits(obj: object, path: str = "$"):
|
||
if isinstance(obj, dict):
|
||
for k, v in obj.items():
|
||
if k in EXEMPT_FIELDS and isinstance(v, str):
|
||
continue
|
||
yield from leak_hits(v, f"{path}.{k}")
|
||
elif isinstance(obj, list):
|
||
for i, v in enumerate(obj):
|
||
yield from leak_hits(v, f"{path}[{i}]")
|
||
elif isinstance(obj, str) and LEAK_RE.search(obj):
|
||
yield path, obj
|
||
|
||
|
||
def assert_blind(pack: object) -> None:
|
||
if not isinstance(pack, list):
|
||
raise ValueError("blind pack must be a list of hit lists")
|
||
for qi, hits in enumerate(pack):
|
||
if not isinstance(hits, list):
|
||
raise ValueError(f"query {qi}: hits must be a list")
|
||
for hi, h in enumerate(hits):
|
||
if not isinstance(h, dict) or set(h) != set(HIT_KEYS):
|
||
raise ValueError(f"query {qi} hit {hi}: keys must be {HIT_KEYS}, got {h!r}")
|
||
if not isinstance(h["rank"], int):
|
||
raise ValueError(f"query {qi} hit {hi}: rank must be int")
|
||
for k in ("url", "title", "snippet"):
|
||
if not isinstance(h[k], str):
|
||
raise ValueError(f"query {qi} hit {hi}: {k} must be str")
|
||
found = list(leak_hits(pack))
|
||
if found:
|
||
bits = "; ".join(f"{p}={v!r}" for p, v in found[:8])
|
||
raise ValueError(f"vendor token outside url/title/snippet: {bits}")
|
||
|
||
|
||
def write_packs(packs: dict[str, list], out_dir: Path, n: int, seed: int | None, k: int = DEFAULT_K) -> Path:
|
||
equalise(packs, k)
|
||
rng = random.Random(seed)
|
||
order = list(CONFIGS)
|
||
rng.shuffle(order)
|
||
key = {lab: cfg for lab, cfg in zip(LABELS, order)}
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
for lab, cfg in key.items():
|
||
pack = packs[cfg]
|
||
assert_blind(pack)
|
||
(out_dir / f"{lab}.json").write_text(json.dumps(pack, indent=2) + "\n")
|
||
key_path = out_dir / "key.json"
|
||
key_path.write_text(json.dumps({"labels": key, "n": n}, indent=2) + "\n")
|
||
return key_path
|
||
|
||
|
||
def self_test() -> int:
|
||
n = 2
|
||
sift_q = {
|
||
"query": "example docs",
|
||
"results": [
|
||
{
|
||
"url": "https://docs.example.com/search",
|
||
"title": "Search method",
|
||
"snippet": "POST /v2/search",
|
||
"score": 0.91,
|
||
"source": "searxng",
|
||
}
|
||
],
|
||
"warnings": ["tavily: 429"],
|
||
"timings": {"search_ms": 40},
|
||
"credits": 1,
|
||
"config": "searxng",
|
||
"sources": ["searxng"],
|
||
}
|
||
hermes_q = {
|
||
"query": "example docs",
|
||
"success": True,
|
||
"data": {
|
||
"web": [
|
||
{
|
||
"url": "https://api.tavily.com/docs",
|
||
"title": "Vendor docs",
|
||
"description": "search_depth basic",
|
||
"position": 1,
|
||
"score": 0.4,
|
||
}
|
||
]
|
||
},
|
||
"warnings": ["brave: timeout"],
|
||
}
|
||
dumps = {
|
||
"searxng": [sift_q, sift_q],
|
||
"searxng_brave": {
|
||
"config": "searxng_brave",
|
||
"queries": [sift_q, sift_q],
|
||
},
|
||
"searxng_brave_tavily": [sift_q, sift_q],
|
||
"hermes_tavily": [hermes_q, hermes_q],
|
||
}
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
raw = Path(tmp) / "raw"
|
||
out = Path(tmp) / "out"
|
||
raw.mkdir()
|
||
for name, payload in dumps.items():
|
||
(raw / f"{name}.json").write_text(json.dumps(payload) + "\n")
|
||
packs = {name: normalize(payload, n) for name, payload in dumps.items()}
|
||
for pack in packs.values():
|
||
assert_blind(pack)
|
||
key_path = write_packs(packs, out, n, seed=1)
|
||
for lab in LABELS:
|
||
blob = json.loads((out / f"{lab}.json").read_text())
|
||
assert_blind(blob)
|
||
assert len(blob) == n
|
||
assert blob[0][0]["url"].startswith("https://")
|
||
keys = set()
|
||
|
||
def walk(o):
|
||
if isinstance(o, dict):
|
||
keys.update(o)
|
||
for v in o.values():
|
||
walk(v)
|
||
elif isinstance(o, list):
|
||
for v in o:
|
||
walk(v)
|
||
|
||
walk(blob)
|
||
if keys != {"rank", "url", "title", "snippet"}:
|
||
raise AssertionError(f"{lab} keys {keys}")
|
||
# titles of correct hits name vendors; url/title/snippet are exempt
|
||
assert_blind(
|
||
[[{"rank": 1, "url": "https://x.example", "title": "tavily guide", "snippet": "ok"}]]
|
||
)
|
||
ok_snip = [
|
||
[{"rank": 1, "url": "https://api.tavily.com/x", "title": "Docs", "snippet": "uses tavily"}]
|
||
]
|
||
assert_blind(ok_snip)
|
||
long = "L" * (SNIPPET_MAX + 40)
|
||
fat = {
|
||
"results": [
|
||
{"url": "https://a.example/" + str(i), "title": "T", "snippet": long}
|
||
for i in range(5)
|
||
]
|
||
}
|
||
thin = {"results": [{"url": "https://b.example/0", "title": "T", "snippet": "short"}]}
|
||
packs2 = {
|
||
"searxng": normalize([fat, fat], 2),
|
||
"searxng_brave": normalize([thin, thin], 2),
|
||
"searxng_brave_tavily": normalize([fat, fat], 2),
|
||
"hermes_tavily": normalize([thin, thin], 2),
|
||
}
|
||
equalise(packs2, k=8)
|
||
for cfg, pack in packs2.items():
|
||
if len(pack[0]) != 1:
|
||
raise AssertionError(f"{cfg} k={len(pack[0])} after equalise")
|
||
if any(len(h["snippet"]) > SNIPPET_MAX for h in pack[0]):
|
||
raise AssertionError("snippet not truncated")
|
||
print("self-test ok", key_path)
|
||
return 0
|
||
|
||
|
||
def main() -> int:
|
||
p = argparse.ArgumentParser(description=__doc__)
|
||
p.add_argument("--self-test", action="store_true")
|
||
p.add_argument(
|
||
"--raw-dir",
|
||
type=Path,
|
||
help="Directory with {searxng,searxng_brave,searxng_brave_tavily,hermes_tavily}.json",
|
||
)
|
||
p.add_argument(
|
||
"--out-dir",
|
||
type=Path,
|
||
default=Path("tests/testdata/bakeoff-out"),
|
||
)
|
||
p.add_argument(
|
||
"--queries",
|
||
type=Path,
|
||
default=Path("tests/testdata/search_queries.txt"),
|
||
)
|
||
p.add_argument(
|
||
"--targets",
|
||
type=Path,
|
||
default=Path("tests/testdata/search_targets.txt"),
|
||
)
|
||
p.add_argument("--seed", type=int, default=None)
|
||
p.add_argument("--k", type=int, default=DEFAULT_K)
|
||
args = p.parse_args()
|
||
if args.self_test:
|
||
return self_test()
|
||
if args.raw_dir is None:
|
||
p.error("--raw-dir is required unless --self-test")
|
||
|
||
queries = load_queries(args.queries)
|
||
targets = load_targets(args.targets)
|
||
if len(queries) != len(targets):
|
||
print(
|
||
f"query/target length mismatch: {len(queries)} vs {len(targets)}",
|
||
file=sys.stderr,
|
||
)
|
||
return 1
|
||
|
||
packs = {}
|
||
for name in CONFIGS:
|
||
f = args.raw_dir / f"{name}.json"
|
||
packs[name] = normalize(json.loads(f.read_text()), len(queries))
|
||
|
||
try:
|
||
key_path = write_packs(packs, args.out_dir, len(queries), args.seed, args.k)
|
||
except ValueError as e:
|
||
print(e, file=sys.stderr)
|
||
return 1
|
||
print(f"wrote {args.out_dir}/{{A,B,C,D}}.json")
|
||
print(f"key sealed at {key_path} — judges do not open this until scoring is done")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|