chore(scripts): add neighbor expansion benchmark sweep

Co-authored-by: Cursor <cursoragent@cursor.com>

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())