diff --git a/scripts/run_neighbor_sweep.py b/scripts/run_neighbor_sweep.py new file mode 100644 index 0000000..22ca749 --- /dev/null +++ b/scripts/run_neighbor_sweep.py @@ -0,0 +1,374 @@ +#!/usr/bin/env python3 +"""Run a Benchmark Sweep: fixed_size × text-embedding-3-large × Neighbor ±0..±3. + +Calls the running FastAPI app over HTTP (same path as the Dashboard). +Doc-major order; gap-fills Process (fixed_size only) when SQLite provenance +says the Model Corpus is not ready; retries each unit twice on failure. + +Usage: + .venv/bin/python scripts/run_neighbor_sweep.py + .venv/bin/python scripts/run_neighbor_sweep.py --base-url http://127.0.0.1:8000 +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from dataclasses import dataclass, field +from typing import Any + +import httpx + +CORPUS_MODEL_ID = "text-embedding-3-large" +STRATEGY = "fixed_size" +TOP_K = 5 +NEIGHBOR_LEVELS = (0, 1, 2, 3) +MAX_ATTEMPTS = 3 # 1 try + 2 retries +DEFAULT_BASE_URL = "http://127.0.0.1:8000" +# Benchmarks with many questions can take a long time (LLM + eval per Q). +REQUEST_TIMEOUT_S = 60 * 60 * 3 # 3 hours + +# App-truth document ↔ questions map (filenames as stored in the app). +SWEEP_PAIRS: list[tuple[str, str]] = [ + ("bazresi.docx", "files/bazresi.json"), + ("customer1.docx", "files/customer1.json"), + ("fire.docx", "files/fire.json"), + ("general-havades-individuals.doc", "files/general-havades-individuals.json"), + ("havades.docx", "files/havades.json"), + ("life-time-individual.docx", "files/life-time-individual.json"), + ("moavenin.docx", "files/moavenin.json"), + ("Refah.docx", "files/refah.json"), + ("website.docx", "files/website.json"), + ("lifetime-compensation.docx", "files/lifetime-compensation.json"), +] + + +@dataclass +class UnitResult: + document: str + neighbor: int + kind: str # "process" | "experiment" + ok: bool + detail: str + experiment_id: str | None = None + attempts: int = 1 + + +@dataclass +class SweepState: + results: list[UnitResult] = field(default_factory=list) + + @property + def failed(self) -> list[UnitResult]: + return [r for r in self.results if not r.ok] + + +def log(msg: str) -> None: + print(msg, flush=True) + + +def is_ready(doc: dict[str, Any]) -> bool: + corpus = doc.get("last_corpus_embedding_model_id") + counts = doc.get("chunk_counts") or {} + fixed = int(counts.get("fixed_size") or 0) + return corpus == CORPUS_MODEL_ID and fixed > 0 + + +def fetch_documents(client: httpx.Client) -> dict[str, dict[str, Any]]: + """Return filename → document dict (paginated).""" + by_name: dict[str, dict[str, Any]] = {} + offset = 0 + limit = 100 + while True: + r = client.get("/documents", params={"offset": offset, "limit": limit}) + r.raise_for_status() + data = r.json() + items = data.get("items") or [] + for doc in items: + by_name[doc["filename"]] = doc + total = int(data.get("total") or 0) + offset += len(items) + if offset >= total or not items: + break + return by_name + + +def error_detail(exc: BaseException) -> str: + if isinstance(exc, httpx.HTTPStatusError): + body = "" + try: + body = exc.response.text[:500] + except Exception: + pass + return f"HTTP {exc.response.status_code}: {body or exc}" + return str(exc) + + +def with_retries( + label: str, + fn, + *, + state: SweepState, + document: str, + neighbor: int, + kind: str, +) -> Any | None: + last_err = "" + for attempt in range(1, MAX_ATTEMPTS + 1): + try: + log(f" [{kind}] {label} (attempt {attempt}/{MAX_ATTEMPTS})") + value = fn() + return value + except Exception as exc: + last_err = error_detail(exc) + log(f" ! failed attempt {attempt}/{MAX_ATTEMPTS}: {last_err}") + if attempt < MAX_ATTEMPTS: + time.sleep(min(5 * attempt, 15)) + state.results.append( + UnitResult( + document=document, + neighbor=neighbor, + kind=kind, + ok=False, + detail=last_err, + attempts=MAX_ATTEMPTS, + ) + ) + return None + + +def ensure_corpus( + client: httpx.Client, + doc: dict[str, Any], + *, + state: SweepState, +) -> dict[str, Any] | None: + """Process fixed_size under CORPUS_MODEL_ID if not ready. Returns updated doc or None.""" + filename = doc["filename"] + if is_ready(doc): + log(f" corpus ready ({CORPUS_MODEL_ID}, fixed_size={doc['chunk_counts'].get('fixed_size')})") + return doc + + log( + f" corpus not ready " + f"(last_corpus={doc.get('last_corpus_embedding_model_id')!r}, " + f"fixed_size={ (doc.get('chunk_counts') or {}).get('fixed_size') }) — processing" + ) + + def _process() -> dict[str, Any]: + r = client.post( + f"/documents/{doc['id']}/process", + json={ + "strategies": [STRATEGY], + "corpus_model_id": CORPUS_MODEL_ID, + }, + ) + r.raise_for_status() + body = r.json() + failed = body.get("strategies_failed") or [] + if failed: + raise RuntimeError(f"process strategies_failed: {json.dumps(failed)}") + # Refresh document record + docs = fetch_documents(client) + updated = docs.get(filename) + if updated is None: + raise RuntimeError(f"document disappeared after process: {filename}") + if not is_ready(updated): + raise RuntimeError( + "process finished but provenance still not ready: " + f"last_corpus={updated.get('last_corpus_embedding_model_id')!r}, " + f"chunk_counts={updated.get('chunk_counts')}" + ) + return updated + + result = with_retries( + f"process {filename}", + _process, + state=state, + document=filename, + neighbor=-1, + kind="process", + ) + if result is None: + return None + state.results.append( + UnitResult( + document=filename, + neighbor=-1, + kind="process", + ok=True, + detail="gap-filled fixed_size", + attempts=1, + ) + ) + return result + + +def run_experiment( + client: httpx.Client, + *, + doc: dict[str, Any], + questions_file: str, + neighbor: int, + state: SweepState, +) -> None: + filename = doc["filename"] + + def _bench() -> dict[str, Any]: + r = client.post( + "/benchmarks", + json={ + "document_id": doc["id"], + "strategies": [STRATEGY], + "questions_file": questions_file, + "top_k": TOP_K, + "neighbor_prev": neighbor, + "neighbor_next": neighbor, + "corpus_model_id": CORPUS_MODEL_ID, + }, + ) + r.raise_for_status() + return r.json() + + body = with_retries( + f"±{neighbor}/{neighbor} on {filename}", + _bench, + state=state, + document=filename, + neighbor=neighbor, + kind="experiment", + ) + if body is None: + return + + exp_id = body.get("experiment_id", "") + state.results.append( + UnitResult( + document=filename, + neighbor=neighbor, + kind="experiment", + ok=True, + detail=( + f"best={body.get('best_strategy')} " + f"latency={body.get('total_latency_seconds')}s " + f"cost=${body.get('estimated_cost_usd')}" + ), + experiment_id=exp_id, + ) + ) + log(f" ok experiment_id={exp_id}") + + +def print_summary(state: SweepState) -> None: + log("") + log("=" * 88) + log("BENCHMARK SWEEP SUMMARY") + log("=" * 88) + log( + f"{'document':<40} {'±N':>4} {'kind':<11} {'status':<6} detail" + ) + log("-" * 88) + for r in state.results: + n = "—" if r.neighbor < 0 else f"±{r.neighbor}" + status = "ok" if r.ok else "FAIL" + detail = r.experiment_id or r.detail + if not r.ok and r.experiment_id is None: + detail = r.detail + elif r.ok and r.experiment_id: + detail = f"id={r.experiment_id} {r.detail}" + log(f"{r.document:<40} {n:>4} {r.kind:<11} {status:<6} {detail}") + + failed = state.failed + log("-" * 88) + exp_ok = sum(1 for r in state.results if r.kind == "experiment" and r.ok) + exp_fail = sum(1 for r in state.results if r.kind == "experiment" and not r.ok) + log(f"Experiments: {exp_ok} ok, {exp_fail} failed") + + if failed: + log("") + log("Manual redo hints (failed units):") + for r in failed: + if r.kind == "process": + log( + f" • Process {r.document}: " + f'POST /documents/{{id}}/process ' + f'{{"strategies":["{STRATEGY}"],"corpus_model_id":"{CORPUS_MODEL_ID}"}}' + ) + else: + log( + f" • Experiment {r.document} ±{r.neighbor}: " + f'POST /benchmarks with strategies=["{STRATEGY}"], ' + f"neighbor_prev={r.neighbor}, neighbor_next={r.neighbor}, " + f"corpus_model_id={CORPUS_MODEL_ID}, questions_file from SWEEP_PAIRS" + ) + log(f" error: {r.detail}") + + +def main() -> int: + parser = argparse.ArgumentParser(description="Run fixed_size Neighbor Expansion Benchmark Sweep") + parser.add_argument( + "--base-url", + default=DEFAULT_BASE_URL, + help=f"FastAPI base URL (default: {DEFAULT_BASE_URL})", + ) + args = parser.parse_args() + base_url = args.base_url.rstrip("/") + + state = SweepState() + log(f"Benchmark Sweep → {base_url}") + log(f"Strategy={STRATEGY} corpus={CORPUS_MODEL_ID} top_k={TOP_K} neighbors={list(NEIGHBOR_LEVELS)}") + log(f"Documents: {len(SWEEP_PAIRS)}") + + timeout = httpx.Timeout(REQUEST_TIMEOUT_S, connect=30.0) + with httpx.Client(base_url=base_url, timeout=timeout) as client: + try: + health = client.get("/admin/health") + health.raise_for_status() + except Exception as exc: + log(f"Cannot reach app at {base_url}: {error_detail(exc)}") + return 2 + + docs = fetch_documents(client) + log(f"Loaded {len(docs)} documents from app") + + for filename, questions_file in SWEEP_PAIRS: + log("") + log(f"=== {filename} ({questions_file}) ===") + doc = docs.get(filename) + if doc is None: + log(f" ! document not found in app — skipping") + state.results.append( + UnitResult( + document=filename, + neighbor=-1, + kind="process", + ok=False, + detail="document not found in GET /documents", + attempts=1, + ) + ) + continue + + updated = ensure_corpus(client, doc, state=state) + if updated is None: + log(f" ! skipping experiments for {filename} (process failed)") + continue + docs[filename] = updated + + for n in NEIGHBOR_LEVELS: + run_experiment( + client, + doc=updated, + questions_file=questions_file, + neighbor=n, + state=state, + ) + + print_summary(state) + return 1 if state.failed else 0 + + +if __name__ == "__main__": + sys.exit(main())