RAGSystem and LoadedDocument

core.rag_system.RAGSystem, core.rag_system.LoadedDocument

Overview

RAGSystem is the concrete, production-oriented implementation of BaseRAGSystem. It wires together a caller-supplied vector database, LLM, chunker, and context manager into a full retrieve-then-generate pipeline, adding: query expansion, retrieval caching, optional re-ranking, bilingual prompt routing, source citation, persistence, and a parallel async API (aadd_documents, aretrieve, agenerate, astream).

Why it exists: to provide an out-of-the-box, opinionated orchestration layer so callers don't have to hand-wire query expansion → search → merge → rerank → context → prompt → generate every time.

When to use it: as the default entry point for building a RAG application with core.

When not to use it: if you need retrieval/generation behavior that doesn't fit the retrieve→context→prompt→generate shape (e.g. multi-hop retrieval, agentic tool use) — in that case, subclass BaseRAGSystem directly instead.

Dependencies

RAGSystem does not implement storage, embeddings, or generation itself. It expects four caller-supplied objects with the following duck-typed interfaces (as inferred from every call site in the source):

Chunk objects returned by the chunker/vector_db are expected to expose a text attribute named one of text, content, or page_content, and a doc_id or document_id attribute (used by Reranker._get_text/_get_doc_id and RAGSystem._format_sources).

LoadedDocument

@dataclass
class LoadedDocument:
    page_content: str
    metadata: Dict[str, Any] = field(default_factory=dict)
    doc_id: Optional[str] = None

A single document ready to be ingested via add_documents().

__post_init__

If doc_id is falsy, one is auto-generated:

h = hashlib.md5(self.page_content.encode()).hexdigest()[:8]
self.doc_id = f"doc_{h}_{int(time.time())}"

[!NOTE] Collision behavior The ID includes a coarse (second-resolution) timestamp plus an 8-character MD5 prefix of the content. Two LoadedDocuments with identical page_content created within the same second will receive the same auto-generated doc_id.

Methods

• to_dict() -> Dict[str, Any] — returns {"doc_id", "page_content", "metadata"}.
• __repr__() -> str — LoadedDocument(doc_id='...', preview='first 80 chars...'...) with newlines stripped from the preview.

RAGSystem constructor

Exact signature

def __init__(
    self,
    vector_db: Any,
    llm: Any,
    chunker: Any,
    context_manager: Any,
    config: Optional[RAGConfig] = None,
    prompt: Optional[PromptInput] = None,
    enable_query_expansion: bool = True,
    query_expansion_variants: int = 3,
) -> None

where PromptInput = Union[Any, Callable[[str, str], str], str].

Parameters

Raises

Initialization flow

1. super().__init__(config) — validates config, sets up stats.
2. Checks all four dependencies are non-None; raises RAGInitializationError naming the missing ones if not.
3. Stores vector_db, llm, chunker, context_manager, and _custom_prompt = prompt.
4. Builds self.prompt_router: a PromptRouter(include_few_shot=config.include_few_shot) only if config.enable_prompt_routing is True and prompt is None; otherwise None.
5. Builds self.reranker: None unless config.enable_reranking is True, in which case a RerankerConfig is constructed from the rerank_* fields of RAGConfig and passed (with llm=self.llm) into Reranker(...).
6. Builds self._query_expander: a QueryExpander(llm=self.llm, max_variants=query_expansion_variants, use_llm=True, fallback_synonyms=True) if enable_query_expansion else None.
7. Builds self._retrieval_cache: a RetrievalCache(maxsize=config.cache_maxsize, ttl=config.cache_ttl if config.cache_ttl > 0 else None) if config.enable_retrieval_cache else None.
8. Logs an INFO summary line with the four feature flags.

Thread safety

Construction itself is not thread-safe in any special sense (it's ordinary Python object construction) and is not expected to be called concurrently.

Friendly API — convenience methods

add_text

def add_text(self, text: str, doc_id: Optional[str] = None,
             metadata: Optional[Dict[str, Any]] = None) -> int

Wraps a single string in a LoadedDocument (auto-doc_id = f"text_{int(time.time())}" if omitted — note this differs from LoadedDocument's own MD5-based auto-ID scheme) and calls add_documents([doc]).

Returns: number of chunks created for that one document (0 if empty/failed).

Example:

n = rag.add_text("Employee handbook excerpt...", doc_id="emp_1")

add_texts

def add_texts(self, texts: Dict[str, str],
              metadata: Optional[Dict[str, Dict[str, Any]]] = None) -> Dict[str, int]

Builds one LoadedDocument per {doc_id: text} entry (with per-doc metadata looked up from the metadata dict) and calls add_documents(docs).

Returns: {doc_id: num_chunks}.

ask

def ask(self, question: str, include_sources: bool = False,
        language: Optional[str] = None, **llm_kwargs: Any) -> str

A friendly alias that forwards directly to generate(question, include_sources=include_sources, language=language, **llm_kwargs). No additional logic.

Document ingestion

add_documents

def add_documents(self, docs: List[LoadedDocument]) -> Dict[str, int]

Chunk and index a list of LoadedDocument objects.

Parameters

Returns

Dict[str, int] — {doc_id: num_chunks}. Failed or empty documents are recorded with 0 chunks rather than raising.

Internal workflow / algorithm

1. If docs is empty, log a warning and return {}.
2. For each document:
   • If page_content is falsy/whitespace-only, log a warning, record 0 chunks, and continue (skip).
   • Otherwise pick the chunk function: chunker.chunk_text if it exists (checked with hasattr), else chunker.chunk.
   • Call it as fn(doc.page_content, doc.doc_id).
   • On success: extend the running all_chunks list and record len(chunks).
   • On any Exception: log an error and record 0 chunks for that document (the exception is swallowed, not re-raised).
3. If any chunks were produced overall:
   • Call self.vector_db.add(all_chunks).
   • Update stats["total_documents"] (count of docs with >0 chunks) and stats["total_chunks"].
   • If a retrieval cache is configured, call self._retrieval_cache.clear() — all cached retrieval results are invalidated whenever any new document is ingested, regardless of relevance.
   • If vector_db.add(...) itself raises, the exception is logged and re-raised (unlike chunking failures).
4. Return the per-document results dict.

Raises

Side effects

• Mutates self.stats["total_documents"] and self.stats["total_chunks"].
• Clears self._retrieval_cache entirely on any successful chunk indexing.
• Emits INFO/DEBUG/WARNING/ERROR log records throughout.

Complexity

Time: O(n · c) where n = number of documents and c = cost of chunking each document (chunker-dependent) plus the vector DB's own add() cost. Space: O(total chunks produced).

Examples

Basic:

rag.add_documents([LoadedDocument(page_content="Some text.", doc_id="doc1")])

Multiple documents:

docs = [
    LoadedDocument(page_content="Doc A content...", doc_id="a"),
    LoadedDocument(page_content="Doc B content...", doc_id="b"),
]
results = rag.add_documents(docs)   # {"a": 4, "b": 2}

Handling a document with no content:

rag.add_documents([LoadedDocument(page_content="   ", doc_id="empty")])
# {"empty": 0}  — logged as a warning, not an error

Error handling — a bad vector DB:

class BrokenVectorDB:
    def add(self, chunks): raise ConnectionError("db down")

rag = RAGSystem(vector_db=BrokenVectorDB(), llm=llm, chunker=chunker,
                context_manager=cm)
try:
    rag.add_documents([LoadedDocument(page_content="text", doc_id="x")])
except ConnectionError:
    print("Vector DB write failed — see logs for details")

Production-ready usage — checking partial failures:

results = rag.add_documents(docs)
failed = [doc_id for doc_id, n in results.items() if n == 0]
if failed:
    logger.warning("Documents that produced zero chunks: %s", failed)

Best practices

• Check the returned dict for 0-chunk entries to detect silently skipped/failed documents.
• Remember that ingesting any document clears the entire retrieval cache — batch your ingestion calls rather than calling add_documents once per document in a hot loop if cache warmth matters.

Common mistakes

• Assuming a chunking exception aborts the whole batch — it doesn't; only that document is affected.
• Passing a Dict[str, str] (matching the abstract base's type hint) instead of List[LoadedDocument] — the concrete RAGSystem.add_documents signature expects a list of LoadedDocument, not a dict.

Related APIs

• add_text, add_texts — convenience wrappers.
• aadd_documents — async counterpart (concurrent chunking).
• DocumentVersionManager.add_or_update — diff-aware alternative that avoids re-indexing unchanged content.

Retrieval

retrieve

def retrieve(self, query: str, top_k: Optional[int] = None) -> List[Tuple[Any, float]]

Retrieve relevant chunks using smart query expansion, optional caching, and optional reranking.

Parameters

Returns

List[Tuple[Any, float]] — (chunk, score) pairs sorted by score descending.

Raises

Algorithm

1. Return [] immediately for an empty query (no exception).
2. Cache lookup: if a RetrievalCache is configured, look up (query, k). On hit, return the cached list immediately — no expansion, search, or reranking is performed.
3. Query expansion: if a QueryExpander is configured, call expand(query) to get [query, variant1, ...]; otherwise use [query] as the sole variant.
4. Per-variant search: for each variant, call vector_db.search(variant, top_k=k, score_threshold=config.min_score).
5. Merge: if there is more than one variant, call merge_retrieval_results(per_variant, top_k=k) (dedup by chunk identity, keep highest score); if there was only one variant, use its results as-is (skips the merge step entirely).
6. Rerank: if a Reranker is configured and results are non-empty, call reranker.rerank(query=query, chunks=results).
7. Cache store: if caching is enabled and results are non-empty, store them under (query, k).
8. Return the results.

Thread safety

RetrievalCache itself is internally thread-safe (RLock-protected). The rest of retrieve()'s logic has no explicit locking; safety depends on the thread safety of the supplied vector_db/llm.

Side effects

• Reads from and writes to self._retrieval_cache.
• Logs at DEBUG/INFO/ERROR levels throughout, including truncated query previews (query[:60]).
• Calls into vector_db.search() — may perform network I/O depending on the implementation.

Complexity

Time: O(v · s + m log m) where v = number of query variants, s = cost of one vector_db.search() call, and m = total merged results (for the sort in merge_retrieval_results), plus reranker cost if enabled (see Reranker.rerank complexity).

Examples

Basic:

results = rag.retrieve("What is the leave policy?")
for chunk, score in results:
    print(score, chunk.text[:60])

Overriding top_k:

results = rag.retrieve("refund policy", top_k=10)

Cache hit vs. miss (observable via stats):

rag.retrieve("same question")
rag.retrieve("same question")   # second call is a cache hit
print(rag.get_cache_stats())    # {'hits': 1, 'misses': 1, ...}

Error handling:

from zerda.rag.core.exceptions import RAGRetrievalError

try:
    rag.retrieve("query")
except RAGRetrievalError as e:
    print(f"Retrieval failed: {e.query!r} — {e}")

Production usage — empty results:

results = rag.retrieve(user_query)
if not results:
    return "I couldn't find anything relevant to your question."

Best practices

• Let retrieve() raise RAGRetrievalError when you need to distinguish "search infrastructure failed" from "search succeeded but found nothing" (empty list, no exception).
• Use invalidate_cache() / get_cache_stats() (below) to monitor and manage cache health in long-running processes.

Common mistakes

• Assuming reranking runs even when chunks came back empty — it's skipped (if self.reranker and results:).
• Forgetting that a single-variant path skips merge_retrieval_results entirely — if you rely on merge_retrieval_results's dedup behavior, note it never runs when query expansion is disabled or produces only the original query.

Related APIs

• QueryExpander.expand
• merge_retrieval_results
• Reranker.rerank
• RetrievalCache

Generation

generate

def generate(self, query: str, include_sources: bool = False,
             language: Optional[str] = None, **llm_kwargs: Any) -> str

End-to-end RAG: retrieve → build context → resolve prompt → call LLM.

Parameters

Returns

str — the answer text, or a human-readable error/status string. This method is designed to never raise — all exceptions are caught internally.

Raises

None by design — see algorithm step 9 below.

Algorithm

1. stats["total_queries"] += 1 (always, even for empty/invalid queries).
2. If the query is empty/whitespace, return "⚠️ Please enter a valid question." immediately (no failure counter increment).
3. Call self.retrieve(query).
   • If it returns no chunks: increment failed_queries, log, and return "No relevant information found.".
4. Build context: self.context_manager.build(query, retrieved).
5. Resolve the prompt via self._resolve_prompt(query, context, language) (see below).
6. Call self.llm.generate(prompt, **llm_kwargs).
7. If include_sources, append self._format_sources(retrieved) (see below).
8. Increment successful_queries, log, and return the answer.
9. Exception handling:
   • RAGRetrievalError (from step 3): increment failed_queries, return "❌ Search error occurred.".
   • Any other Exception (from context building, prompt resolution, or LLM call): increment failed_queries, log the error, and return f"❌ There was an error: {exc}".

Thread safety

stats mutation is unsynchronized — see BaseRAGSystem thread safety caveat.

Side effects

• Mutates self.stats.
• Reads/writes the retrieval cache (via retrieve()).
• Calls llm.generate() — typically a network request.
• Extensive logging including truncated query previews.

Complexity

Dominated by retrieve()'s complexity plus the LLM call latency; context building and prompt resolution are effectively O(context length).

Examples

Basic:

answer = rag.generate("What is the return policy?")

With sources:

answer = rag.generate("What is the return policy?", include_sources=True)
print(answer)
# ...answer text...
#
# 📚 Sources:
# • policy_1  (score: 0.87)

Forcing a language:

answer = rag.generate("ما هي سياسة الإرجاع؟", language="ar")

Passing LLM kwargs:

answer = rag.generate("Summarize the policy.", temperature=0.2, max_tokens=200)

Production-ready usage — treating "no results" and "error" as distinct UX states:

answer = rag.generate(user_query)
if answer.startswith("❌"):
    log_and_alert(answer)
elif answer == "No relevant information found.":
    show_no_results_ui()
else:
    show_answer(answer)

Best practices

• Because generate() never raises, check the returned string's prefix/content ("❌", "No relevant information found.") if you need to branch on failure modes programmatically, rather than wrapping the call in try/except.
• Use ask() as a shorter alias when you don't need to think about the method name.

Common mistakes

• Wrapping generate() in a broad try/except expecting it to raise on LLM failures — it won't; failures are returned as strings.
• Forgetting that total_queries increments even for a blank/whitespace query.

Related APIs

• ask — alias.
• _resolve_prompt, _format_sources — internal helpers used by this method.
• agenerate — async counterpart.
• astream — token-streaming counterpart.

Prompt management

set_prompt

def set_prompt(self, prompt: Optional[PromptInput]) -> None

Replace or remove the custom prompt after construction.

• Passing a non-None value sets self._custom_prompt = prompt and disables self.prompt_router (set to None), regardless of config.enable_prompt_routing.
• Passing None: - If config.enable_prompt_routing is True, rebuilds self.prompt_router = PromptRouter(include_few_shot=config.include_few_shot). - Otherwise self.prompt_router stays None.

Example:

from dataclasses import dataclass

@dataclass
class SimpleTemplate:
    template: str
    def format(self, context, question):
        return self.template.format(context=context, question=question)

rag.set_prompt(SimpleTemplate("Context:\n{context}\n\nQ: {question}\nA:"))
# ... later ...
rag.set_prompt(None)   # restore PromptRouter / built-in default

_resolve_prompt (internal)

def _resolve_prompt(self, query: str, context: str, language: Optional[str] = None) -> str

Priority order, exactly as implemented:

1. Custom prompt (self._custom_prompt is not None):
   • If it has a .format attribute, try cp.format(context=context, question=query); on KeyError/TypeError, retry with cp.format(context=context, query=query) (supports templates that name the parameter query instead of question).
   • Else if it's callable, call cp(context, query).
   • Else if it's a str, call cp.format(context=context, question=query).
   • Else, log a warning and fall through to the router/default path below.
2. PromptRouter (if configured): calls prompt_router.build(query=query, context=context, language=language). Also separately calls detect_language/detect_question_type purely for a debug log line (these results are not reused — build() redoes its own detection internally).
3. Built-in fallback (used only when there is no custom prompt and no PromptRouter, i.e. enable_prompt_routing=False and prompt=None): detects language via a private helper and calls self._build_prompt(...).

_detect_language_simple (internal)

def _detect_language_simple(self, text: str) -> str

Returns "ar" if any character in text falls in the Unicode range \u0600–\u06FF, else "en". This is a simpler check than PromptRouter.detect_language's ratio-based approach (see PromptRouter) — it only checks for any Arabic character, not a proportion.

_build_prompt (internal)

def _build_prompt(self, query: str, context: str, language_pr: str) -> str

Hard-coded Arabic/English instruction templates (distinct from PromptRouter's templates — this is the last-resort fallback used only when prompt routing is disabled entirely).

Internal helpers

_format_sources

def _format_sources(self, chunks: List[Tuple[Any, float]], max_sources: int = 3) -> str

Builds a "📚 Sources:\n• {doc_id} (score: {score:.2f})" block, deduplicated by doc_id (or document_id, or "unknown" if neither attribute exists), capped at max_sources (hardcoded to 3 at the call site in generate() — the parameter is not currently exposed to callers of generate()). Returns "" if no lines were produced.

Persistence

remove_document

def remove_document(self, doc_id: str) -> int

Calls self.vector_db.remove_by_doc_id(doc_id). On success (removed > 0), decrements stats["total_documents"] by 1 and stats["total_chunks"] by removed (both floored at 0 via max(0, ...)). On any exception, logs an error and returns 0 rather than propagating.

Returns: number of chunks removed (0 if not found or on error).

save

def save(self, path: str) -> None

1. Creates the directory path (and parents) if needed.
2. Calls self.vector_db.save(Path(path) / "vector_db").
3. Writes self.stats as pretty-printed JSON to Path(path) / "stats.json".

Raises

load (classmethod)

@classmethod
def load(cls, path: str, vector_db: Any, llm: Any, chunker: Any,
         context_manager: Any, config: Optional[RAGConfig] = None,
         enable_query_expansion: bool = True) -> "RAGSystem"

1. Raises FileNotFoundError if path doesn't exist.
2. Calls vector_db.load(Path(path) / "vector_db") — wraps any failure in VectorDBPersistenceError("load", ...).
3. Constructs a new RAGSystem with the now-populated vector_db and the other supplied components.
4. If stats.json exists under path, updates the new instance's stats dict with its contents.

[!WARNING] You must supply live components to load() load() does not reconstruct llm, chunker, or context_manager from disk — only the vector database's own persisted state (via vector_db.load()) and the stats counters are restored. You must pass working llm/chunker/context_manager instances just as you would to the constructor.

invalidate_cache

def invalidate_cache(self) -> None

Manually clears the retrieval cache (no-op if caching is disabled). Useful when documents are updated by a process other than add_documents/remove_document (which already clear/adjust the cache themselves).

get_cache_stats

def get_cache_stats(self) -> Dict[str, Any]

Returns {} if caching is disabled, otherwise self._retrieval_cache.stats() — see RetrievalCache.stats.

get_stats (override)

def get_stats(self) -> Dict[str, Any]

Extends BaseRAGSystem.get_stats() with:

{
  "vector_db_size":  getattr(self.vector_db, "size", 0),
  "llm_type":        type(self.llm).__name__,
  "chunker_type":    type(self.chunker).__name__,
  "query_expansion": self._query_expander is not None,
  "reranking":       self.reranker is not None,
  "cache":           self.get_cache_stats(),
}

__repr__

RAGSystem(docs={n}, chunks={n}, queries={n}, expansion={bool})

Context manager

cleanup (override)

Iterates ("vector_db", "llm", "chunker", "context_manager") and calls .cleanup() on each component that defines it, catching and logging (not propagating) any exception per-component. Logs RAGSystem cleanup complete when done.

__enter__ / __exit__

Synchronous context-manager support; __exit__ calls cleanup() and returns False (exceptions propagate normally).

Async API

RAGSystem overrides all of BaseRAGSystem's thread-wrapped async defaults with native asyncio implementations.

async aadd_documents

async def aadd_documents(self, docs: Dict[str, str]) -> Dict[str, int]

Chunks all documents concurrently via asyncio.gather over per-document coroutines (each of which runs the sync chunk function in a thread via asyncio.to_thread). After gathering, calls vector_db.aadd(all_chunks) if the vector DB defines it, else falls back to asyncio.to_thread(vector_db.add, all_chunks).

[!WARNING] Parameter naming mismatch Despite the abstract base and this override both typing docs as Dict[str, str], RAGSystem's synchronous add_documents expects List[LoadedDocument]. aadd_documents iterates docs.items(), so — unlike the sync method — it genuinely expects a Dict[str, str] ({doc_id: text}) as documented in its own docstring. The sync and async ingestion methods on the same class accept different input shapes.

async aretrieve

async def aretrieve(self, query: str, top_k: Optional[int] = None) -> List[Tuple]

Mirrors retrieve()'s algorithm (expand → search each variant concurrently via asyncio.gather → merge → rerank) but:

• Does not consult or populate the retrieval cache at all (no cache interaction in the async path).
• Uses vector_db.asearch(...) if defined, else asyncio.to_thread(vector_db.search, ...).
• Uses asyncio.to_thread for the (synchronous) Reranker.rerank call.

async agenerate

async def agenerate(self, query: str, include_sources: bool = False,
                    language: Optional[str] = None, **llm_kwargs: Any) -> str

Mirrors generate(): calls aretrieve, builds context and resolves the prompt via asyncio.to_thread, then calls llm.generate_async(...) if defined, else asyncio.to_thread(llm.generate, ...). Also never raises — same never-raises design as generate(), catching all exceptions and returning an error string.

async astream

async def astream(self, query: str, language: Optional[str] = None, **llm_kwargs: Any)

An async generator that yields answer tokens incrementally.

Algorithm:

1. aretrieve(query); if empty, yield "No relevant information found." and return.
2. Build context and resolve the prompt (both via asyncio.to_thread).
3. Call self.llm.astream(prompt, **llm_kwargs) and branch on what it returns:
   • Async generator → iterate and yield each token directly.
   • Coroutine → await it; if the awaited result is itself an async generator, iterate it; if it's a str, yield it word-by-word with a trailing space (asyncio.sleep(0) between words to yield control); otherwise yield str(actual).
   • Sync iterable (has __iter__) → iterate and yield str(token) with asyncio.sleep(0) between tokens.
   • Anything else → yield str(result) as a single chunk.

Example:

async for token in rag.astream("Explain the refund policy."):
    print(token, end="", flush=True)

async __aenter__ / async __aexit__

Async context-manager support; __aexit__ calls the synchronous cleanup() and returns False.