rag.rag 

async aindex(text, source=None, checksum=None, **metadata)[source]

Index a document into the vector store and/or BM25 index.

Chunks the text, computes embeddings, and stores them. Skips indexing if checksum matches the previously stored value.

Parameters:

text (str) – Document text to index.
source (str | None) – Source identifier for deduplication and filtering.
checksum (str | None) – If provided with source, skip indexing when the checksum matches the previously stored value.
**metadata (str) – Additional metadata key-value pairs attached to every chunk.

Returns:

Number of chunks created (0 if text is empty or checksum matches).

Return type:

index_many(documents)[source]

Index multiple documents sequentially (sync wrapper).

Parameters:: documents (list[dict[str, Any]]) – List of dicts with keys matching aindex parameters (text, source, checksum, etc.).
Returns:: Total number of chunks created across all documents.
Return type:: int

async aindex_many(documents)[source]

Index multiple documents sequentially.

Parameters:: documents (list[dict[str, Any]]) – List of dicts with keys matching aindex parameters (text, source, checksum, etc.).
Returns:: Total number of chunks created across all documents.
Return type:: int

chunk(text, **metadata)[source]

Split text into chunks without indexing.

Parameters:

text (str) – Text to chunk.
**metadata (str) – Metadata attached to each chunk.

Returns:

List of TextChunk instances.

Return type:

list[TextChunk]

search(query, top_k=5, **filters)[source]

Search for relevant chunks (sync wrapper).

Safe to call from within a running event loop (e.g. Jupyter).

Parameters:

query (str) – Search query string.
top_k (int) – Maximum number of results to return.
**filters (str) – Metadata key-value filters passed to the store.

Returns:

Ranked list of search hits.

Return type:

async asearch(query, top_k=5, **filters)[source]

Search for relevant chunks using vector, BM25, or hybrid search.

Applies query expansion and reranking if configured.

Parameters:

query (str) – Search query string.
top_k (int) – Maximum number of results to return.
**filters (str) – Metadata key-value filters passed to the store.

Returns:

Ranked list of search hits.

Return type:

query(question, generate_fn=None, top_k=5, prompt_template=None, max_context_chars=None, allow_llm_on_empty=True, generate_timeout=None, **filters)[source]

Retrieve context and generate an answer (sync wrapper).

Safe to call from within a running event loop (e.g. Jupyter).

Parameters:

question (str)
generate_fn (Callable[[str], str | GenerationResult] | None)
top_k (int)
prompt_template (str | None)
max_context_chars (int | None)
allow_llm_on_empty (bool)
generate_timeout (float | None)
filters (str)

Return type:

async aquery(question, generate_fn=None, top_k=5, prompt_template=None, max_context_chars=None, allow_llm_on_empty=True, generate_timeout=None, **filters)[source]

Retrieve context and generate an answer.

Searches for relevant chunks via asearch.
Formats hits into a context string (format_hits).
Fills the prompt template with {context} and {question}.
Calls generate_fn in a thread (asyncio.to_thread).

Parameters:

question (str) – The user question.
generate_fn (Callable[[str], str | GenerationResult] | None) – A sync callable that takes the formatted prompt and returns an answer string or GenerationResult. When None, uses the default set on the RAG instance.
top_k (int) – Number of search results to retrieve.
prompt_template (str | None) – Must contain {context} and {question} placeholders. Unknown placeholders resolve to empty string. Defaults to DEFAULT_RAG_PROMPT.
max_context_chars (int | None) – Truncates the formatted context to this many characters. Note that each hit includes a header line (source, relevance) that counts toward the budget, so small values may exclude all content. None means no limit.
allow_llm_on_empty (bool) – When False and no search hits are found, skip the LLM call and return an empty QueryResult. Defaults to True (backward compat).
generate_timeout (float | None) – Maximum seconds to wait for generate_fn to complete. Raises TimeoutError if exceeded. The underlying LLM request continues running in its thread (cannot be cancelled), so the API cost may still be incurred. Defaults to None (no timeout).
**filters (str) – Passed through to the vector store where clause.

Returns:

A QueryResult with the answer, search hits, deduplicated sources, the full prompt sent to generate_fn, and generation metadata (usage, cost, duration).

Return type:

delete(source)[source]

Delete all chunks associated with a source from the store and BM25 index.

Parameters:: source (str) – Source identifier to delete.
Return type:: None

count()[source]

Return the total number of indexed chunks.

Checks the vector store first, then falls back to BM25.

Returns:: Chunk count, or 0 if neither store is configured.
Return type:: int

rebuild_bm25_from_store()[source]

Rebuild the BM25 index from the persistent vector store.

Clears the current BM25 index and re-adds every chunk the store exposes via get_all(). Chunks whose chunking_strategy metadata is set and differs from the current chunker are skipped (legacy chunks with no strategy tag are included for backward compatibility).

Returns:: Number of documents added to the BM25 index.
Raises:: RuntimeError – if BM25 or the store is not configured.
Return type:: int

Classes

class RAG(embed='mistral/mistral-embed', store=None, chunker='recursive', chunk_size=1000, chunk_overlap=200, bm25_alpha=None, bm25_autorebuild=True, reranker=None, query_expander=None, generate_fn=None)[source]

Bases: object

End-to-end retrieval-augmented generation pipeline.

Combines embedding, chunking, vector/BM25 search, reranking, query expansion, and LLM generation into a single interface.

Parameters:

embed (Embeddings | str) – Embedding model instance or model name string (e.g. "mistral/mistral-embed").
store (VectorStore | None) – Vector store for persistent embeddings. If None, only BM25 search is available (when bm25_alpha is set).
chunker (str) – Name of the chunking strategy ("recursive", "character", "markdown", "code", "hierarchical").
chunk_size (int) – Target chunk size in characters.
chunk_overlap (int) – Overlap between consecutive chunks in characters.
bm25_alpha (float | None) – If set, enables BM25 alongside vector search with this weight for the vector component (0–1). None disables BM25.
reranker (str | None) – Reranker strategy name ("cross_encoder", "diversity", "noop"). None disables reranking.
query_expander (Callable[[str], list[str]] | None) – Callable that takes a query string and returns a list of expanded query strings.
generate_fn (Callable[[str], str | GenerationResult] | None) – Default generation function for query() and aquery(). Takes a prompt string, returns an answer string or GenerationResult.
bm25_autorebuild (bool)

Example

Using litellm for generation:

from ffai.rag import RAG, litellm_generate_fn

rag = RAG(embed="mistral/mistral-embed")
fn = litellm_generate_fn(
    model="mistral/mistral-small-latest", api_key=key,
)
result = rag.query("What is Python?", generate_fn=fn)
print(result.usage.input_tokens)
print(result.cost_usd)

classmethod from_config(*, bm25_only=False, api_key=None, **overrides)[source]

Create a RAG instance from the project configuration file.

Reads all settings from config/main.yaml under the rag: key. The embedding model’s API key is resolved in order: the api_key parameter, the provider-specific environment variable (e.g. MISTRAL_API_KEY), and finally None (which will raise at embed time if no key is found).

Parameters:

bm25_only (bool) – If True, skip vector store creation and use BM25 search only.
api_key (str | None) – API key for the embedding model provider. When None, the key is read from a provider-specific environment variable (e.g. MISTRAL_API_KEY).
**overrides (Any) – Override any constructor parameter from config.

Returns:

Configured RAG instance.

Return type:

RAG

set_generate_fn(generate_fn)[source]

Set or replace the default generation function.

Also re-wires the query expander if one is configured.

Parameters:: generate_fn (Callable[[str], str | GenerationResult]) – Sync callable that takes a prompt and returns an answer string or GenerationResult.
Return type:: None

index(text, source=None, checksum=None, **metadata)[source]

Index a document (sync wrapper).

Safe to call from within a running event loop (e.g. Jupyter).

Parameters:

text (str) – Document text to index.
source (str | None) – Source identifier for deduplication and filtering.
checksum (str | None) – If provided with source, skip indexing when the checksum matches the previously stored value.
**metadata (str) – Additional metadata key-value pairs attached to every chunk.

Returns:

Number of chunks created.

Return type:

async aindex(text, source=None, checksum=None, **metadata)[source]

Index a document into the vector store and/or BM25 index.

Chunks the text, computes embeddings, and stores them. Skips indexing if checksum matches the previously stored value.

Parameters:

text (str) – Document text to index.
source (str | None) – Source identifier for deduplication and filtering.
checksum (str | None) – If provided with source, skip indexing when the checksum matches the previously stored value.
**metadata (str) – Additional metadata key-value pairs attached to every chunk.

Returns:

Number of chunks created (0 if text is empty or checksum matches).

Return type:

index_many(documents)[source]

Index multiple documents sequentially (sync wrapper).

Parameters:: documents (list[dict[str, Any]]) – List of dicts with keys matching aindex parameters (text, source, checksum, etc.).
Returns:: Total number of chunks created across all documents.
Return type:: int

async aindex_many(documents)[source]

Index multiple documents sequentially.

Parameters:: documents (list[dict[str, Any]]) – List of dicts with keys matching aindex parameters (text, source, checksum, etc.).
Returns:: Total number of chunks created across all documents.
Return type:: int

chunk(text, **metadata)[source]

Split text into chunks without indexing.

Parameters:

text (str) – Text to chunk.
**metadata (str) – Metadata attached to each chunk.

Returns:

List of TextChunk instances.

Return type:

list[TextChunk]

search(query, top_k=5, **filters)[source]

Search for relevant chunks (sync wrapper).

Safe to call from within a running event loop (e.g. Jupyter).

Parameters:

query (str) – Search query string.
top_k (int) – Maximum number of results to return.
**filters (str) – Metadata key-value filters passed to the store.

Returns:

Ranked list of search hits.

Return type:

async asearch(query, top_k=5, **filters)[source]

Search for relevant chunks using vector, BM25, or hybrid search.

Applies query expansion and reranking if configured.

Parameters:

query (str) – Search query string.
top_k (int) – Maximum number of results to return.
**filters (str) – Metadata key-value filters passed to the store.

Returns:

Ranked list of search hits.

Return type:

query(question, generate_fn=None, top_k=5, prompt_template=None, max_context_chars=None, allow_llm_on_empty=True, generate_timeout=None, **filters)[source]

Retrieve context and generate an answer (sync wrapper).

Safe to call from within a running event loop (e.g. Jupyter).

Parameters:

question (str)
generate_fn (Callable[[str], str | GenerationResult] | None)
top_k (int)
prompt_template (str | None)
max_context_chars (int | None)
allow_llm_on_empty (bool)
generate_timeout (float | None)
filters (str)

Return type:

async aquery(question, generate_fn=None, top_k=5, prompt_template=None, max_context_chars=None, allow_llm_on_empty=True, generate_timeout=None, **filters)[source]

Retrieve context and generate an answer.

Searches for relevant chunks via asearch.
Formats hits into a context string (format_hits).
Fills the prompt template with {context} and {question}.
Calls generate_fn in a thread (asyncio.to_thread).

Parameters:

question (str) – The user question.
generate_fn (Callable[[str], str | GenerationResult] | None) – A sync callable that takes the formatted prompt and returns an answer string or GenerationResult. When None, uses the default set on the RAG instance.
top_k (int) – Number of search results to retrieve.
prompt_template (str | None) – Must contain {context} and {question} placeholders. Unknown placeholders resolve to empty string. Defaults to DEFAULT_RAG_PROMPT.
max_context_chars (int | None) – Truncates the formatted context to this many characters. Note that each hit includes a header line (source, relevance) that counts toward the budget, so small values may exclude all content. None means no limit.
allow_llm_on_empty (bool) – When False and no search hits are found, skip the LLM call and return an empty QueryResult. Defaults to True (backward compat).
generate_timeout (float | None) – Maximum seconds to wait for generate_fn to complete. Raises TimeoutError if exceeded. The underlying LLM request continues running in its thread (cannot be cancelled), so the API cost may still be incurred. Defaults to None (no timeout).
**filters (str) – Passed through to the vector store where clause.

Returns:

A QueryResult with the answer, search hits, deduplicated sources, the full prompt sent to generate_fn, and generation metadata (usage, cost, duration).

Return type: