"""Example (b): a stochastic RAG pipeline with mocked LLM steps sampled via Gumbel-max. Steps (each "LLM" picks one option from logits with ``noise.categorical``): * ``rewrite`` -- query rewrite: precise / vague / wrong_entity * ``retrieve`` -- document: gold / near_miss / stale / irrelevant * ``answer`` -- correct / hallucinated / stale_fact / abstain * ``verify`` -- as-run passes the answer through; the fix re-grounds ungrounded answers on a gold doc and abstains otherwise. Each fix is a "better model": it shifts logits, it does not make the step deterministic. A case has a difficulty ``d`` in [0, 1] that lowers the logits of good options. 200 synthetic cases; the failing ones are analyzed. The script reports 1. counterfactual blame with the *recorded* noise (exact replay), 2. the same with noise *abducted* from the observed choices only (posterior Gumbels), and the interventional ("prior", fresh noise) version, 3. the variance of Shapley estimates with common random numbers (CRN, same worlds for every coalition) vs independent re-sampling per coalition, at equal numbers of pipeline runs. Run: ``python -m agentblame.demos.stochastic_rag [--replicates 8] [--worlds 4] [--var-cases 100]`` """ from __future__ import annotations import argparse import os import time from concurrent.futures import ProcessPoolExecutor import numpy as np from agentblame import (Case, Pipeline, blame_batch, exact_shapley, format_batch, mc_shapley, single_swap, step, write_html) from agentblame.attribution import make_game from agentblame.noise import NoiseSource ANSWERS = ["correct", "hallucinated", "stale_fact", "abstain"] # --- rewrite --------------------------------------------------------------- @step("rewrite", cost=1.0) def rewrite(question, difficulty, noise): d = difficulty return noise.categorical({"precise": 1.2 - 1.2 * d, "vague": 0.5, "wrong_entity": -0.2 + 0.8 * d}) @rewrite.variant("strong_model", cost=3.0) def rewrite_fixed(question, difficulty, noise): d = difficulty return noise.categorical({"precise": 3.0 - d, "vague": 0.0, "wrong_entity": -1.5}) # --- retrieve -------------------------------------------------------------- _RET = {"precise": [2.0, 0.8, 0.8, -1.0], "vague": [0.3, 0.8, 1.0, 0.2], "wrong_entity": [-1.5, 0.2, 0.3, 1.5]} _RET_FIX = {"precise": [3.5, 0.5, -0.5, -2.0], "vague": [2.0, 0.8, 0.0, -1.0], "wrong_entity": [0.0, 0.5, -0.5, 0.8]} _DOCS = ["gold", "near_miss", "stale", "irrelevant"] @step("retrieve", cost=2.0) def retrieve(rewrite, noise): return noise.categorical(dict(zip(_DOCS, _RET[rewrite]))) @retrieve.variant("reranker", cost=4.0) def retrieve_fixed(rewrite, noise): return noise.categorical(dict(zip(_DOCS, _RET_FIX[rewrite]))) # --- answer ---------------------------------------------------------------- _ANS = {"gold": [2.0, 0.3, -1.0, -0.5], "near_miss": [0.3, 1.0, -1.0, 0.2], "stale": [-1.0, 0.2, 2.0, -0.5], "irrelevant": [-1.5, 1.2, -1.0, 0.8]} _ANS_FIX = {"gold": [3.5, -0.5, -2.0, -1.0], "near_miss": [0.8, -0.2, -1.0, 1.0], "stale": [-0.5, -0.5, 1.0, 0.8], "irrelevant": [-2.0, -0.5, -1.5, 1.5]} def _answer(table, doc, difficulty, noise): lg = list(table[doc]) lg[0] -= difficulty return noise.categorical(dict(zip(ANSWERS, lg))) @step("answer", cost=2.0) def answer(retrieve, difficulty, noise): return _answer(_ANS, retrieve, difficulty, noise) @answer.variant("strong_model", cost=6.0) def answer_fixed(retrieve, difficulty, noise): return _answer(_ANS_FIX, retrieve, difficulty, noise) # --- verify ---------------------------------------------------------------- @step("verify", cost=1.0) def verify(answer, retrieve): return answer @verify.variant("grounding_check", cost=2.0) def verify_fixed(answer, retrieve, noise): if answer in ("hallucinated", "stale_fact"): if retrieve == "gold": return noise.categorical({"correct": 1.5, "abstain": 0.0}, key="regrounded") return "abstain" return answer def metric(out) -> float: return 1.0 if out["verify"] == "correct" else 0.0 def build_pipeline() -> Pipeline: return Pipeline([rewrite, retrieve, answer, verify], metric, name="stochastic_rag") def make_cases(n: int = 200, seed: int = 0) -> list[Case]: rng = np.random.default_rng(seed) return [Case(f"q{i:03d}", {"question": f"synthetic question {i}", "difficulty": round(float(rng.uniform()), 3)}, seed=i) for i in range(n)] def variance_experiment(failing: list[Case], mode: str, n_worlds: int, replicates: int) -> dict: """Estimate per-case blame ``replicates`` times with common random numbers (CRN) and with independent worlds per coalition, at equal numbers of pipeline runs. Three estimators: * ``fix-alone`` -- v({i}) - v(empty), each value averaged over n_worlds worlds * ``exact-avg`` -- exact Shapley of the world-averaged game (16 * n_worlds runs) * ``mc-perm`` -- permutation Shapley, a fresh world per permutation, round(16 * n_worlds / 5) permutations (about the same number of runs) Returns, per estimator and sampling scheme, the variance across replicates (averaged over cases and steps) and the step-execution cost per replicate. """ pipe = build_pipeline() # fresh pipeline: this runs in a worker process n_perm = round(16 * n_worlds / 5) res = {} for crn in (True, False): est = {k: np.zeros((replicates, len(failing), 4)) for k in ("fix-alone", "exact-avg", "mc-perm")} cost = {k: 0.0 for k in est} for r in range(replicates): pipe.clear_cache() for c, case in enumerate(failing): c0 = pipe.meter.total g = make_game(pipe, case, mode=mode, n_worlds=n_worlds, crn=crn, seed=1000 + r) sh = exact_shapley(g) est["exact-avg"][r, c] = list(sh.values.values()) est["fix-alone"][r, c] = list(single_swap(g).values()) cost["exact-avg"] += pipe.meter.total - c0 c0 = pipe.meter.total g = make_game(pipe, case, mode=mode, crn=crn, seed=2000 + r) mc = mc_shapley(g, n_perm, sample_worlds=True, seed=r) est["mc-perm"][r, c] = list(mc.values.values()) cost["mc-perm"] += pipe.meter.total - c0 cost["fix-alone"] = float("nan") # by-product of exact-avg for k, e in est.items(): res[(k, "CRN" if crn else "independent")] = { "per_case_var": float(e.var(axis=0, ddof=1).mean()), "mean_phi": e.mean(axis=(0, 1)).round(3).tolist(), "cost": cost[k] / replicates, } pipe.clear_cache() return res def main(argv=None): ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) ap.add_argument("--cases", type=int, default=200) ap.add_argument("--worlds", type=int, default=4, help="noise worlds per coalition value") ap.add_argument("--replicates", type=int, default=8) ap.add_argument("--var-cases", type=int, default=100, help="failing cases used in the variance experiment") ap.add_argument("--out", default="agentblame_out") args = ap.parse_args(argv) t0 = time.time() pipe = build_pipeline() cases = make_cases(args.cases) rec = blame_batch(pipe, cases, mode="recorded") print("1) Counterfactual blame, recorded noise (exact replay of each failure)") print(format_batch(rec), "\n") post = blame_batch(pipe, cases, mode="posterior", n_worlds=8) print("2a) Counterfactual blame, noise abducted from observed choices (8 posterior worlds)") print(format_batch(post), "\n") prior = blame_batch(pipe, cases, mode="prior", n_worlds=8) print("2b) Interventional blame, fresh noise (8 prior worlds) -- a different question") print(format_batch(prior), "\n") failing = [c for c in cases if pipe.run(c.inputs, noise=NoiseSource(c.seed)).score < 1] failing = failing[:args.var_cases] lines = [f"3) Variance of per-case blame estimates: {len(failing)} failing cases, " f"{args.worlds} worlds per coalition, {args.replicates} replicates", f"{'mode':<10}{'estimator':<11}{'Var CRN':>10}{'Var indep':>11}{'indep/CRN':>11}" f"{'cost CRN':>10}{'cost indep':>12}{'Var*cost':>10} mean estimate, CRN " f"(rewrite, retrieve, answer, verify)"] modes = ("posterior", "prior") with ProcessPoolExecutor(max_workers=2) as ex: # at most 2 workers (shared machine) futs = {m: ex.submit(variance_experiment, failing, m, args.worlds, args.replicates) for m in modes} var = {m: f.result() for m, f in futs.items()} for mode in modes: res = var[mode] for est in ("fix-alone", "exact-avg", "mc-perm"): a_, b_ = res[(est, "CRN")], res[(est, "independent")] eff = (b_["per_case_var"] * b_["cost"]) / (a_["per_case_var"] * a_["cost"]) costs = (f"{a_['cost']:>10.0f}{b_['cost']:>12.0f}{eff:>9.2f}x" if np.isfinite(eff) else f"{'(by-product of exact-avg)':>31}") lines.append(f"{mode:<10}{est:<11}{a_['per_case_var']:>10.5f}{b_['per_case_var']:>11.5f}" f"{b_['per_case_var'] / a_['per_case_var']:>10.2f}x{costs} {a_['mean_phi']}") lines.append("Var = variance across replicates, averaged over cases and steps; cost = step " "cost units charged per replicate (cache hits are free); Var*cost = " "(Var*cost)_indep / (Var*cost)_CRN, i.e. efficiency gain of CRN at equal budget. " "Recorded-noise replay has zero estimator variance (it is deterministic).") text = "\n".join(lines) print(text) os.makedirs(args.out, exist_ok=True) path = write_html( os.path.join(args.out, "stochastic_rag.html"), "Stochastic RAG: blame over 200 test cases", [("Counterfactual blame (recorded noise)", rec), ("Counterfactual blame (abducted posterior noise)", post), ("Interventional blame (fresh noise)", prior), ("Variance: common random numbers vs independent re-sampling", text)], intro="rewrite → retrieve → answer → verify; every step is a mocked LLM sampling from " "logits via the Gumbel-max trick. Fixes are stronger models / a reranker / a " "grounding check.") print(f"\nHTML report: {path}\nruntime {time.time() - t0:.1f}s") return rec, post, prior, var if __name__ == "__main__": main()