Skip to content

SelfConsistency and RAGAgent

SelfConsistency

class SelfConsistency:
    def __init__(
        self,
        n=5,
        *,
        temperature=0.8,
        model=None,
        system="Think step by step, then end with 'Answer: <final answer>'.",
        extract=default_extract,
    ): ...

Ask the same question n times in parallel at a higher temperature and return the majority answer.

"""Majority vote over parallel samples (SelfConsistency), and a separate LLM critic loop."""

import asyncio
import json

from aaa.algos import GeneratorCritic, SelfConsistency
from aaa.llm import FakeLLMClient, LLMHandler
from aaa.llm.fake import text
from aaa.runtime import Engine

samples = iter(["... Answer: 42", "Answer: 41", "so Answer: 42.", "Answer: 42", "Answer: 7"])


def critic_model(request):
    system = request.messages[0].content if request.messages[0].role == "system" else ""
    if "strict reviewer" in system:  # the critic: accept only the revised draft
        ok = "v2" in request.messages[-1].content
        return text(json.dumps({"ok": ok, "feedback": "" if ok else "mention v2"}))
    revised = any("mention v2" in (m.content or "") for m in request.messages)
    return text("draft v2" if revised else "draft v1")


async def main() -> None:
    async with Engine() as engine:
        LLMHandler(FakeLLMClient(lambda request: text(next(samples)))).attach(engine)
        vote = await engine.run(SelfConsistency(n=5), {"task": "6 x 7?"})
    print("vote:", vote.value)

    async with Engine() as engine:
        LLMHandler(FakeLLMClient(critic_model)).attach(engine)
        out = await engine.run(GeneratorCritic(), {"task": "write a draft"})
    print("critic loop:", out.value)


asyncio.run(main())
vote: {'answer': '42', 'votes': 3, 'samples': 5, 'agreement': 0.6}
critic loop: {'candidate': 'draft v2', 'iterations': 2}
  • extract(text) -> str pulls the comparable answer from each sample. The default takes the last Answer: ... line, else the last line.
  • Answers are normalized (lowercased, punctuation trimmed, trailing period dropped) before voting, so 42, 42. and Answer: 42. agree.
  • Ties go to the earliest-seen answer. agreement is the winner's share.
  • Individual sample failures are tolerated; if all fail the run fails.
  • Cost is n times a single call, but latency is one call (they run concurrently).

Use it where answers are exactly comparable (arithmetic, classification, extracted fields). It does nothing useful for open-ended prose.

RAGAgent

class RAGAgent:
    def __init__(self, index, *, k=5, **tool_loop_kwargs): ...

A ToolLoop preconfigured with a search_documents tool over any Index and a system prompt that demands grounding and citations ("answer using only retrieved evidence; cite ids like [d1]; say so if the evidence is insufficient").

"""Agentic RAG: the model decides when to search your documents, then answers with citations."""

import asyncio

from aaa.algos import RAGAgent
from aaa.data import Doc, InMemoryIndex
from aaa.llm import FakeLLMClient, LLMHandler
from aaa.llm.fake import text, tool_call
from aaa.runtime import Engine
from aaa.tools import ToolHandler


class WordEmbedder:  # any object with `async embed(list[str]) -> list[vector]`; use a real model in prod
    VOCAB = ("paris", "france", "tokyo", "japan", "capital")

    async def embed(self, texts):
        return [[float(w in t.lower()) for w in self.VOCAB] for t in texts]


def model(request):
    last = request.messages[-1]
    if last.role == "user":
        return tool_call("search_documents", query="capital of France")
    return text("Paris is the capital of France [d1].")


async def main() -> None:
    index = InMemoryIndex(WordEmbedder())  # BM25 keyword search fused with vector similarity
    await index.upsert(
        [
            Doc(id="d1", text="Paris is the capital of France.", meta={"source": "atlas"}),
            Doc(id="d2", text="Tokyo is the capital of Japan.", meta={"source": "atlas"}),
        ]
    )
    agent = RAGAgent(index, k=1)
    async with Engine() as engine:
        LLMHandler(FakeLLMClient(model)).attach(engine)
        ToolHandler(agent.tools).attach(engine)
        outcome = await engine.run(agent, {"task": "What is the capital of France?"})
    print(outcome.value["text"])


asyncio.run(main())
Paris is the capital of France [d1].
  • The model decides when and what to retrieve, and may search several times with different queries.
  • Results carry id, score, text (truncated) and meta, so citations can point back to your source.
  • Attach the agent's tool executor: ToolHandler(agent.tools).attach(engine).
  • Retrieval quality is the index's: BM25 by default, hybrid with an Embedder. For large corpora implement the Index protocol over your vector database.

Agentic, not prefetched

This is tool-based retrieval. It does not retrieve before the first model call. If you want a fixed retrieve-then-read pipeline, build it with a WorkflowGraph (retrieve node, generate node).