Multi-agent RAG pipelines

pipeline.py"""Multi-agent RAG pipeline — ORCHESTRATION only, fully offline & deterministic.
No API key, no cluster: fake embed/retrieve + grader/synth/verify agents + the loop.
Run: python3 pipeline.py"""
from dataclasses import dataclass, field

# ---- shared state (the blackboard every agent reads/writes) ----------------
@dataclass
class State:
    question: str
    plan: str = ""
    sources: tuple = ()
    candidates: list = field(default_factory=list)   # retriever output
    context: list = field(default_factory=list)       # grader-kept chunks
    draft: str = ""
    cites: list = field(default_factory=list)
    verdict: str = ""
    tries: int = 0
    trace: list = field(default_factory=list)

CORPUS = {
    "refund":  "Refunds are accepted within 30 days of purchase.",
    "reset":   "Reset your password in Settings -> Security.",
    "shipping":"Standard shipping takes 3-5 business days.",
}

def _score(query, text):            # deterministic fake "relevance": prefix-term overlap
    q = query.lower().split(); t = text.lower().split()
    # count query words that prefix-match some doc word ("refund" ~ "refunds")
    return sum(any(w.startswith(qw) or qw.startswith(w) for w in t) for qw in q)

# ---- the specialized agents (each does ONE job) ----------------------------
def router(s: State) -> State:
    s.plan = "lookup"; s.sources = ("vector",)
    s.trace.append(("router", s.plan)); return s

def retriever(s: State) -> State:                     # pull candidates from the index
    s.candidates = [(k, v, _score(s.question, v)) for k, v in CORPUS.items()]
    s.tries += 1
    s.trace.append(("retriever", f"try#{s.tries} pulled {len(s.candidates)}")); return s

def grader(s: State, keep_min=1) -> State:            # drop junk, keep the relevant
    s.context = [(k, v) for k, v, sc in s.candidates if sc >= keep_min]
    s.trace.append(("grader", f"kept {len(s.context)}/{len(s.candidates)}")); return s

def synthesizer(s: State) -> State:                   # grounded, cited draft
    if s.context:
        s.draft = " ".join(v for _, v in s.context)
        s.cites = [k for k, _ in s.context]
    else:
        s.draft = ""; s.cites = []
    s.trace.append(("synth", f"draft={bool(s.draft)} cites={s.cites}")); return s

def verifier(s: State) -> State:                      # faithful to context? every
    ctx = " ".join(v for _, v in s.context).lower()   # answer word must appear in ctx
    unsupported = [w for w in s.draft.lower().split() if w not in ctx]
    s.verdict = "faithful" if s.draft and not unsupported else "unfaithful"
    s.trace.append(("verifier", s.verdict)); return s

# ---- orchestrator: the state machine with BOUNDED back-edges ---------------
def run(question, max_tries=3):
    s = State(question=question)
    s = router(s)
    while True:
        s = retriever(s)
        s = grader(s)
        if not s.context and s.tries < max_tries:     # weak retrieval -> loop back
            s.question += " policy"                    # (a real router would rewrite)
            continue
        s = synthesizer(s)
        s = verifier(s)
        if s.verdict == "faithful":
            return {"answer": s.draft, "cites": s.cites, "trace": s.trace}
        if s.tries >= max_tries:                        # bound the send-back loop
            return {"answer": "I don't have enough grounded information to answer.",
                    "cites": [], "trace": s.trace}
        # unfaithful but budget left -> send back to retrieve for more context

if __name__ == "__main__":
    import json
    print(json.dumps(run("what is the refund window"), indent=2))
    print(json.dumps(run("who is the president"), indent=2))   # ungroundable -> refuse

graph_langgraph.py# needs: pip install langgraph  + a model/API key. Node fns = the agents above.
from langgraph.graph import StateGraph, START, END

g = StateGraph(dict)                       # dict state == the blackboard
g.add_node("router", router)
g.add_node("retriever", retriever)
g.add_node("grader", grader)
g.add_node("synth", synthesizer)
g.add_node("verify", verifier)

g.add_edge(START, "router")
g.add_edge("router", "retriever")
g.add_edge("retriever", "grader")
# grader decides: enough context -> synth, else loop back to retriever (bounded in state)
g.add_conditional_edges("grader",
    lambda s: "synth" if s["context"] or s["tries"] >= 3 else "retriever",
    {"synth": "synth", "retriever": "retriever"})
g.add_edge("synth", "verify")
# verifier decides: faithful -> END, else send back to retriever (bounded)
g.add_conditional_edges("verify",
    lambda s: END if s["verdict"] == "faithful" or s["tries"] >= 3 else "retriever",
    {END: END, "retriever": "retriever"})
app = g.compile()
# CrewAI (needs: pip install crewai + key) models the same as role-based agents with a
# manager process; see M1. AutoGen (M2) would use message-passing between the agents.

budget.py"""A shared budget guard for the pipeline. Offline & deterministic.
Bounds BOTH the loop (iterations) and the spend (model calls). Run: python3 budget.py"""
class Budget:
    def __init__(self, max_iters=3, max_calls=12):
        self.max_iters = max_iters; self.max_calls = max_calls
        self.iters = 0; self.calls = 0

    def spend(self, n=1):                 # every agent call goes through here
        self.calls += n
        if self.calls > self.max_calls:
            raise BudgetExceeded(f"call budget blown: {self.calls}/{self.max_calls}")

    def next_iter(self) -> bool:          # False -> stop looping, give up gracefully
        self.iters += 1
        return self.iters <= self.max_iters

class BudgetExceeded(Exception):
    pass

def guarded_pipeline(question, budget: Budget):
    """Model the loop with the guard deciding when to stop."""
    trace = []
    while budget.next_iter():
        budget.spend(2)                   # e.g. retriever + grader call
        grounded = "refund" in question.lower()   # fake: only 'refund' is groundable
        trace.append((budget.iters, budget.calls, "hit" if grounded else "miss"))
        if grounded:
            budget.spend(2)               # synth + verify
            return {"answer": "Refunds within 30 days.", "cost": budget.calls, "trace": trace}
        question += " policy"             # rewrite & retry within budget
    # loop exhausted WITHOUT a grounded hit -> the honest fallback
    return {"answer": "I don't know — no grounded context found.",
            "cost": budget.calls, "trace": trace}

if __name__ == "__main__":
    print(guarded_pipeline("what is the refund window", Budget()))
    print(guarded_pipeline("what is the meaning of life", Budget()))
    try:
        guarded_pipeline("x", Budget(max_iters=99, max_calls=3))
    except BudgetExceeded as e:
        print("STOPPED:", e)

eval.py"""Pipeline evaluation, offline & deterministic (no judge model needed).
Per-stage recall@k for the retriever + an end-to-end faithfulness gate. Run: python3 eval.py"""

# ---- per-stage: retriever recall@k (FA5 set-overlap idea) ------------------
def recall_at_k(retrieved_ids, gold_ids, k):
    top = set(retrieved_ids[:k]); gold = set(gold_ids)
    return len(top & gold) / len(gold) if gold else 0.0

# ---- end-to-end: faithfulness gate (what the verifier agent enforces) ------
import re
STOP = {"is","are","the","a","an","of","in","to","and","our","within"}

def _words(text):                    # lowercase, strip punctuation ("days." -> "days")
    return re.findall(r"[a-z0-9]+", text.lower())

def faithfulness(answer, context):
    """Fraction of answer content-words supported by the retrieved context."""
    ctx = set(w for c in context for w in _words(c)) - STOP
    ans = [w for w in _words(answer) if w not in STOP]
    if not ans:
        return 0.0
    return sum(w in ctx for w in ans) / len(ans)

def faithfulness_gate(answer, context, threshold=0.8):
    score = faithfulness(answer, context)
    return {"score": round(score, 2),
            "verdict": "pass" if score >= threshold else "FAIL (ungrounded)"}

def evaluate_pipeline(cases):
    """cases: list of (retrieved_ids, gold_ids, answer, context)."""
    recalls, passes = [], 0
    for rid, gid, ans, ctx in cases:
        recalls.append(recall_at_k(rid, gid, k=3))
        if faithfulness_gate(ans, ctx)["verdict"] == "pass":
            passes += 1
    return {"mean_recall@3": round(sum(recalls)/len(recalls), 2),
            "grounded_answer_rate": round(passes/len(cases), 2)}

if __name__ == "__main__":
    print(faithfulness_gate("Refunds are accepted within 30 days.",
                            ["Refunds are accepted within 30 days of purchase."]))
    print(faithfulness_gate("The CEO is Bob.",                       # ungrounded
                            ["Refunds are accepted within 30 days."]))
    cases = [
        (["refund","reset","x"], ["refund"], "Refunds within 30 days.",
         ["Refunds are accepted within 30 days."]),
        (["a","b","c"], ["shipping"], "The CEO is Bob.",             # bad retrieval + ungrounded
         ["Refunds are accepted within 30 days."]),
    ]
    print(evaluate_pipeline(cases))