Extending RAGBench

Systems are discovered through a registry and describe themselves, so adding one is: write the class, declare its options, register it, and regenerate the docs. To benchmark a pipeline you already have (a service, or a function written with any framework) without making it a RAGBench system, use external_http or python_callable instead: see benchmarking external systems.

1. Declare the options

Options are a Pydantic model that forbids unknown keys, so a typo in a config fails at load time with a suggestion. Reuse the building blocks in src/ragbench/rag_systems/options.py (TopKOptions, VectorOptions, FusionOptions, …) and give every field a description — it becomes the option table in systems.md.

from pydantic import Field

from ragbench.rag_systems.options import TopKOptions, VectorOptions


class MyOptions(VectorOptions, TopKOptions):
    my_threshold: float = Field(default=0.3, ge=0, le=1, description="Drop candidates scoring below this.")

2. Write the class

Create a file under src/ragbench/rag_systems/ and inherit from BaseRAGSystem. Register it with @SYSTEMS.register(...) and describe it with a SystemSpec:

from ragbench.config.schema import SystemConfig
from ragbench.documents.schema import Document
from ragbench.rag_systems.base import BaseRAGSystem, IngestionResult, RetrievalResult
from ragbench.rag_systems.components import build_chunker, build_embedder, build_vector_index
from ragbench.rag_systems.spec import SystemSpec
from ragbench.registry import SYSTEMS


@SYSTEMS.register("my_rag")
class MyRAGSystem(BaseRAGSystem):
    """One-line description of the strategy."""

    spec = SystemSpec(
        type="my_rag",
        title="My RAG",
        summary="One line for the README table",
        best_for="When this strategy wins",
        cost_profile="low",        # free | low | medium | high
        latency_profile="fast",    # fast | medium | slow
        requires_llm=False,        # does retrieval itself call an LLM?
        agentic=False,             # does it run a multi-step loop?
        # retrieves=False,        # only for systems that fetch nothing: their retrieval columns are blank, not 0
        # chunker=None,           # for systems with no `chunker:` section
        options=MyOptions,
    )
    options: MyOptions             # validated copy of the config's `retrieval:` section

    def __init__(self, config: SystemConfig, force_mock: bool = False):
        super().__init__(config, force_mock=force_mock)
        self.chunker = build_chunker(config.chunker)
        self.embedding_model = build_embedder(config.models, force_mock)
        self.store = build_vector_index(self.embedding_model, self.options, self.name)

    def ingest(self, documents: list[Document]) -> IngestionResult:
        # build whatever indexes you need; report cost on IngestionResult.cost
        ...

    def fetch_context(self, question: str, top_k: int | None = None) -> RetrievalResult:
        k = self.options.resolve_top_k(top_k)
        ...

Use the shared builders (build_chunker, build_embedder, build_vector_index, build_reranker) rather than constructing chunkers, embedding models or stores yourself — they apply the project defaults and mock mode. Read tunables from self.options, never from raw dicts.

The top_k passed to fetch_context is the retrieval depth: return that many ranked chunks (the evaluator scores retrieval on all of them), while only the first context_k are given to the generator. Make sure any internal candidate pool is at least top_k deep.

Prompts belong in the harness. The default answer_question sends self.harness.answer_system() as the generator's system message, so a new system that only changes retrieval gets an editable harness for free. A system that sends a prompt of its own declares its default texts with default_harness and reads them from self.harness, so a harness: patch (and the optimizer) can edit them:

from ragbench.harness.prompts import ANSWER_PROMPTS

    @classmethod
    def default_harness(cls, options):
        return {**ANSWER_PROMPTS, "query_rewrite_prompt": "Rewrite the question as a keyword query."}

    def _rewrite(self, question: str) -> str:
        messages = [{"role": "system", "content": self.harness.rewrite_prompt}, {"role": "user", "content": f"Question: {question}"}]
        ...

Return {} from default_harness for a system whose prompts live elsewhere. Knobs need nothing: a harness: knob sets the retrieval: / chunker: option of the same name when your option model has it.

Record steps

answer_question() traces every question. Wrap each piece of work in self.trace.step(kind, name) so the report can show where time and money go (kinds: retrieve, rerank, llm, tool, embed, route, grade, generate; the final generate step is recorded for you):

with self.trace.step("retrieve", "vector_search", top_k=k) as step:
    step.set_input(question)
    result = self.store.search(question, top_k=k)
    step.set_chunks(result.chunks, result.cost)   # or step.set_llm(llm_result, messages, cost=...) / step.set_cost(...)

Step costs must add up to the cost you return. If they do not, answer_question appends a visible untracked step with the difference, so a gap shows up in the stage columns instead of hiding. A system that records no steps still gets a coarse fetch_context step. Outside answer_question (for example a direct fetch_context call) self.trace is a silent no-op.

BaseRAGSystem.answer_question() already wires fetch_context into the configured LLM, tracks cost, and returns an AnswerResult. Override it only when your system needs custom generation (query rewriting, multi-hop).

Keep the system dataset-agnostic and mock-safe: it must work on any user-supplied corpus and run without an API key (use the builders and self.llm, which fall back to mocks). HyDERAG (hyde_rag.py) is a compact example of a system that adds an LLM step while staying mock-safe.

If your system sends the LLM a new kind of prompt, teach the mock what to answer with MockLLM.register_responder(predicate, respond, name=...) (predicate(prompt, json_mode) -> bool, respond(prompt) -> str; the newest registration wins, MockLLM.unregister_responder(name) removes it). Put the phrase that identifies your prompt in models/prompts.py so the system and the mock share it. For LLM calls made at ingestion, wrap them in rag_systems/llm_ingest.map_llm_calls so one failed call degrades one item instead of the run (ContextualRAG and HierarchicalRAG show both).

Writing an agentic system with AgentLoop

A system that decides, step by step, whether to keep searching (grade and retry, multi-hop, tool use) should not hand-roll its loop. ragbench.agents.AgentLoop owns the stopping rules; you write a policy, a function that does one iteration of work and returns what to do next:

from ragbench.agents import BUDGET, Abort, AgentBudget, AgentLoop, Continue, Finish

def policy(state):                      # state.iteration is 1 on the first call; state.query is the current query
    chunks = my_search(state.query)     # record work with self.trace.step(...) as usual
    if enough(chunks):
        return Finish()                 # or Finish(answer) when the policy writes the answer itself
    if state.over_budget():             # check before every extra LLM call
        return Abort(BUDGET)
    return Continue(better_query(state.query))

state = AgentLoop(AgentBudget(max_steps=4, max_cost_usd=0.02), tracer).run(policy, question=question)
state.metadata()  # {"steps", "hops", "budget_exhausted", "termination": finished | budget | max_steps | abort, ...state.info}

The loop stops at max_steps whatever the policy returns (never more than 50), and before an iteration when this question's traced spending has reached max_cost_usd or max_tokens. It never raises for a spent budget: it ends with termination="budget" and you answer from the evidence gathered so far. Each decision is recorded as a free route step (agent_decision, agent_stop). Put the loop inside fetch_context and return the chunks; AgenticRAG (rag_systems/agentic_base.py) does the wiring for you: it sums everything the agent spent into the retrieval cost, offers self._ask(kind, name, messages, parse, ...) for traced JSON control calls, and adds llm_calls to the agent block of AnswerResult.metadata, which the evaluator persists as agent in per_question_results.jsonl (summary columns avg_steps, avg_llm_calls, budget_exhausted_rate). For a system whose model calls tools, agents/tool_agent.py (ToolAgent, native or ReAct-JSON) and rag_systems/tool_agent_base.py (ToolAgentRAG) already implement the loop, the budgets and the trace; AgentSearchRAG and GrepAgentRAG are two short subclasses (see tools.md). Keep prompt text in agents/prompts.py and reply parsing in agents/replies.py; corrective_rag.py and iterative_rag.py are the worked examples.

A system that calls tools builds a ToolBox from its tools: list; see tools.md for the tool types, custom tools and the safety model.

3. Register the import

Add the import to src/ragbench/rag_systems/__init__.py so the module loads (and its decorator runs) with the package. Then regenerate the docs, which derive from your SystemSpec and option model:

python scripts/generate_docs.py

CI (--check) and tests/test_docs_fresh.py fail when the generated files are stale.

4. Add a config and run it

Copy one of the configs/*.yaml files and change the systems: block (the schema is in configuration.md), then:

ragbench run --config configs/my_rag.yaml
ragbench compare --config configs/all.yaml   # head-to-head with the others

Add a test under tests/ that exercises ingestion and retrieval against the demo dataset in mock mode (no API key). See tests/test_bm25_system.py for a minimal example.

Chunkers and rerankers

Both use the same mechanism: decorate the class with @CHUNKERS.register("name", aliases=(...)) (one file per chunker under documents/chunkers/) or @RERANKERS.register("name") (one file per reranker under models/rerankers/). Names are then accepted in configs (chunker.type, retrieval.reranker) and unknown names are rejected with a did-you-mean suggestion.

A chunker subclasses BaseChunker and implements _spans(document) -> list[Span], returning character ranges in document order; the base class turns them into chunks (no empty chunks, start_char/end_char that reproduce the text, deterministic ids), so a new chunker only decides where to cut. Declare the ChunkerConfig fields it honours beyond the common ones in options (min_chunk_size, prefix_heading, breakpoint_percentile), its defaults in default_chunk_size / default_chunk_overlap, and needs_embedder = True if it embeds text (charge that cost through self.last_cost). documents/chunkers/_split.py has the splitting and packing helpers; recursive.py is the smallest structure-aware example.

A reranker is any class with name and rerank(question, chunks, top_k) -> RerankResult. Return at most top_k chunks ranked 1..n, each with metadata["reranker"] and metadata["original_rank"], and charge any per-call spend in RerankResult.cost.rerank_cost. Optional class attributes tell create_reranker how to build it: needs_llm = True (gets the system's LLM), takes_model = True (gets retrieval.reranker_model), offline_fallback = "<name>" (mock runs use that reranker instead, so --mock never downloads a model), and import_modules (libraries imported on the main thread before worker threads start). models/rerankers/cross_encoder.py is a complete example: lazy, shared model load, serialized prediction, actionable error when the extra is missing.

The built-in cross_encoder reranker is local. A hosted reranking API (Cohere, Voyage, Jina, ...) is deliberately not built in: each has its own per-search pricing that would need a verified price row, and nothing in CI can exercise it. Ship one as a plugin instead (entry point below): call the API inside rerank, sort by the returned relevance scores, and put the spend in rerank_cost.

Document loaders

A loader turns one file into documents: loader(path, context) -> list[Document], registered for a file extension with @LOADERS.register(".ext", aliases=(".other",)) (one module per format under documents/loaders/). Build documents with make_document(path, context, title, text, extra_metadata, doc_id=None) so ids follow the usual rules (files of many rows use doc_id=f"{base}#{row}"). Decode bytes with decode_bytes (or read_text_file) to get the shared encoding fallbacks; report recoverable problems with context.warn(message); raise DocumentLoadError for a file that cannot be used (it is raised or skipped according to dataset.on_error); raise MissingExtraError when an optional library is missing (always fatal). Put character spans your format knows about (PDF pages) in metadata["page_spans"] as [start, end, page] triples to give chunks a page.

Vector backends

A backend is a class registered with @VECTOR_BACKENDS.register("name") (one file per backend under stores/index/). It is constructed with (collection_name=, persist_directory=), receives unit-length vectors in build(ids, vectors, payloads), and returns (row index, cosine score) pairs, best first, from search(query, top_k). Declare the class attributes backend, approximate, persistent (honours persist_directory) and import_modules (libraries to import on the main thread before worker threads start). Check the optional library in __init__ with require_extra(...) so a missing extra fails when the system is built, and never fall back to another backend. stores/index/faiss_index.py is the smallest complete example.

Model providers

A provider turns the part of a model ref after provider: into a model object. Write a factory and register it in LLM_PROVIDERS (chat, returns an LLM) and/or EMBEDDERS (returns an EmbeddingModel):

from ragbench.models.llms import LLM, LLMResult
from ragbench.registry import LLM_PROVIDERS


class MyLLM(LLM):
    provider = "my_provider"          # part of the disk-cache key: two providers may serve the same model name

    def __init__(self, model: str):
        self.model_name = model       # what reports and price lookups show

    def generate(self, messages, *, temperature=0.0, max_tokens=None, json_mode=False, tools=None) -> LLMResult: ...


@LLM_PROVIDERS.register("my_provider")
def make_llm(model: str, *, providers) -> LLM:    # `providers`: the config's endpoint entries
    return MyLLM(model)

my_provider:some-model is then a valid ref everywhere. Messages and tools use the OpenAI chat shapes; translate them inside the provider and fill LLMResult.tool_calls (see models/providers/anthropic.py for a worked translation). Keep retries (call_with_retry) and limiter_for("<provider>") around every request, import optional SDKs lazily (raise MissingExtraError naming the extra), and add a price row to models/cost.py if the provider bills per token.

Tools

A tool is a typed function an agent can call. Most are just a Python function pointed to from the config (tools: [{name: my_tool, path: "pkg.module:function"}]); tools.md covers the signature rules, the context argument, side effects and the safety model. To ship a tool as a package, register a class in TOOLS (@TOOLS.register("name"), see src/ragbench/tools/calculator.py for the smallest one) or use the entry point below.

Plugins from other packages

A separate package can contribute components without touching this repository by declaring entry points. This table is the one list of the groups; each component is described in its own section above.

Entry-point group	Contributes	The entry point names
ragbench.systems	RAG systems	a BaseRAGSystem subclass
ragbench.chunkers	chunkers	a BaseChunker subclass
ragbench.rerankers	rerankers	a reranker class
ragbench.loaders	document loaders	a loader function; the entry point's name is the file extension with its dot (".epub", quoted in TOML)
ragbench.vector_backends	vector backends	a backend class
ragbench.llm_providers	chat model providers	a factory f(model, *, providers) -> LLM
ragbench.embedders	embedding providers	a factory returning an EmbeddingModel
ragbench.decision_models	decision models for decision_gate	a factory f(model) -> DecisionModel
ragbench.tools	agent tools	a tool class

[project.entry-points."ragbench.systems"]
my_rag = "my_package.rag:MyRAGSystem"

[project.entry-points."ragbench.loaders"]
".epub" = "my_package.epub:load_epub"

Plugins are loaded when the component's package is imported (ragbench.rag_systems for systems, chunkers and rerankers). A plugin that fails to import is skipped with a warning, and a name that is already taken is ignored. Give a plugin system a spec to get config validation and documentation; without one its retrieval: section is accepted as-is.

Decision models

decision_gate asks a decision model typed yes/no questions and acts on the probabilities. A backend is a class with decide(text, predicates) -> DecisionResult, registered with @DECISION_MODELS.register("name") (or DECISION_MODELS.add), and then selectable as retrieval.decision_model.provider:

from collections.abc import Sequence

from ragbench.models.cost import CostBreakdown
from ragbench.models.decisions import DecisionModel, DecisionResult, Predicate
from ragbench.registry import DECISION_MODELS


class MyDecisionModel(DecisionModel):
    provider = "my_decisions"          # part of the cache key
    model_name = "my-model-v1"

    def decide(self, text: str, predicates: Sequence[Predicate]) -> DecisionResult:
        probabilities = call_my_api(text, [(p.name, p.instructions) for p in predicates])  # name -> P(true), or None for "no answer"
        tokens = ...  # input tokens you were billed for
        cost = CostBreakdown(decision_input_tokens=tokens, decision_cost=tokens / 1_000_000 * MY_PRICE_PER_1M)
        return DecisionResult(probabilities=probabilities, model=self.model_name, input_tokens=tokens, cost=cost)


DECISION_MODELS.add("my_decisions", lambda model: MyDecisionModel())   # factory: model name from the config, or None

Return None, never 0, for a question the model declined or you could not parse: the gate keeps a chunk it knows nothing about, and a made-up 0 would silently drop evidence. Book spend as decision_cost so it shows up in its own column and in the stage tables. Wrap the backend in CachedDecisionModel inside your factory to get the run's disk cache (same rules as the chat cache: hits are charged at the current price, the original latency is replayed, errors are never cached), call limiter_for(...) and call_with_retry as models/decisions/openai_decisions.py does if the API is rate-limited, and register an offline stand-in or rely on --mock, which always selects the mock backend. The registry is the extension slot for any decision-only service: its primitives (yes/no probability, choice, score) map onto Predicate and DecisionResult, and a client written against a documented API is all it needs. Package it as a plugin through the ragbench.decision_models entry-point group.

How a run flows

flowchart LR
    D[Documents] --> L[Loaders] --> C[Chunkers]
    C --> S[Systems: indexes, retrieval, agents and tools]
    S --> T[Step trace: latency, tokens, cost]
    Q[Questions and qrels] --> E[Evaluator]
    S --> E
    T --> E
    E --> M[Metrics, judge, statistics]
    M --> R[Recommendation and reports]

The evaluator asks every system the same questions, scores each answer, and writes the files all reports are rebuilt from (ragbench report RUN).

The evaluator is the ragbench/evaluation/ package. evaluator.py orchestrates a run (dataset, checkpoints, spending cap, scheduling through runtime/executor.py); plan.py holds the EvaluationPlan, which systems run on which questions and how many trials of each; question_runner.py answers and scores one question; rows.py turns it into rows; summary.py aggregates them per system; outputs.py writes every file. A feature that changes what a run covers (a subset of systems or questions, a split, repeated trials) belongs in the plan, not in a new evaluator flag.