p3: bake-off collector that refuses a mislabelled dump
run.py records one named config at a time. Sift dumps fail if sizes.sources_ok does not match --config. Hermes Tavily uses the live-tool request shape, not sift. Raw dumps stay gitignored.
This commit is contained in:
@@ -0,0 +1,286 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Collect one named bake-off config into a raw dump. Does not enable sources.
|
||||
|
||||
Sift configs POST /v1/search against whatever orion currently has enabled.
|
||||
The run fails if sizes.sources_ok does not match the named config, so a
|
||||
mis-set orion cannot be labelled as another configuration.
|
||||
|
||||
Config 4 (hermes_tavily) calls Tavily the way Hermes does (not sift).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import tomllib
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
CONFIGS = (
|
||||
"searxng",
|
||||
"searxng_brave",
|
||||
"searxng_brave_tavily",
|
||||
"hermes_tavily",
|
||||
)
|
||||
EXPECTED_SOURCES = {
|
||||
"searxng": frozenset({"searxng"}),
|
||||
"searxng_brave": frozenset({"searxng", "brave_llm"}),
|
||||
"searxng_brave_tavily": frozenset({"searxng", "brave_llm", "tavily"}),
|
||||
}
|
||||
DEFAULT_K = 8
|
||||
HERMES_TAVILY_URL = "https://api.tavily.com/search"
|
||||
# Hermes plugins/web/tavily/provider.py _SEARCH_PAYLOAD + max_results.
|
||||
HERMES_TAVILY_EXTRA = {"include_raw_content": False, "include_images": False}
|
||||
|
||||
|
||||
def load_queries(path: Path) -> list[str]:
|
||||
return [ln.strip() for ln in path.read_text().splitlines() if ln.strip()]
|
||||
|
||||
|
||||
def load_sift_client(url: str | None, key_file: str | None) -> tuple[str, str]:
|
||||
if url and key_file:
|
||||
return url.rstrip("/"), Path(key_file).expanduser().read_text().strip()
|
||||
cfg_path = os.environ.get("SIFT_CLIENT")
|
||||
if not cfg_path:
|
||||
cfg_path = str(Path.home() / ".config/sift/client.toml")
|
||||
raw = Path(cfg_path).read_bytes()
|
||||
cfg = tomllib.loads(raw.decode())
|
||||
base = (url or cfg.get("url") or "").rstrip("/")
|
||||
kf = key_file or cfg.get("key_file")
|
||||
if not base or not kf:
|
||||
raise SystemExit("sift url and key_file required")
|
||||
key = Path(kf).expanduser().read_text().strip()
|
||||
return base, key
|
||||
|
||||
|
||||
def post_json(url: str, body: dict, headers: dict, timeout: float) -> tuple[int, dict, float]:
|
||||
data = json.dumps(body).encode()
|
||||
req = urllib.request.Request(url, data=data, method="POST")
|
||||
req.add_header("Content-Type", "application/json")
|
||||
for k, v in headers.items():
|
||||
req.add_header(k, v)
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
raw = resp.read()
|
||||
status = resp.status
|
||||
except urllib.error.HTTPError as e:
|
||||
raw = e.read()
|
||||
status = e.code
|
||||
elapsed = (time.perf_counter() - t0) * 1000
|
||||
try:
|
||||
env = json.loads(raw.decode())
|
||||
except json.JSONDecodeError:
|
||||
env = {"error": "non-json"}
|
||||
if not isinstance(env, dict):
|
||||
env = {"error": "non-object"}
|
||||
return status, env, elapsed
|
||||
|
||||
|
||||
def sift_one(base: str, key: str, query: str, k: int, timeout: float) -> dict:
|
||||
status, env, ms = post_json(
|
||||
base + "/v1/search",
|
||||
{"query": query, "k": k},
|
||||
{"Authorization": "Bearer " + key},
|
||||
timeout,
|
||||
)
|
||||
results = env.get("results") if isinstance(env.get("results"), list) else []
|
||||
sizes = env.get("sizes") if isinstance(env.get("sizes"), dict) else {}
|
||||
timings = env.get("timings") if isinstance(env.get("timings"), dict) else {}
|
||||
warnings = env.get("warnings") if isinstance(env.get("warnings"), list) else []
|
||||
return {
|
||||
"query": query,
|
||||
"http_status": status,
|
||||
"results": results,
|
||||
"warnings": warnings,
|
||||
"sizes": sizes,
|
||||
"timings": timings,
|
||||
"cached": bool(env.get("cached")),
|
||||
"latency_ms": round(ms, 1),
|
||||
"error": env.get("error") if status != 200 else None,
|
||||
}
|
||||
|
||||
|
||||
def check_sources(config: str, rows: list[dict]) -> str | None:
|
||||
want = EXPECTED_SOURCES[config]
|
||||
for i, row in enumerate(rows):
|
||||
sizes = row.get("sizes") or {}
|
||||
ok = sizes.get("sources_ok") or []
|
||||
failed = sizes.get("sources_failed") or []
|
||||
got = frozenset(ok)
|
||||
if got != want:
|
||||
return (
|
||||
f"query {i}: sources_ok={sorted(got)} want={sorted(want)} "
|
||||
f"failed={failed} (orion config does not match --config {config})"
|
||||
)
|
||||
if failed:
|
||||
return f"query {i}: sources_failed={failed}"
|
||||
return None
|
||||
|
||||
|
||||
def hermes_tavily_one(url: str, key: str, query: str, k: int, timeout: float) -> dict:
|
||||
payload = {"query": query, "max_results": k, **HERMES_TAVILY_EXTRA}
|
||||
headers = {"Authorization": "Bearer " + key, "X-Client-Name": "hermes-agent"}
|
||||
status, env, ms = post_json(url, payload, headers, timeout)
|
||||
raw_results = env.get("results")
|
||||
results = raw_results if isinstance(raw_results, list) else []
|
||||
web = []
|
||||
for i, r in enumerate(results, 1):
|
||||
if not isinstance(r, dict):
|
||||
continue
|
||||
web.append(
|
||||
{
|
||||
"url": r.get("url") or "",
|
||||
"title": r.get("title") or "",
|
||||
"description": r.get("content") or r.get("snippet") or "",
|
||||
"position": i,
|
||||
}
|
||||
)
|
||||
return {
|
||||
"query": query,
|
||||
"success": status == 200,
|
||||
"http_status": status,
|
||||
"data": {"web": web},
|
||||
"latency_ms": round(ms, 1),
|
||||
"credits": env.get("usage") if isinstance(env.get("usage"), dict) else None,
|
||||
"error": env.get("error") if status != 200 else None,
|
||||
}
|
||||
|
||||
|
||||
def write_dump(path: Path, config: str, rows: list[dict]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
payload = {"config": config, "queries": rows}
|
||||
path.write_text(json.dumps(payload, indent=2) + "\n")
|
||||
|
||||
|
||||
def self_test() -> int:
|
||||
import tempfile
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from threading import Thread
|
||||
|
||||
class H(BaseHTTPRequestHandler):
|
||||
def log_message(self, format, *args):
|
||||
return
|
||||
|
||||
def do_POST(self):
|
||||
n = int(self.headers.get("Content-Length") or 0)
|
||||
_ = self.rfile.read(n)
|
||||
if self.path == "/v1/search":
|
||||
body = {
|
||||
"results": [
|
||||
{
|
||||
"url": "https://docs.example.com",
|
||||
"title": "Docs",
|
||||
"snippet": "ok",
|
||||
"score": 0.2,
|
||||
"source": "searxng",
|
||||
}
|
||||
],
|
||||
"warnings": [],
|
||||
"sizes": {
|
||||
"results": 1,
|
||||
"sources_ok": ["searxng"],
|
||||
"sources_failed": [],
|
||||
},
|
||||
"timings": {"search_ms": 3},
|
||||
"cached": False,
|
||||
}
|
||||
else:
|
||||
body = {
|
||||
"results": [
|
||||
{
|
||||
"url": "https://example.com/a",
|
||||
"title": "Hit",
|
||||
"content": "snippet",
|
||||
}
|
||||
],
|
||||
"usage": {"credits": 1},
|
||||
}
|
||||
raw = json.dumps(body).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(raw)))
|
||||
self.end_headers()
|
||||
self.wfile.write(raw)
|
||||
|
||||
srv = ThreadingHTTPServer(("127.0.0.1", 0), H)
|
||||
Thread(target=srv.serve_forever, daemon=True).start()
|
||||
host = f"http://127.0.0.1:{srv.server_address[1]}"
|
||||
try:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
qpath = Path(tmp) / "q.txt"
|
||||
qpath.write_text("one\ntwo\n")
|
||||
raw_dir = Path(tmp) / "raw"
|
||||
keyf = Path(tmp) / "key"
|
||||
keyf.write_text("k\n")
|
||||
# matching sources_ok
|
||||
rows = [
|
||||
sift_one(host, "k", "one", 8, 5),
|
||||
sift_one(host, "k", "two", 8, 5),
|
||||
]
|
||||
err = check_sources("searxng", rows)
|
||||
if err:
|
||||
raise AssertionError(err)
|
||||
bad = check_sources("searxng_brave", rows)
|
||||
if not bad:
|
||||
raise AssertionError("mismatch not caught")
|
||||
ht = hermes_tavily_one(host + "/search", "k", "one", 8, 5)
|
||||
if ht["data"]["web"][0]["description"] != "snippet":
|
||||
raise AssertionError(ht)
|
||||
write_dump(raw_dir / "searxng.json", "searxng", rows)
|
||||
print("self-test ok")
|
||||
finally:
|
||||
srv.shutdown()
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--config", choices=CONFIGS)
|
||||
p.add_argument("--raw-dir", type=Path, default=Path("tests/testdata/bakeoff-raw"))
|
||||
p.add_argument("--queries", type=Path, default=Path("tests/testdata/search_queries.txt"))
|
||||
p.add_argument("--k", type=int, default=DEFAULT_K)
|
||||
p.add_argument("--timeout", type=float, default=60)
|
||||
p.add_argument("--sift-url")
|
||||
p.add_argument("--sift-key-file")
|
||||
p.add_argument("--tavily-url", default=HERMES_TAVILY_URL)
|
||||
p.add_argument("--self-test", action="store_true")
|
||||
args = p.parse_args()
|
||||
if args.self_test:
|
||||
return self_test()
|
||||
if not args.config:
|
||||
p.error("--config is required unless --self-test")
|
||||
queries = load_queries(args.queries)
|
||||
if not queries:
|
||||
print("no queries", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
rows: list[dict] = []
|
||||
if args.config == "hermes_tavily":
|
||||
key = os.environ.get("TAVILY_API_KEY", "").strip()
|
||||
if not key:
|
||||
print("TAVILY_API_KEY required for hermes_tavily", file=sys.stderr)
|
||||
return 1
|
||||
for q in queries:
|
||||
rows.append(hermes_tavily_one(args.tavily_url, key, q, args.k, args.timeout))
|
||||
else:
|
||||
base, key = load_sift_client(args.sift_url, args.sift_key_file)
|
||||
for q in queries:
|
||||
rows.append(sift_one(base, key, q, args.k, args.timeout))
|
||||
err = check_sources(args.config, rows)
|
||||
if err:
|
||||
print(err, file=sys.stderr)
|
||||
return 1
|
||||
|
||||
out = args.raw_dir / f"{args.config}.json"
|
||||
write_dump(out, args.config, rows)
|
||||
print(f"wrote {out} n={len(rows)}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user