ragaudit-cli 0.3.0__py3-none-any.whl

ragaudit/tune/loop.py

@@ -0,0 +1,208 @@
+"""The closed tuning loop: search configs on the dev split, then score baseline and best once on the test split.
+
+1. Trial 0 is the baseline config. Each later trial asks the strategy for
+   a config, evaluates it on the **dev split only** (cached results are
+   reused) and keeps the best dev score (ties go to the earlier trial).
+2. When the search stops (max trials, space exhausted, or the LLM budget),
+   the baseline and the best config are evaluated on the **locked test
+   split, once each**, through ragaudit.runner.locked. Only aggregate
+   scores come back.
+3. If the best config's dev score beats its test score by more than the
+   threshold, the run is flagged as likely overfit to the dev set.
+
+The loop sees questions only through DevSet; it never touches the test
+split except through `evaluate_test_split`.
+
+LLM budget: --max-llm-calls covers the judge and the proposer. Room is
+reserved up front for the final test evaluation (two configs x the test
+set), and the search stops before a trial that might not fit, so the
+honest test score is always produced.
+"""
+
+import uuid
+from collections.abc import Callable
+
+from ragaudit.adapters.base import AdapterFactory
+from ragaudit.llm.factory import BudgetedCompleter, CallBudget, Completer
+from ragaudit.models.evaluation import FailureClass, QuestionResult
+from ragaudit.models.question import Split, utc_now
+from ragaudit.models.tune import Trial, TuneResult
+from ragaudit.runner.evaluate import compute_config_hash, run_evaluation
+from ragaudit.runner.locked import count_test_questions, evaluate_test_split
+from ragaudit.scoring.correctness import judge_correctness
+from ragaudit.storage.store import Store
+from ragaudit.tune.devset import DevSet
+from ragaudit.tune.search_space import SearchSpace, canonical
+from ragaudit.tune.strategies import make_strategy
+
+DEFAULT_MAX_TRIALS = 20
+DEFAULT_OVERFIT_THRESHOLD = 10.0
+FAILED_EXAMPLES_FOR_PROPOSER = 5
+
+
+class TuneError(Exception):
+    """Raised when tuning can't start (no dev questions, budget too small for the final evaluation, ...)."""
+
+
+def _judge_for(client: Completer):
+    def judge(question: str, expected_answer: str, candidate_answer: str):
+        return judge_correctness(client, question, expected_answer, candidate_answer)
+
+    return judge
+
+
+def _feedback(best: Trial, best_results: list[QuestionResult], dev: DevSet) -> dict:
+    """Dev-only feedback for the proposer: the best config's failure counts and a few failed dev examples."""
+    by_id = {q.question_id: q for q in dev.questions}
+    failed = [r for r in best_results if r.failure_class != FailureClass.PASS][:FAILED_EXAMPLES_FOR_PROPOSER]
+    return {
+        "best_config": best.config,
+        "best_dev_score": best.dev.score,
+        "failure_counts": best.dev.failure_counts,
+        "failed_examples": [
+            {
+                "question": by_id[r.question_id].question,
+                "gold_answer": by_id[r.question_id].gold_answer,
+                "pipeline_answer": r.answer,
+                "failure_class": r.failure_class.value,
+                "gold_chunk_retrieved": r.retrieval_hit,
+            }
+            for r in failed
+            if r.question_id in by_id
+        ],
+    }
+
+
+def run_tune(
+    store: Store,
+    adapter_factory: AdapterFactory,
+    adapter_sha256: str,
+    adapter_path: str,
+    space: SearchSpace,
+    judge_client: Completer,
+    index_version: str,
+    *,
+    strategy: str = "grid",
+    proposer_client: Completer | None = None,
+    max_trials: int = DEFAULT_MAX_TRIALS,
+    max_llm_calls: int | None = None,
+    seed: int = 0,
+    k: int = 5,
+    overfit_threshold: float = DEFAULT_OVERFIT_THRESHOLD,
+    on_trial: Callable[[Trial, Trial], None] | None = None,
+) -> TuneResult:
+    dev = DevSet(store)
+    if not dev.questions:
+        raise TuneError("no active dev questions; run `ragaudit generate --corpus <path>` first.")
+    n_test = count_test_questions(store)
+    if n_test == 0:
+        raise TuneError("no active test questions to score the final configs on; generate more questions first.")
+    if max_trials < 1:
+        raise TuneError("--max-trials must be at least 1")
+
+    final_reserve = 2 * n_test
+    if max_llm_calls is not None and max_llm_calls < final_reserve + len(dev):
+        raise TuneError(
+            f"--max-llm-calls {max_llm_calls} is too small: the baseline needs up to {len(dev)} judge calls on dev "
+            f"and the final test evaluation reserves {final_reserve}."
+        )
+    search_budget = CallBudget(None if max_llm_calls is None else max_llm_calls - final_reserve)
+    judge = _judge_for(BudgetedCompleter(judge_client, search_budget))
+    proposer = BudgetedCompleter(proposer_client, search_budget) if proposer_client is not None else None
+    chooser = make_strategy(strategy, space, seed, proposer)
+
+    result = TuneResult(
+        tune_id=str(uuid.uuid4()),
+        created_at=utc_now(),
+        strategy=strategy,
+        search_space=space.knobs,
+        adapter_path=adapter_path,
+        judge_model=judge_client.model,
+        index_version=index_version,
+        k=k,
+        seed=seed,
+        max_trials=max_trials,
+        overfit_threshold=overfit_threshold,
+        stopped_reason="max_trials",
+    )
+    store.save_tune_run(result.tune_id, result.created_at, strategy, space.knobs)
+
+    def config_hash(config: dict) -> str:
+        return compute_config_hash(config, adapter_sha256, judge_client.model, k)
+
+    def evaluate_dev(config: dict) -> tuple:
+        return run_evaluation(
+            store,
+            dev.questions,
+            Split.DEV,
+            adapter_factory(config),
+            judge,
+            config=config,
+            config_hash=config_hash(config),
+            index_version=index_version,
+            k=k,
+            full=True,
+        )
+
+    trials: list[Trial] = []
+    best_results: list[QuestionResult] = []
+    next_proposal = (space.baseline, "baseline", None)
+    while next_proposal is not None:
+        config, source, rationale = next_proposal
+        summary, results = evaluate_dev(config)
+        trial = Trial(
+            index=len(trials),
+            config=config,
+            config_hash=config_hash(config),
+            source=source,
+            rationale=rationale,
+            dev=summary,
+            created_at=utc_now(),
+        )
+        trials.append(trial)
+        store.save_tune_trial(result.tune_id, trial)
+        if trial.dev.score > trials[result.best_index].dev.score or trial.index == 0:
+            result.best_index = trial.index
+            best_results = results
+        if on_trial is not None:
+            on_trial(trial, trials[result.best_index])
+
+        if len(trials) >= max_trials:
+            result.stopped_reason = "max_trials"
+            break
+        worst_case_trial = len(dev) + (2 if proposer is not None else 0)
+        if search_budget.remaining is not None and search_budget.remaining < worst_case_trial:
+            result.stopped_reason = "max_llm_calls"
+            break
+        proposal = chooser.propose(trials, _feedback(trials[result.best_index], best_results, dev))
+        if proposal is None:
+            result.stopped_reason = "search_space_exhausted"
+            break
+        next_proposal = (proposal.config, proposal.source, proposal.rationale)
+
+    result.trials = trials
+    final_budget = CallBudget(None)
+    final_judge = _judge_for(BudgetedCompleter(judge_client, final_budget))
+
+    def evaluate_test(config: dict):
+        return evaluate_test_split(
+            store,
+            adapter_factory(config),
+            final_judge,
+            config=config,
+            config_hash=config_hash(config),
+            index_version=index_version,
+            k=k,
+            full=True,
+        )
+
+    result.baseline_test = evaluate_test(result.baseline.config)
+    same_config = canonical(result.best.config) == canonical(result.baseline.config)
+    result.best_test = result.baseline_test if same_config else evaluate_test(result.best.config)
+
+    result.overfit_gap = round(result.best.dev.score - result.best_test.score, 2)
+    result.overfit_warning = result.overfit_gap > overfit_threshold
+    result.llm_calls = search_budget.used + final_budget.used
+    result.finished_at = utc_now()
+    store.finish_tune_run(result.tune_id, result.finished_at, result.model_dump_json())
+    return result

ragaudit/tune/search_space.py

@@ -0,0 +1,140 @@
+"""The tuning search space: a YAML file of knobs the pipeline accepts, each with a finite list of values.
+
+    knobs:
+      chunk_size: [20, 40, 80]
+      top_k: [1, 3, 5]
+      hybrid_weight: {min: 0.0, max: 1.0, step: 0.25}   # expands to 0.0, 0.25, ..., 1.0
+      reranker: [none, overlap]
+    baseline:            # optional; defaults to the first value of each knob
+      chunk_size: 40
+      top_k: 1
+
+Every config a strategy proposes must assign every knob a value from its
+list; `contains` checks that.
+"""
+
+import itertools
+import json
+import random
+from collections.abc import Iterator
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from ragaudit import yaml_io
+
+MAX_RANGE_VALUES = 1000
+
+
+class SearchSpaceError(ValueError):
+    """Raised for a malformed search-space file or an out-of-space config."""
+
+
+def _expand_range(knob: str, spec: dict) -> list:
+    try:
+        low, high, step = spec["min"], spec["max"], spec["step"]
+    except KeyError as exc:
+        raise SearchSpaceError(f"knob {knob!r}: a range needs min, max and step") from exc
+    if not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in (low, high, step)):
+        raise SearchSpaceError(f"knob {knob!r}: min, max and step must be numbers")
+    if step <= 0 or high < low:
+        raise SearchSpaceError(f"knob {knob!r}: need step > 0 and max >= min")
+    count = int(round((high - low) / step)) + 1
+    if count > MAX_RANGE_VALUES:
+        raise SearchSpaceError(f"knob {knob!r}: range expands to {count} values (max {MAX_RANGE_VALUES})")
+    integral = all(isinstance(v, int) for v in (low, high, step))
+    return [low + i * step if integral else round(low + i * step, 10) for i in range(count)]
+
+
+def _same_value(a, b) -> bool:
+    if isinstance(a, (int, float)) and isinstance(b, (int, float)) and not isinstance(a, bool) and not isinstance(b, bool):
+        return abs(a - b) < 1e-9
+    return a == b
+
+
+def canonical(config: dict) -> str:
+    """A stable string identity for a config, for "already tried?" checks."""
+    return json.dumps(config, sort_keys=True, default=str)
+
+
+class SearchSpace(BaseModel):
+    knobs: dict[str, list]
+    baseline: dict
+
+    @property
+    def size(self) -> int:
+        size = 1
+        for values in self.knobs.values():
+            size *= len(values)
+        return size
+
+    def normalize(self, config: dict) -> dict:
+        """Return config with each value replaced by the space's own value; raise if anything is outside the space."""
+        if set(config) != set(self.knobs):
+            missing, extra = set(self.knobs) - set(config), set(config) - set(self.knobs)
+            raise SearchSpaceError(f"config must set exactly the declared knobs (missing {sorted(missing)}, extra {sorted(extra)})")
+        normalized = {}
+        for knob, values in self.knobs.items():
+            match = next((v for v in values if _same_value(v, config[knob])), None)
+            if match is None and not any(_same_value(v, config[knob]) for v in values):
+                raise SearchSpaceError(f"{knob}={config[knob]!r} is outside the search space {values}")
+            normalized[knob] = match
+        return normalized
+
+    def contains(self, config: dict) -> bool:
+        try:
+            self.normalize(config)
+        except SearchSpaceError:
+            return False
+        return True
+
+    def grid(self) -> Iterator[dict]:
+        """Every config, in a fixed order (the first knob varies slowest)."""
+        names = list(self.knobs)
+        for combo in itertools.product(*(self.knobs[n] for n in names)):
+            yield dict(zip(names, combo))
+
+    def random_config(self, rng: random.Random) -> dict:
+        return {knob: rng.choice(values) for knob, values in self.knobs.items()}
+
+
+def parse_search_space(data) -> SearchSpace:
+    if not isinstance(data, dict):
+        raise SearchSpaceError("search space must be a mapping with a `knobs` section")
+    raw_knobs = data.get("knobs", data if "baseline" not in data else None)
+    if not isinstance(raw_knobs, dict) or not raw_knobs:
+        raise SearchSpaceError("search space needs a non-empty `knobs` mapping")
+
+    knobs: dict[str, list] = {}
+    for knob, spec in raw_knobs.items():
+        if isinstance(spec, dict):
+            values = _expand_range(knob, spec)
+        elif isinstance(spec, list):
+            values = list(spec)
+        else:
+            values = [spec]
+        if not values:
+            raise SearchSpaceError(f"knob {knob!r} has no values")
+        knobs[str(knob)] = values
+
+    baseline_spec = data.get("baseline") or {}
+    if not isinstance(baseline_spec, dict):
+        raise SearchSpaceError("`baseline` must be a mapping of knob: value")
+    unknown = set(baseline_spec) - set(knobs)
+    if unknown:
+        raise SearchSpaceError(f"baseline sets unknown knob(s): {sorted(unknown)}")
+    space = SearchSpace(knobs=knobs, baseline={})
+    baseline = {knob: baseline_spec.get(knob, values[0]) for knob, values in knobs.items()}
+    space.baseline = space.normalize(baseline)
+    return space
+
+
+def load_search_space(path: str | Path) -> SearchSpace:
+    path = Path(path)
+    if not path.is_file():
+        raise SearchSpaceError(f"search space file not found: {path}")
+    try:
+        data = yaml_io.load(path)
+    except yaml_io.YAMLError as exc:
+        raise SearchSpaceError(f"could not parse {path}: {exc}") from exc
+    return parse_search_space(data)

ragaudit/tune/strategies.py

@@ -0,0 +1,197 @@
+"""How `tune` picks the next config: grid, seeded random, or an LLM reading dev-set failures.
+
+Every strategy only ever sees dev-split information: trial configs, dev
+scores, dev failure-class counts and a few failed dev examples. Proposals
+are validated against the declared search space; an LLM proposal that is
+outside it (or already tried) gets one corrective retry, then a seeded
+random config is used instead and the trial is marked `llm_fallback`.
+"""
+
+import itertools
+import json
+import random
+import re
+from ragaudit.llm.factory import Completer
+from ragaudit.models.tune import Trial
+from ragaudit.tune.search_space import SearchSpace, SearchSpaceError, canonical
+
+STATE_MARKER = "STATE_JSON:"
+
+PROPOSER_SYSTEM_PROMPT = (
+    "You are tuning a RAG pipeline. You will get the search space (every knob and its allowed values), the "
+    "configs tried so far with their dev-set scores (pass rate in points: the gold passage was retrieved AND "
+    "the answer was correct), the failure-class counts of the best config so far, and a few failed dev "
+    "examples. Failure classes: retrieval_miss = the gold passage was not in the top k; generation_miss = it "
+    "was retrieved but the answer was still wrong.\n"
+    "\n"
+    "Propose ONE new config, not yet tried, that is likely to score higher. Every knob must be set, and every "
+    "value must be one of that knob's allowed values exactly. Reason from the failures: retrieval misses call "
+    "for retrieval knobs, generation misses for generation knobs.\n"
+    "\n"
+    "Respond with strict JSON only, no prose, no markdown fences, in the form: "
+    '{"config": {knob: value, ...}, "rationale": str}'
+)
+
+RETRIEVAL_KNOB_ORDER = ["top_k", "hybrid_weight", "reranker", "chunk_size", "chunk_overlap"]
+GENERATION_KNOB_ORDER = ["prompt_variant", "chunk_size", "chunk_overlap", "top_k", "reranker"]
+
+
+class Proposed:
+    """A strategy's next config, where it came from, and (for llm) why."""
+
+    def __init__(self, config: dict, source: str, rationale: str | None = None) -> None:
+        self.config = config
+        self.source = source
+        self.rationale = rationale
+
+
+def _tried(trials: list[Trial]) -> set[str]:
+    return {canonical(t.config) for t in trials}
+
+
+class GridStrategy:
+    name = "grid"
+
+    def __init__(self, space: SearchSpace) -> None:
+        self._grid = space.grid()
+
+    def propose(self, trials: list[Trial], feedback: dict) -> Proposed | None:
+        tried = _tried(trials)
+        for config in self._grid:
+            if canonical(config) not in tried:
+                return Proposed(config, "grid")
+        return None
+
+
+class RandomStrategy:
+    name = "random"
+
+    def __init__(self, space: SearchSpace, seed: int) -> None:
+        self.space = space
+        self.rng = random.Random(seed)
+        self._order: list[dict] | None = None
+
+    def propose(self, trials: list[Trial], feedback: dict) -> Proposed | None:
+        tried = _tried(trials)
+        if self.space.size <= 100_000:
+            if self._order is None:
+                self._order = list(self.space.grid())
+                self.rng.shuffle(self._order)
+            for config in self._order:
+                if canonical(config) not in tried:
+                    return Proposed(config, "random")
+            return None
+        for _ in range(1000):
+            config = self.space.random_config(self.rng)
+            if canonical(config) not in tried:
+                return Proposed(config, "random")
+        return None
+
+
+def build_proposer_prompt(space: SearchSpace, trials: list[Trial], feedback: dict) -> str:
+    state = {
+        "search_space": space.knobs,
+        "trials": [
+            {"config": t.config, "dev_score": t.dev.score, "failure_counts": t.dev.failure_counts} for t in trials
+        ],
+        "best_config": feedback.get("best_config"),
+        "best_dev_score": feedback.get("best_dev_score"),
+        "failure_counts": feedback.get("failure_counts", {}),
+        "failed_dev_examples": feedback.get("failed_examples", []),
+    }
+    return f"{STATE_MARKER}\n{json.dumps(state, indent=1, default=str)}"
+
+
+def parse_proposal(raw: str) -> tuple[dict, str | None] | None:
+    text = raw.strip()
+    fenced = re.match(r"^```(?:json)?\s*(.*?)\s*```$", text, re.DOTALL)
+    try:
+        data = json.loads(fenced.group(1) if fenced else text)
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(data, dict) or not isinstance(data.get("config"), dict):
+        return None
+    rationale = data.get("rationale")
+    return data["config"], rationale if isinstance(rationale, str) else None
+
+
+class LLMStrategy:
+    name = "llm"
+
+    def __init__(self, space: SearchSpace, client: Completer, seed: int) -> None:
+        self.space = space
+        self.client = client
+        self._fallback = RandomStrategy(space, seed)
+
+    def propose(self, trials: list[Trial], feedback: dict) -> Proposed | None:
+        tried = _tried(trials)
+        if len(tried) >= self.space.size:
+            return None
+        prompt = build_proposer_prompt(self.space, trials, feedback)
+        problem = None
+        for _ in range(2):
+            user = prompt if problem is None else f"{prompt}\n\nYour previous proposal was rejected: {problem}. Try again."
+            parsed = parse_proposal(self.client.complete(PROPOSER_SYSTEM_PROMPT, user))
+            if parsed is None:
+                problem = "it was not valid JSON of the form {\"config\": {...}, \"rationale\": str}"
+                continue
+            config, rationale = parsed
+            try:
+                config = self.space.normalize(config)
+            except SearchSpaceError as exc:
+                problem = str(exc)
+                continue
+            if canonical(config) in tried:
+                problem = f"config {canonical(config)} was already tried"
+                continue
+            return Proposed(config, "llm", rationale)
+
+        fallback = self._fallback.propose(trials, feedback)
+        if fallback is None:
+            return None
+        return Proposed(fallback.config, "llm_fallback", f"LLM proposal rejected ({problem}); used a random config")
+
+
+def make_strategy(name: str, space: SearchSpace, seed: int, client: Completer | None = None):
+    if name == "grid":
+        return GridStrategy(space)
+    if name == "random":
+        return RandomStrategy(space, seed)
+    if name == "llm":
+        if client is None:
+            raise ValueError("the llm strategy needs an LLM client")
+        return LLMStrategy(space, client, seed)
+    raise ValueError(f"unknown strategy {name!r}; expected grid, random or llm")
+
+
+def mock_proposal(user: str) -> dict:
+    """What MockLLM answers to the proposer prompt: coordinate ascent steered by the dominant failure class.
+
+    Starting from the best config, change one knob, picking the knob list
+    for whichever failure class dominated the most recent trial (the
+    latest evidence), and prefer the next value up before going down.
+    Deterministic.
+    """
+    state = json.loads(user.split(STATE_MARKER, 1)[1].split("\n\nYour previous proposal", 1)[0])
+    space: dict[str, list] = state["search_space"]
+    tried = {canonical(t["config"]) for t in state["trials"]}
+    best = state["best_config"] or {knob: values[0] for knob, values in space.items()}
+    failures = (state["trials"][-1].get("failure_counts") if state["trials"] else None) or state.get("failure_counts") or {}
+    retrieval_heavy = failures.get("retrieval_miss", 0) >= failures.get("generation_miss", 0)
+    preferred = RETRIEVAL_KNOB_ORDER if retrieval_heavy else GENERATION_KNOB_ORDER
+    order = [k for k in preferred if k in space] + [k for k in space if k not in preferred]
+    dominant = "retrieval misses" if retrieval_heavy else "generation misses"
+
+    for knob in order:
+        values = space[knob]
+        position = values.index(best[knob]) if best[knob] in values else 0
+        for candidate in values[position + 1 :] + list(reversed(values[:position])):
+            config = dict(best, **{knob: candidate})
+            if canonical(config) not in tried:
+                return {"config": config, "rationale": f"{dominant} dominate, so try {knob}={candidate}"}
+
+    for combo in itertools.product(*space.values()):
+        config = dict(zip(space, combo))
+        if canonical(config) not in tried:
+            return {"config": config, "rationale": "every single-knob change was tried; next untried config"}
+    return {"config": best, "rationale": "search space exhausted"}