Federated RAG

Overview

FederatedRAG queries multiple independently-registered RAG backends in parallel, aggregates their answers, caches results, applies circuit-breaking per source, and supports streaming from the best-performing source. This is the natural composition point for the other nine subsystems in this package — any object exposing retrieve/generate (including another AgenticRAG, DomainSpecificRAG, etc.) can be registered as a source.

AggregationMethod (str, Enum)

Values (from federat.py): the class docstring documents four named strategies plus a default — WEIGHTED, VOTING, RANKING, MERGE, and SIMPLE are all referenced as concrete aggregation/streaming code paths (_weighted, _voting, _ranking, _merge, _simple methods exist on FederatedRAG).

CircuitBreaker

@dataclass
class CircuitBreaker:
    failure_threshold: int = ...
    recovery_timeout: float = ...

    @property
    def state(self) -> CircuitState: ...
    def allow(self) -> bool: ...
    def record_success(self) -> None: ...
    def record_failure(self) -> None: ...

State machine: CLOSED → OPEN after failure_threshold consecutive failures; OPEN → HALF_OPEN after recovery_timeout seconds; HALF_OPEN → CLOSED on success or back to OPEN on failure. Each FederatedSource owns its own CircuitBreaker instance, so one failing backend doesn't affect others.

TTLCache (federated_rag.federat_cache)

class TTLCache:
    def __init__(self, ttl: float = 60.0, max_size: int = 256): ...
    def get(self, query: str, context: Optional[Dict]) -> Optional[Any]: ...
    def set(self, query: str, context: Optional[Dict], value: Any) -> None: ...
    def invalidate(self) -> None: ...

Described in its docstring as thread-safe, keyed on (query, context). This is a separate implementation from agentic_rag.cache.TTLCache — the two are not interchangeable and don't share cache entries even if used for the same backend.

StreamChunk

@dataclass
class StreamChunk:
    text: str
    source: str
    done: bool = False
    metadata: Optional[Dict[str, Any]] = None
    error: Optional[str] = None

Yielded by FederatedRAG.stream()/astream(). error is non-None when a source fails mid-stream; done=True marks the final chunk for a source.

FederatedSource

@dataclass
class FederatedSource:
    name: str
    rag_system: Any
    weight: float = 1.0
    enabled: bool = True
    timeout: float = 10.0
    metadata: Optional[Dict] = None
    circuit_breaker: CircuitBreaker = field(default_factory=CircuitBreaker)

Internal registration record created by FederatedRAG.add_source(). supports_streaming() checks for astream → stream → generate_stream on the wrapped backend, in that priority order.

FederatedRAG

Constructor

def __init__(
    self,
    aggregation_method: AggregationMethod | str = AggregationMethod.WEIGHTED,
    min_sources: int = 1,
    cache_ttl: float = 60.0,
    cache_max_size: int = 256,
    stream_race_timeout: float = 2.0,
    max_concurrency: int = 2,
    provider_concurrency: int = 1,
    pre_query_hook: Optional[Callable[[str, Optional[Dict]], None]] = None,
    post_query_hook: Optional[Callable[[Dict[str, Any]], None]] = None,
)

add_source(name, rag_system, weight=1.0, timeout=10.0, metadata=None) -> FederatedRAG

Registers a backend; returns self for chaining:

fed = FederatedRAG().add_source("docs", rag1, weight=1.5).add_source("wiki", rag2, weight=1.0)

Also: remove_source(name), enable_source(name, enabled=True), get_stats() -> Dict[str, Dict] (per-source query count, errors, latency).

query(query, context=None, top_k=5) -> Dict[str, Any]

Synchronous buffered query — safe inside or outside a running event loop, per its docstring (it detects and handles both cases rather than requiring the caller to manage async manually).

query_async(query, context=None, top_k=5) -> Dict[str, Any]

Async buffered query: queries all enabled sources concurrently (respecting max_concurrency/provider_concurrency and each source's circuit breaker via _query_source, which retries on rate-limit errors with exponential backoff), then aggregates via the configured method.

stream(...) / astream(...)

for chunk in fed.stream("Explain RAG"):
    print(chunk.text, end="", flush=True)

async for chunk in fed.astream("Explain RAG"):
    print(chunk.text, end="", flush=True)

Both yield StreamChunk. stream() works even inside environments with an already-running loop (e.g. Jupyter), per its docstring.

Fallback streaming behavior

"If no stream method is found, the federation falls back to calling generate() and re-emits the answer in small word-level fragments."

This means even backends with no native streaming support work with FederatedRAG.stream()/astream() — just without true token-level latency benefits.

_call_generate — the two backend calling conventions

FederatedRAG supports two backend shapes automatically:

• Pattern A (separate retrieve/generate): calls retrieve(query, top_k=top_k) then generate(query) — the only way to correctly honor top_k when generation retrieves internally.
• Pattern B (single generate(query, ...)): used when the backend has no retrieve() method; only passes kwargs the method actually declares (top_k, context, language), inspected via signature reflection — never positionally, to avoid TypeError.

Async priority for generation calls: agenerate > generate (if it's a coroutine function) > generate (sync, run in a thread).

Other methods

Complete example

from zerda.rag.types.federated_rag import FederatedRAG, AggregationMethod

fed = FederatedRAG(
    aggregation_method=AggregationMethod.RANKING,
    min_sources=1,
    max_concurrency=4,
    provider_concurrency=2,
)
fed.add_source("internal_docs", internal_rag, weight=2.0, timeout=8.0)
fed.add_source("public_wiki", wiki_rag, weight=1.0, timeout=15.0)

result = fed.query("What is retrieval-augmented generation?", top_k=5)
print(result)
print(fed.get_stats())

Best practices

• Use RANKING when sources have meaningfully different quality per query (it avoids paying for generation on every source).
• Use MERGE when you want a composite answer citing multiple sources explicitly, and are fine paying the latency cost.
• Tune provider_concurrency separately from max_concurrency if multiple sources share the same underlying LLM provider, to avoid provider-side rate limiting.

Common mistakes

• Registering sources with the same backend class but different provider metadata expectations without setting the "provider" key — _provider_of falls back to the class name, which can accidentally group unrelated sources under one concurrency limit.
• Choosing VOTING and expecting streaming — it always falls back to a single buffered chunk.

Related APIs

• Agentic RAG, Domain-Specific RAG, Self-Improving RAG — all valid rag_system values to register as sources.