chunking_strategies_evaluation/scripts/run_neighbor_sweep.py

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