#!/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())