rag

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:

int

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:

int

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:

list[SearchHit]

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:

list[SearchHit]

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:
Return type:

QueryResult

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.

  1. Searches for relevant chunks via asearch.

  2. Formats hits into a context string (format_hits).

  3. Fills the prompt template with {context} and {question}.

  4. 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:

QueryResult

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

class BM25Index(k1=1.5, b=0.75, epsilon=0.25)[source]

Bases: object

BM25 sparse index for keyword-based retrieval.

Implements the Okapi BM25 algorithm for text relevance scoring. Used alongside vector search for hybrid retrieval.

Parameters:

  • k1 (float) – BM25 k1 parameter (term frequency saturation).

  • b (float) – BM25 b parameter (document length normalization).

  • epsilon (float) – Small value to avoid division by zero.

tokenize(text)[source]

Tokenize text into terms.

Parameters:

text (str) – Text to tokenize.

Returns:

List of lowercase tokens.

Return type:

list[str]

add_document(doc_id, content, metadata=None)[source]

Add a document to the index.

Parameters:

  • doc_id (str) – Unique document identifier.

  • content (str) – Document text content.

  • metadata (dict[str, Any] | None) – Optional metadata dictionary.

Return type:

None

add_documents(documents, id_key='id', content_key='content')[source]

Add multiple documents to the index.

Parameters:

  • documents (list[dict[str, Any]]) – List of document dictionaries.

  • id_key (str) – Key for document ID in each dict.

  • content_key (str) – Key for document content in each dict.

Returns:

Number of documents added.

Return type:

int

delete_document(doc_id)[source]

Delete a document from the index.

Parameters:

doc_id (str) – Document ID to delete.

Returns:

True if document was deleted, False if not found.

Return type:

bool

search(query, n_results=5)[source]

Search for documents using BM25 scoring.

Parameters:

  • query (str) – Search query.

  • n_results (int) – Maximum number of results.

Returns:

List of result dicts with doc_id, score, content, metadata.

Return type:

list[dict[str, Any]]

clear()[source]

Clear all documents from the index.

Return type:

None

delete_by_metadata(key, value)[source]

Delete all documents matching a metadata key-value pair.

Parameters:

  • key (str) – Metadata key to match.

  • value (Any) – Metadata value to match.

Returns:

Number of documents deleted.

Return type:

int

count()[source]

Get the number of documents in the index.

Return type:

int

get_stats()[source]

Get index statistics.

Return type:

dict[str, Any]

class CharacterChunker(chunk_size=1000, chunk_overlap=200, metadata=None, respect_word_boundaries=True)[source]

Bases: ChunkerBase

Character-based chunking with word-boundary awareness.

Parameters:

  • chunk_size (int) – Maximum characters per chunk.

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • respect_word_boundaries (bool) – Whether to prefer word boundaries for splits.

chunk(text, metadata=None)[source]

Split text into overlapping chunks.

Parameters:

  • text (str) – The text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

class ChunkDeduplicator(mode='exact', similarity_threshold=0.95)[source]

Bases: object

Detect and filter duplicate/near-duplicate chunks.

Supports multiple deduplication modes: - exact: Hash-based exact content matching - similarity: Embedding cosine similarity threshold

Parameters:

  • mode (str) – Deduplication mode (“exact” or “similarity”).

  • similarity_threshold (float) – Threshold for similarity-based dedup (0.0-1.0).

Example

>>> dedup = ChunkDeduplicator(mode="exact")
>>> chunks, embeddings = dedup.filter_duplicates(chunks, embeddings)
compute_hash(content)[source]

Compute content hash for exact deduplication.

Parameters:

content (str) – Chunk text content.

Returns:

SHA256 hash (first 16 chars).

Return type:

str

is_duplicate(content, embedding=None)[source]

Check if chunk is a duplicate of previously seen content.

Parameters:

  • content (str) – Chunk text content.

  • embedding (list[float] | None) – Optional embedding vector for similarity-based dedup.

Returns:

True if chunk is a duplicate.

Return type:

bool

filter_duplicates(chunks, embeddings)[source]

Filter duplicate chunks from a batch.

Parameters:

  • chunks (list[Any]) – List of chunk objects (must have .content attribute or be string).

  • embeddings (list[list[float]]) – List of embedding vectors.

Returns:

Tuple of (filtered_chunks, filtered_embeddings).

Return type:

tuple[list[Any], list[list[float]]]

clear()[source]

Clear seen hashes and embeddings.

Return type:

None

get_stats()[source]

Get deduplicator statistics.

Returns:

Dict with mode, threshold, and counts.

Return type:

dict[str, Any]

class ChunkerBase(chunk_size=1000, chunk_overlap=200, metadata=None)[source]

Bases: ABC

Abstract base class for text chunking strategies.

Parameters:

  • chunk_size (int) – Maximum size of each chunk (interpretation varies by strategy).

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata to attach to all chunks.

abstractmethod chunk(text, metadata=None)[source]

Split text into chunks.

Parameters:

  • text (str) – The text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk (merged with default).

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

property name: str

Return the chunker strategy name.

class ClientAdapter(client, **kwargs)[source]

Bases: object

Wrap an FFAI client into a callable that returns a GenerationResult.

Handles both synchronous and asynchronous generate_response methods, automatically awaiting coroutines when necessary.

Parameters:

  • client (Any) – An FFAI client instance with a generate_response method.

  • **kwargs (Any) – Extra keyword arguments forwarded to generate_response on every call.

__call__(prompt)[source]

Generate a response for the given prompt.

Parameters:

prompt (str) – The prompt string to send to the client.

Returns:

A GenerationResult containing the response text, usage, cost, and duration metadata.

Return type:

GenerationResult

class CodeChunker(chunk_size=1000, chunk_overlap=200, metadata=None, language='python', split_by='function')[source]

Bases: ChunkerBase

Code-aware chunking that splits by functions/classes.

Parameters:

  • chunk_size (int) – Maximum characters per chunk.

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • language (str) – Programming language for parsing hints.

  • split_by (str) – Strategy: “function”, “class”, or “module”.

LANGUAGE_PATTERNS: dict[str, dict[str, str]] = {'generic': {'class': '^(class|struct|interface)\\s+\\w+', 'comment': '^(#|//|/\\*)', 'function': '^(function|func|def|fn)\\s+\\w+', 'import': '^(import|use|require)'}, 'go': {'class': '^type\\s+\\w+\\s+struct', 'comment': '^(//|/\\*)', 'function': '^func\\s+(\\(\\w+\\s+\\*?\\w+\\)\\s+)?\\w+\\s*\\(', 'import': '^import\\s+'}, 'java': {'class': '^\\s*(public\\s+)?(abstract\\s+)?class\\s+\\w+|interface\\s+\\w+', 'comment': '^(//|/\\*)', 'function': '^\\s*(public|private|protected|static)?\\s*\\w+\\s+\\w+\\s*\\(', 'import': '^import\\s+'}, 'javascript': {'class': '^(export\\s+)?class\\s+\\w+', 'comment': '^(//|/\\*|\\*)', 'function': '^(async\\s+)?function\\s+\\w+|const\\s+\\w+\\s*=\\s*(async\\s+)?\\([^)]*\\)\\s*=>|export\\s+(async\\s+)?function', 'import': '^(import\\s+|export\\s+|require\\s*\\()'}, 'python': {'class': '^class\\s+\\w+[\\(:]', 'comment': '^(#|\'\'\'|\\"\\"\\")', 'decorator': '^@\\w+', 'function': '^(async\\s+)?def\\s+\\w+\\s*\\(', 'import': '^(import\\s+|from\\s+\\S+\\s+import)'}, 'rust': {'class': '^(pub\\s+)?struct\\s+\\w+|impl\\s+\\w+', 'comment': '^(//|/\\*|\\*)', 'function': '^(pub\\s+)?(async\\s+)?fn\\s+\\w+', 'import': '^use\\s+'}, 'typescript': {'class': '^(export\\s+)?(abstract\\s+)?class\\s+\\w+|interface\\s+\\w+|type\\s+\\w+', 'comment': '^(//|/\\*|\\*)', 'function': '^(async\\s+)?(function\\s+\\w+|const\\s+\\w+\\s*=\\s*(async\\s+)?\\([^)]*\\)\\s*(:\\s*\\w+)?\\s*=>|export\\s+(async\\s+)?function)', 'import': '^(import\\s+|export\\s+|require\\s*\\()'}}
chunk(text, metadata=None)[source]

Split code by structural boundaries.

Parameters:

  • text (str) – The code text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

class ContextualEmbeddings(context_template=None, max_context_length=200)[source]

Bases: object

Generate embeddings with document context prepended.

Prepends document context (title, summary, or preceding content) to each chunk before embedding, improving semantic understanding.

Parameters:

  • context_prefix – Template for context prefix.

  • max_context_length (int) – Maximum characters for context prefix.

  • context_template (str | None)

DEFAULT_CONTEXT_TEMPLATE = 'Document: {title}\n\nSection: {section}\n\n{chunk}'
prepare_chunk_for_embedding(chunk_content, document_title=None, section_header=None, document_summary=None, preceding_context=None)[source]

Prepare a chunk with context for embedding.

Parameters:

  • chunk_content (str) – The chunk text content.

  • document_title (str | None) – Document title or name.

  • section_header (str | None) – Section header if available.

  • document_summary (str | None) – Brief document summary.

  • preceding_context (str | None) – Text immediately preceding this chunk.

Returns:

Context-enhanced text for embedding.

Return type:

str

prepare_chunks_batch(chunks, document_title=None, document_summary=None)[source]

Prepare multiple chunks with context for embedding.

Parameters:

  • chunks (list[dict[str, Any]]) – List of chunk dictionaries with ‘content’ and optional metadata.

  • document_title (str | None) – Document title for all chunks.

  • document_summary (str | None) – Document summary (currently unused but available).

Returns:

List of context-enhanced texts for embedding.

Return type:

list[str]

class CrossEncoderReranker(model_name='cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512, fastembed_model_name=None)[source]

Bases: RerankerBase

Re-ranker using cross-encoder models.

Tries sentence-transformers first, then falls back to fastembed (ONNX-based). Requires at least one of:

pip install sentence-transformers pip install fastembed

Parameters:

  • model_name (str) – Cross-encoder model name.

  • max_length (int) – Maximum sequence length (sentence-transformers only).

  • fastembed_model_name (str | None) – Model name for fastembed fallback.

rerank(query, results, n_results=None)[source]

Re-rank results using cross-encoder scoring.

Parameters:

  • query (str) – Original search query.

  • results (list[dict[str, Any]]) – Search results to re-rank.

  • n_results (int | None) – Number of results to return.

Returns:

Re-ranked results with updated scores.

Return type:

list[dict[str, Any]]

class DiversityReranker(lambda_param=0.7)[source]

Bases: RerankerBase

Re-ranker that promotes result diversity.

Re-orders results to maximize diversity based on content similarity. Uses MMR (Maximal Marginal Relevance) style selection.

Parameters:

lambda_param (float) – Balance between relevance and diversity (0-1). Higher = more relevance, lower = more diversity.

rerank(query, results, n_results=None)[source]

Re-rank results for diversity.

Parameters:

  • query (str) – Original search query (not used, kept for interface).

  • results (list[dict[str, Any]]) – Search results to re-rank.

  • n_results (int | None) – Number of results to return.

Returns:

Diversified results.

Return type:

list[dict[str, Any]]

class Embeddings(model='mistral/mistral-embed', api_key=None, api_base=None, cache_enabled=True, cache_size=256, device='cpu', **kwargs)[source]

Bases: object

Compute text embeddings using local or remote models.

Supports local sentence-transformer models (prefix local/) and remote providers via LiteLLM (e.g. mistral/mistral-embed, openai/text-embedding-3-small). Results are optionally cached in an LRU-style in-memory cache.

Parameters:

  • model (str) – Model identifier. Use local/<name> for local sentence-transformer models, or <provider>/<model> for remote APIs.

  • api_key (str | None) – API key for the remote provider. If None, the key is read from a provider-specific environment variable.

  • api_base (str | None) – Optional custom API base URL.

  • cache_enabled (bool) – Whether to cache embedding results in memory.

  • cache_size (int) – Maximum number of entries in the embedding cache.

  • device (str) – Torch device for local models (e.g. "cpu", "cuda").

  • **kwargs (Any) – Extra keyword arguments forwarded to the LiteLLM embedding call.

async aembed(texts)[source]

Compute embeddings for one or more texts asynchronously.

Parameters:

texts (str | list[str]) – A single text string or a list of text strings.

Returns:

List of embedding vectors, one per input text.

Return type:

list[list[float]]

async aembed_single(text)[source]

Compute the embedding for a single text asynchronously.

Parameters:

text (str) – The text to embed.

Returns:

The embedding vector.

Return type:

list[float]

embed(texts)[source]

Compute embeddings for one or more texts synchronously.

Parameters:

texts (str | list[str]) – A single text string or a list of text strings.

Returns:

List of embedding vectors, one per input text.

Return type:

list[list[float]]

embed_single(text)[source]

Compute the embedding for a single text synchronously.

Parameters:

text (str) – The text to embed.

Returns:

The embedding vector.

Return type:

list[float]

clear_cache()[source]

Clear the embedding cache.

Returns:

Number of entries that were in the cache before clearing.

Return type:

int

cache_stats()[source]

Return statistics about the embedding cache.

Returns:

Dictionary with keys enabled, max_size, and entries.

Return type:

dict[str, Any]

static cosine_similarity(a, b)[source]

Compute cosine similarity between two vectors.

Parameters:
Returns:

Cosine similarity score in the range [-1, 1]. Returns 0.0 if either vector has zero magnitude.

Return type:

float

property provider: str
property is_local: bool
class GenerationResult(text, usage=None, cost_usd=0.0, duration_ms=None)[source]

Bases: object

Result of a RAG generation call.

Return this from generate_fn to preserve token usage, cost, and timing metadata. Returning a plain string silently discards these metrics and emits a warning.

Parameters:
text

Generated answer text.

Type:

str

usage

Provider-specific token usage object (e.g. TokenUsage).

Type:

Any | None

cost_usd

Estimated cost in USD.

Type:

float

duration_ms

Wall-clock generation duration in milliseconds.

Type:

float | None

text: str
usage: Any | None = None
cost_usd: float = 0.0
duration_ms: float | None = None
class HierarchicalChunker(chunk_size=400, chunk_overlap=100, metadata=None, parent_chunk_size=1500, max_levels=2)[source]

Bases: ChunkerBase

Hierarchical chunking with parent-child relationships.

Parameters:

  • chunk_size (int) – Maximum characters per leaf chunk.

  • chunk_overlap (int) – Overlap between leaf chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • parent_chunk_size (int) – Size of parent chunks (larger).

  • max_levels (int) – Maximum hierarchy depth.

chunk(text, metadata=None)[source]

Split text into hierarchical chunks with parent-child relationships.

Parameters:

  • text (str) – The text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of HierarchicalTextChunk objects with parent-child links.

Return type:

list[HierarchicalTextChunk]

get_parent_chunks(chunks)[source]

Filter to return only parent chunks.

Parameters:

chunks (list[HierarchicalTextChunk])

Return type:

list[HierarchicalTextChunk]

get_child_chunks(chunks)[source]

Filter to return only child chunks.

Parameters:

chunks (list[HierarchicalTextChunk])

Return type:

list[HierarchicalTextChunk]

get_chunks_with_parent_context(child_chunks, all_chunks)[source]

Get child chunks with their parent context for retrieval.

Parameters:
Returns:

List of dicts with child content and parent context.

Return type:

list[dict[str, Any]]

class HierarchicalIndex(include_parent_context=True)[source]

Bases: object

Index for hierarchical chunk storage and retrieval.

Stores both parent and child chunks, enabling: - Fine-grained search on child chunks - Parent context retrieval for matched children

Parameters:

include_parent_context (bool) – Whether to include parent content in results.

add_chunk(chunk_id, content, embedding=None, parent_id=None, hierarchy_level=0, metadata=None)[source]

Add a chunk to the hierarchical index.

Parameters:

  • chunk_id (str) – Unique chunk identifier.

  • content (str) – Chunk text content.

  • embedding (list[float] | None) – Chunk embedding vector (optional).

  • parent_id (str | None) – Parent chunk ID (None for root/parent chunks).

  • hierarchy_level (int) – Level in hierarchy (0=parent, 1+=children).

  • metadata (dict[str, Any] | None) – Optional metadata dictionary.

Return type:

None

get_chunk(chunk_id)[source]

Get a chunk by ID.

Parameters:

chunk_id (str) – Chunk identifier.

Returns:

Chunk data dictionary or None if not found.

Return type:

dict[str, Any] | None

get_parent(chunk_id)[source]

Get the parent of a chunk.

Parameters:

chunk_id (str) – Child chunk identifier.

Returns:

Parent chunk data or None if not found or no parent.

Return type:

dict[str, Any] | None

get_children(parent_id)[source]

Get all children of a parent chunk.

Parameters:

parent_id (str) – Parent chunk identifier.

Returns:

List of child chunk data dictionaries.

Return type:

list[dict[str, Any]]

get_parent_chunks()[source]

Get all parent chunks (level 0).

Returns:

List of parent chunk data dictionaries.

Return type:

list[dict[str, Any]]

get_child_chunks()[source]

Get all child chunks (level > 0).

Returns:

List of child chunk data dictionaries.

Return type:

list[dict[str, Any]]

get_child_embeddings()[source]

Get IDs and embeddings for all child chunks.

Returns:

Tuple of (chunk_ids, embeddings).

Return type:

tuple[list[str], list[list[float]]]

enhance_results_with_context(results, include_parent=None)[source]

Enhance search results with parent context.

Parameters:

  • results (list[dict[str, Any]]) – Search results with chunk IDs.

  • include_parent (bool | None) – Override default include_parent_context.

Returns:

Enhanced results with parent content added.

Return type:

list[dict[str, Any]]

delete_chunk(chunk_id)[source]

Delete a chunk and update relationships.

Parameters:

chunk_id (str) – Chunk identifier to delete.

Returns:

True if deleted, False if not found.

Return type:

bool

delete_by_reference(reference_name)[source]

Delete all chunks for a document reference.

Parameters:

reference_name (str) – Document reference name in metadata.

Returns:

Number of chunks deleted.

Return type:

int

clear()[source]

Clear all chunks from the index.

Return type:

None

count()[source]

Get total number of chunks.

Return type:

int

count_parents()[source]

Get number of parent chunks.

Return type:

int

count_children()[source]

Get number of child chunks.

Return type:

int

get_stats()[source]

Get index statistics.

Return type:

dict[str, Any]

class HierarchicalTextChunk(content, chunk_index, start_char, end_char, metadata=None, id='', parent_id=None, child_ids=None, hierarchy_level=0)[source]

Bases: TextChunk

A text chunk that participates in a parent-child hierarchy.

Used by hierarchical chunking strategies that maintain relationships between larger parent chunks and smaller child chunks.

Parameters:
id

Unique identifier for this chunk.

Type:

str

parent_id

Identifier of the parent chunk, or None if this is a root.

Type:

str | None

child_ids

Identifiers of child chunks, or an empty list.

Type:

list[str] | None

hierarchy_level

Depth in the hierarchy (0 = root).

Type:

int

id: str = ''
parent_id: str | None = None
child_ids: list[str] | None = None
hierarchy_level: int = 0
class HybridSearch(vector_search_fn=None, bm25_search_fn=None, alpha=0.6, rrf_k=60)[source]

Bases: object

Hybrid search combining vector similarity and BM25 keyword matching.

Uses reciprocal rank fusion (RRF) to combine results from multiple retrieval methods.

Parameters:

  • vector_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Function that takes (query, n_results) and returns vector results.

  • bm25_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Function that takes (query, n_results) and returns BM25 results.

  • alpha (float) – Weight for vector search (1-alpha for BM25). Default 0.6.

  • rrf_k (int) – RRF constant for rank fusion. Default 60.

search(query, n_results=5, mode='hybrid')[source]

Perform search using the specified mode.

Parameters:

  • query (str) – Search query.

  • n_results (int) – Maximum number of results.

  • mode (str) – Search mode - “vector”, “bm25”, or “hybrid”.

Returns:

List of search results with merged scores.

Return type:

list[dict[str, Any]]

set_alpha(alpha)[source]

Set the alpha parameter (vector weight).

Parameters:

alpha (float) – Weight for vector search (0.0 to 1.0).

Return type:

None

set_search_functions(vector_search_fn=None, bm25_search_fn=None)[source]

Set or update search functions.

Parameters:
Return type:

None

class MarkdownChunker(chunk_size=1000, chunk_overlap=200, metadata=None, split_headers=None, preserve_structure=True, max_chunk_fallback=True)[source]

Bases: ChunkerBase

Markdown-aware chunking that splits by headers.

Parameters:

  • chunk_size (int) – Maximum characters per chunk.

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • split_headers (list[str] | None) – List of header levels to split on (e.g., [“h1”, “h2”]).

  • preserve_structure (bool) – Whether to include header in chunk content.

  • max_chunk_fallback (bool) – Whether to further split large sections.

HEADER_PATTERNS = {'h1': '^# .+', 'h2': '^## .+', 'h3': '^### .+', 'h4': '^#### .+', 'h5': '^##### .+', 'h6': '^###### .+'}
chunk(text, metadata=None)[source]

Split markdown text by headers while respecting size limits.

Parameters:

  • text (str) – The markdown text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

class NoopReranker[source]

Bases: RerankerBase

Pass-through re-ranker that does nothing.

Used when re-ranking is disabled but a reranker interface is expected.

rerank(query, results, n_results=None)[source]

Return results unchanged.

Parameters:

  • query (str) – Original search query (ignored).

  • results (list[dict[str, Any]]) – Search results.

  • n_results (int | None) – Number of results to return.

Returns:

Original results, optionally truncated.

Return type:

list[dict[str, Any]]

class QueryExpander(llm_generate_fn=None, n_variations=3, include_original=True)[source]

Bases: object

Expand queries using LLM for multi-query retrieval.

Generates multiple reformulations of a search query to improve recall by catching different phrasings and aspects of the topic.

Parameters:

  • llm_generate_fn (Callable[[str], Any] | None) – Function that takes a prompt string and returns LLM response.

  • n_variations (int) – Number of query variations to generate (default: 3).

  • include_original (bool) – Whether to include original query in results (default: True).

Example

>>> def mock_llm(prompt): return "1. What is auth?\n2. How to authenticate?"
>>> expander = QueryExpander(llm_generate_fn=mock_llm, n_variations=2)
>>> queries = expander.expand("authentication methods")
>>> # Returns: ["authentication methods", "What is auth?", "How to authenticate?"]
expand(query)[source]

Generate query variations for improved retrieval.

Parameters:

query (str) – Original search query.

Returns:

List of query variations including original (if include_original=True).

Return type:

list[str]

set_llm_function(fn)[source]

Set or update the LLM generate function.

Parameters:

fn (Callable[[str], Any]) – Function that takes a prompt and returns LLM response.

Return type:

None

class QueryResult(answer, hits, sources, prompt, usage=None, cost_usd=0.0, duration_ms=None)[source]

Bases: object

Result of a RAG query combining search and generation.

Parameters:
answer

Generated answer text.

Type:

str

hits

Search hits used as context.

Type:

list[ffai.rag.types.SearchHit]

sources

Deduplicated source identifiers.

Type:

list[str]

prompt

Full prompt sent to the generation function.

Type:

str

usage

Provider-specific token usage object.

Type:

Any | None

cost_usd

Estimated generation cost in USD.

Type:

float

duration_ms

Generation duration in milliseconds.

Type:

float | None

answer: str
hits: list[SearchHit]
sources: list[str]
prompt: str
usage: Any | None = None
cost_usd: float = 0.0
duration_ms: float | None = None
class RecursiveChunker(chunk_size=1000, chunk_overlap=200, metadata=None, separators=None, keep_separator=True)[source]

Bases: ChunkerBase

Recursive chunking that splits hierarchically by separators.

Parameters:

  • chunk_size (int) – Maximum characters per chunk.

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • separators (list[str] | None) – List of separators in order of preference.

  • keep_separator (bool) – Whether to keep the separator with chunks.

DEFAULT_SEPARATORS = ['\n\n\n', '\n\n', '\n', '. ', '! ', '? ', '; ', ', ', ' ', '']
chunk(text, metadata=None)[source]

Split text recursively using hierarchical separators.

Parameters:

  • text (str) – The text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

class RerankerBase[source]

Bases: object

Base class for re-rankers.

rerank(query, results, n_results=None)[source]

Re-rank search results.

Parameters:

  • query (str) – Original search query.

  • results (list[dict[str, Any]]) – Search results to re-rank.

  • n_results (int | None) – Number of results to return (None = all).

Returns:

Re-ranked results.

Return type:

list[dict[str, Any]]

class SearchHit(content, score, source='', metadata=<factory>, parent_content=None, id='')[source]

Bases: object

Single search result from a RAG query.

Parameters:
content

Matched chunk text.

Type:

str

score

Relevance score (higher is better).

Type:

float

source

Source document identifier.

Type:

str

metadata

Additional metadata from indexing.

Type:

dict[str, Any]

parent_content

Parent chunk text for hierarchical context.

Type:

str | None

id

Unique chunk identifier.

Type:

str

content: str
score: float
source: str = ''
metadata: dict[str, Any]
parent_content: str | None = None
id: str = ''
class TextChunk(content, chunk_index, start_char, end_char, metadata=None)[source]

Bases: object

A single chunk of text produced by a chunker.

Parameters:
content

The text content of the chunk.

Type:

str

chunk_index

Zero-based index of this chunk in the sequence.

Type:

int

start_char

Character offset where this chunk begins in the source.

Type:

int

end_char

Character offset where this chunk ends in the source.

Type:

int

metadata

Optional metadata attached to this chunk.

Type:

dict[str, Any] | None

content: str
chunk_index: int
start_char: int
end_char: int
metadata: dict[str, Any] | None = None
VectorStore

alias of ChromaVectorStore

class VectorStoreBase[source]

Bases: ABC

Abstract base class for vector store backends.

All backends must implement the core CRUD + search methods. The where parameter in asearch accepts a backend-neutral filter dict with string keys and values. Each backend translates this to its native filter format internally.

Example backend-neutral filters:

{"source": "doc1"}
{"chunking_strategy": "recursive", "source": "doc1"}

For compound filters, backends should support an $and key:

{"$and": [{"source": "doc1"}, {"chunking_strategy": "recursive"}]}
abstract property name: str

Backend identifier (e.g. "chroma", "pgvector").

abstractmethod async aadd(ids, texts, embeddings, metadatas)[source]

Add documents with pre-computed embeddings to the store.

Parameters:

  • ids (list[str]) – Unique identifiers for each document.

  • texts (list[str]) – Document text content.

  • embeddings (list[list[float]]) – Pre-computed embedding vectors.

  • metadatas (list[dict[str, Any]]) – Metadata dicts (must include source key).

Returns:

Number of documents added.

Return type:

int

abstractmethod async asearch(query_embedding, top_k=5, where=None)[source]

Search for documents by vector similarity.

Parameters:

  • query_embedding (list[float]) – Query vector.

  • top_k (int) – Maximum number of results.

  • where (dict[str, Any] | None) – Metadata filter dict (e.g. {"source": "doc1"}).

Returns:

Ranked search hits.

Return type:

list[SearchHit]

abstractmethod delete_by_source(source)[source]

Delete all chunks matching source.

Parameters:

source (str)

Return type:

None

abstractmethod delete_by_source_and_strategy(source, strategy)[source]

Delete chunks matching both source and chunking_strategy.

Parameters:
Return type:

None

abstractmethod count()[source]

Return the total number of stored chunks.

Return type:

int

abstractmethod clear()[source]

Delete all stored data and recreate the collection.

Return type:

None

abstractmethod list_sources()[source]

Return a sorted list of indexed source names.

Return type:

list[str]

abstractmethod get_all()[source]

Return all stored documents as dicts with id, content, metadata keys.

Return type:

list[dict[str, Any]]

abstractmethod needs_reindex(source, checksum, strategy='default')[source]

Check whether source needs re-indexing.

Parameters:

  • source (str) – Source identifier.

  • checksum (str) – Expected document checksum.

  • strategy (str) – Chunking strategy name.

Returns:

True if the source has changed or is not indexed.

Return type:

bool

chunk_text(text, strategy='recursive', chunk_size=1000, chunk_overlap=200, metadata=None, **kwargs)[source]

Convenience function to chunk text with a single call.

Parameters:

  • text (str) – Text to chunk.

  • strategy (str) – Chunking strategy name.

  • chunk_size (int) – Maximum chunk size.

  • chunk_overlap (int) – Overlap between chunks.

  • metadata (dict[str, Any] | None) – Metadata to attach to chunks.

  • **kwargs (Any) – Additional strategy-specific parameters.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

format_hits(hits, max_chars=None, include_parent=True)[source]

Format search hits into a numbered, human-readable text string.

Each hit is rendered with its index, source, relevance score, content, and optionally a snippet of its parent context.

Parameters:

  • hits (list[SearchHit]) – List of SearchHit objects to format.

  • max_chars (int | None) – Optional character budget; hits beyond the budget are omitted.

  • include_parent (bool) – Whether to include parent-context snippets.

Returns:

A formatted string, or an empty string if hits is empty.

Return type:

str

fuse_search_results(result_lists, n_results=5, dedupe_by='id')[source]

Fuse results from multiple searches with deduplication.

Combines results from multiple query variations, removing duplicates while preserving relevance ordering.

Parameters:

  • result_lists (list[list[dict[str, Any]]]) – List of search result lists from different queries.

  • n_results (int) – Maximum number of final results to return.

  • dedupe_by (str) – Field to use for deduplication (default: “id”).

Returns:

Fused and deduplicated list of results.

Return type:

list[dict[str, Any]]

get_chunker(strategy='recursive', chunk_size=1000, chunk_overlap=200, metadata=None, **kwargs)[source]

Get a chunker instance based on strategy name.

Parameters:

  • strategy (str) – Chunking strategy name (character, recursive, markdown, code, hierarchical).

  • chunk_size (int) – Maximum chunk size.

  • chunk_overlap (int) – Overlap between chunks.

  • metadata (dict[str, Any] | None) – Default metadata for chunks.

  • **kwargs (Any) – Additional strategy-specific parameters.

Returns:

Configured chunker instance.

Raises:

ValueError – If strategy name is not recognized.

Return type:

ChunkerBase

get_reranker(reranker_type='none', **kwargs)[source]

Get a reranker by type name.

Parameters:

  • reranker_type (str) – Type of reranker (“cross_encoder”, “diversity”, “none”).

  • **kwargs (Any) – Additional arguments for the reranker.

Returns:

Configured reranker instance.

Return type:

RerankerBase

get_store(backend='chroma', **kwargs)[source]

Get a vector store instance by backend name.

Parameters:

  • backend (str) – Backend name ("chroma", "pgvector", "qdrant", "sqlite_vss").

  • **kwargs (Any) – Backend-specific constructor arguments.

Returns:

Configured vector store instance.

Raises:

  • ValueError – If backend name is not recognized.

  • ImportError – If the backend’s dependency is not installed.

Return type:

VectorStoreBase

is_store_available(name)[source]

Check if a backend’s dependencies are installed.

Parameters:

name (str)

Return type:

bool

list_available_stores()[source]

List backend names whose dependencies are installed.

Return type:

list[str]

list_chunkers()[source]

List available chunking strategies.

Returns:

List of strategy names.

Return type:

list[str]

list_stores()[source]

List all known backend names (regardless of availability).

Return type:

list[str]

litellm_generate_fn(model, api_key=None, temperature=0.5, max_tokens=1024, **kwargs)[source]
Parameters:
Return type:

Callable[[str], GenerationResult]

reciprocal_rank_fusion(result_lists, k=60, weights=None)[source]

Combine multiple result lists using reciprocal rank fusion.

Parameters:

  • result_lists (list[list[dict[str, Any]]]) – List of result lists to fuse.

  • k (int) – RRF constant.

  • weights (list[float] | None) – Optional weights for each result list.

Returns:

Fused and sorted results.

Return type:

list[dict[str, Any]]

Classes

class BM25Index(k1=1.5, b=0.75, epsilon=0.25)[source]

Bases: object

BM25 sparse index for keyword-based retrieval.

Implements the Okapi BM25 algorithm for text relevance scoring. Used alongside vector search for hybrid retrieval.

Parameters:

  • k1 (float) – BM25 k1 parameter (term frequency saturation).

  • b (float) – BM25 b parameter (document length normalization).

  • epsilon (float) – Small value to avoid division by zero.

tokenize(text)[source]

Tokenize text into terms.

Parameters:

text (str) – Text to tokenize.

Returns:

List of lowercase tokens.

Return type:

list[str]

add_document(doc_id, content, metadata=None)[source]

Add a document to the index.

Parameters:

  • doc_id (str) – Unique document identifier.

  • content (str) – Document text content.

  • metadata (dict[str, Any] | None) – Optional metadata dictionary.

Return type:

None

add_documents(documents, id_key='id', content_key='content')[source]

Add multiple documents to the index.

Parameters:

  • documents (list[dict[str, Any]]) – List of document dictionaries.

  • id_key (str) – Key for document ID in each dict.

  • content_key (str) – Key for document content in each dict.

Returns:

Number of documents added.

Return type:

int

delete_document(doc_id)[source]

Delete a document from the index.

Parameters:

doc_id (str) – Document ID to delete.

Returns:

True if document was deleted, False if not found.

Return type:

bool

search(query, n_results=5)[source]

Search for documents using BM25 scoring.

Parameters:

  • query (str) – Search query.

  • n_results (int) – Maximum number of results.

Returns:

List of result dicts with doc_id, score, content, metadata.

Return type:

list[dict[str, Any]]

clear()[source]

Clear all documents from the index.

Return type:

None

delete_by_metadata(key, value)[source]

Delete all documents matching a metadata key-value pair.

Parameters:

  • key (str) – Metadata key to match.

  • value (Any) – Metadata value to match.

Returns:

Number of documents deleted.

Return type:

int

count()[source]

Get the number of documents in the index.

Return type:

int

get_stats()[source]

Get index statistics.

Return type:

dict[str, Any]

class CharacterChunker(chunk_size=1000, chunk_overlap=200, metadata=None, respect_word_boundaries=True)[source]

Bases: ChunkerBase

Character-based chunking with word-boundary awareness.

Parameters:

  • chunk_size (int) – Maximum characters per chunk.

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • respect_word_boundaries (bool) – Whether to prefer word boundaries for splits.

chunk(text, metadata=None)[source]

Split text into overlapping chunks.

Parameters:

  • text (str) – The text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

property name: str

Return the chunker strategy name.

class ChunkDeduplicator(mode='exact', similarity_threshold=0.95)[source]

Bases: object

Detect and filter duplicate/near-duplicate chunks.

Supports multiple deduplication modes: - exact: Hash-based exact content matching - similarity: Embedding cosine similarity threshold

Parameters:

  • mode (str) – Deduplication mode (“exact” or “similarity”).

  • similarity_threshold (float) – Threshold for similarity-based dedup (0.0-1.0).

Example

>>> dedup = ChunkDeduplicator(mode="exact")
>>> chunks, embeddings = dedup.filter_duplicates(chunks, embeddings)
compute_hash(content)[source]

Compute content hash for exact deduplication.

Parameters:

content (str) – Chunk text content.

Returns:

SHA256 hash (first 16 chars).

Return type:

str

is_duplicate(content, embedding=None)[source]

Check if chunk is a duplicate of previously seen content.

Parameters:

  • content (str) – Chunk text content.

  • embedding (list[float] | None) – Optional embedding vector for similarity-based dedup.

Returns:

True if chunk is a duplicate.

Return type:

bool

filter_duplicates(chunks, embeddings)[source]

Filter duplicate chunks from a batch.

Parameters:

  • chunks (list[Any]) – List of chunk objects (must have .content attribute or be string).

  • embeddings (list[list[float]]) – List of embedding vectors.

Returns:

Tuple of (filtered_chunks, filtered_embeddings).

Return type:

tuple[list[Any], list[list[float]]]

clear()[source]

Clear seen hashes and embeddings.

Return type:

None

get_stats()[source]

Get deduplicator statistics.

Returns:

Dict with mode, threshold, and counts.

Return type:

dict[str, Any]

class ChunkerBase(chunk_size=1000, chunk_overlap=200, metadata=None)[source]

Bases: ABC

Abstract base class for text chunking strategies.

Parameters:

  • chunk_size (int) – Maximum size of each chunk (interpretation varies by strategy).

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata to attach to all chunks.

abstractmethod chunk(text, metadata=None)[source]

Split text into chunks.

Parameters:

  • text (str) – The text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk (merged with default).

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

property name: str

Return the chunker strategy name.

class ClientAdapter(client, **kwargs)[source]

Bases: object

Wrap an FFAI client into a callable that returns a GenerationResult.

Handles both synchronous and asynchronous generate_response methods, automatically awaiting coroutines when necessary.

Parameters:

  • client (Any) – An FFAI client instance with a generate_response method.

  • **kwargs (Any) – Extra keyword arguments forwarded to generate_response on every call.

__call__(prompt)[source]

Generate a response for the given prompt.

Parameters:

prompt (str) – The prompt string to send to the client.

Returns:

A GenerationResult containing the response text, usage, cost, and duration metadata.

Return type:

GenerationResult

class CodeChunker(chunk_size=1000, chunk_overlap=200, metadata=None, language='python', split_by='function')[source]

Bases: ChunkerBase

Code-aware chunking that splits by functions/classes.

Parameters:

  • chunk_size (int) – Maximum characters per chunk.

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • language (str) – Programming language for parsing hints.

  • split_by (str) – Strategy: “function”, “class”, or “module”.

LANGUAGE_PATTERNS: dict[str, dict[str, str]] = {'generic': {'class': '^(class|struct|interface)\\s+\\w+', 'comment': '^(#|//|/\\*)', 'function': '^(function|func|def|fn)\\s+\\w+', 'import': '^(import|use|require)'}, 'go': {'class': '^type\\s+\\w+\\s+struct', 'comment': '^(//|/\\*)', 'function': '^func\\s+(\\(\\w+\\s+\\*?\\w+\\)\\s+)?\\w+\\s*\\(', 'import': '^import\\s+'}, 'java': {'class': '^\\s*(public\\s+)?(abstract\\s+)?class\\s+\\w+|interface\\s+\\w+', 'comment': '^(//|/\\*)', 'function': '^\\s*(public|private|protected|static)?\\s*\\w+\\s+\\w+\\s*\\(', 'import': '^import\\s+'}, 'javascript': {'class': '^(export\\s+)?class\\s+\\w+', 'comment': '^(//|/\\*|\\*)', 'function': '^(async\\s+)?function\\s+\\w+|const\\s+\\w+\\s*=\\s*(async\\s+)?\\([^)]*\\)\\s*=>|export\\s+(async\\s+)?function', 'import': '^(import\\s+|export\\s+|require\\s*\\()'}, 'python': {'class': '^class\\s+\\w+[\\(:]', 'comment': '^(#|\'\'\'|\\"\\"\\")', 'decorator': '^@\\w+', 'function': '^(async\\s+)?def\\s+\\w+\\s*\\(', 'import': '^(import\\s+|from\\s+\\S+\\s+import)'}, 'rust': {'class': '^(pub\\s+)?struct\\s+\\w+|impl\\s+\\w+', 'comment': '^(//|/\\*|\\*)', 'function': '^(pub\\s+)?(async\\s+)?fn\\s+\\w+', 'import': '^use\\s+'}, 'typescript': {'class': '^(export\\s+)?(abstract\\s+)?class\\s+\\w+|interface\\s+\\w+|type\\s+\\w+', 'comment': '^(//|/\\*|\\*)', 'function': '^(async\\s+)?(function\\s+\\w+|const\\s+\\w+\\s*=\\s*(async\\s+)?\\([^)]*\\)\\s*(:\\s*\\w+)?\\s*=>|export\\s+(async\\s+)?function)', 'import': '^(import\\s+|export\\s+|require\\s*\\()'}}
chunk(text, metadata=None)[source]

Split code by structural boundaries.

Parameters:

  • text (str) – The code text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

property name: str

Return the chunker strategy name.

class ContextualEmbeddings(context_template=None, max_context_length=200)[source]

Bases: object

Generate embeddings with document context prepended.

Prepends document context (title, summary, or preceding content) to each chunk before embedding, improving semantic understanding.

Parameters:

  • context_prefix – Template for context prefix.

  • max_context_length (int) – Maximum characters for context prefix.

  • context_template (str | None)

DEFAULT_CONTEXT_TEMPLATE = 'Document: {title}\n\nSection: {section}\n\n{chunk}'
prepare_chunk_for_embedding(chunk_content, document_title=None, section_header=None, document_summary=None, preceding_context=None)[source]

Prepare a chunk with context for embedding.

Parameters:

  • chunk_content (str) – The chunk text content.

  • document_title (str | None) – Document title or name.

  • section_header (str | None) – Section header if available.

  • document_summary (str | None) – Brief document summary.

  • preceding_context (str | None) – Text immediately preceding this chunk.

Returns:

Context-enhanced text for embedding.

Return type:

str

prepare_chunks_batch(chunks, document_title=None, document_summary=None)[source]

Prepare multiple chunks with context for embedding.

Parameters:

  • chunks (list[dict[str, Any]]) – List of chunk dictionaries with ‘content’ and optional metadata.

  • document_title (str | None) – Document title for all chunks.

  • document_summary (str | None) – Document summary (currently unused but available).

Returns:

List of context-enhanced texts for embedding.

Return type:

list[str]

class CrossEncoderReranker(model_name='cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512, fastembed_model_name=None)[source]

Bases: RerankerBase

Re-ranker using cross-encoder models.

Tries sentence-transformers first, then falls back to fastembed (ONNX-based). Requires at least one of:

pip install sentence-transformers pip install fastembed

Parameters:

  • model_name (str) – Cross-encoder model name.

  • max_length (int) – Maximum sequence length (sentence-transformers only).

  • fastembed_model_name (str | None) – Model name for fastembed fallback.

rerank(query, results, n_results=None)[source]

Re-rank results using cross-encoder scoring.

Parameters:

  • query (str) – Original search query.

  • results (list[dict[str, Any]]) – Search results to re-rank.

  • n_results (int | None) – Number of results to return.

Returns:

Re-ranked results with updated scores.

Return type:

list[dict[str, Any]]

class DiversityReranker(lambda_param=0.7)[source]

Bases: RerankerBase

Re-ranker that promotes result diversity.

Re-orders results to maximize diversity based on content similarity. Uses MMR (Maximal Marginal Relevance) style selection.

Parameters:

lambda_param (float) – Balance between relevance and diversity (0-1). Higher = more relevance, lower = more diversity.

rerank(query, results, n_results=None)[source]

Re-rank results for diversity.

Parameters:

  • query (str) – Original search query (not used, kept for interface).

  • results (list[dict[str, Any]]) – Search results to re-rank.

  • n_results (int | None) – Number of results to return.

Returns:

Diversified results.

Return type:

list[dict[str, Any]]

class GenerationResult(text, usage=None, cost_usd=0.0, duration_ms=None)[source]

Bases: object

Result of a RAG generation call.

Return this from generate_fn to preserve token usage, cost, and timing metadata. Returning a plain string silently discards these metrics and emits a warning.

Parameters:
text

Generated answer text.

Type:

str

usage

Provider-specific token usage object (e.g. TokenUsage).

Type:

Any | None

cost_usd

Estimated cost in USD.

Type:

float

duration_ms

Wall-clock generation duration in milliseconds.

Type:

float | None

text: str
usage: Any | None = None
cost_usd: float = 0.0
duration_ms: float | None = None
class HierarchicalChunker(chunk_size=400, chunk_overlap=100, metadata=None, parent_chunk_size=1500, max_levels=2)[source]

Bases: ChunkerBase

Hierarchical chunking with parent-child relationships.

Parameters:

  • chunk_size (int) – Maximum characters per leaf chunk.

  • chunk_overlap (int) – Overlap between leaf chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • parent_chunk_size (int) – Size of parent chunks (larger).

  • max_levels (int) – Maximum hierarchy depth.

chunk(text, metadata=None)[source]

Split text into hierarchical chunks with parent-child relationships.

Parameters:

  • text (str) – The text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of HierarchicalTextChunk objects with parent-child links.

Return type:

list[HierarchicalTextChunk]

get_parent_chunks(chunks)[source]

Filter to return only parent chunks.

Parameters:

chunks (list[HierarchicalTextChunk])

Return type:

list[HierarchicalTextChunk]

get_child_chunks(chunks)[source]

Filter to return only child chunks.

Parameters:

chunks (list[HierarchicalTextChunk])

Return type:

list[HierarchicalTextChunk]

get_chunks_with_parent_context(child_chunks, all_chunks)[source]

Get child chunks with their parent context for retrieval.

Parameters:
Returns:

List of dicts with child content and parent context.

Return type:

list[dict[str, Any]]

property name: str

Return the chunker strategy name.

class HierarchicalIndex(include_parent_context=True)[source]

Bases: object

Index for hierarchical chunk storage and retrieval.

Stores both parent and child chunks, enabling: - Fine-grained search on child chunks - Parent context retrieval for matched children

Parameters:

include_parent_context (bool) – Whether to include parent content in results.

add_chunk(chunk_id, content, embedding=None, parent_id=None, hierarchy_level=0, metadata=None)[source]

Add a chunk to the hierarchical index.

Parameters:

  • chunk_id (str) – Unique chunk identifier.

  • content (str) – Chunk text content.

  • embedding (list[float] | None) – Chunk embedding vector (optional).

  • parent_id (str | None) – Parent chunk ID (None for root/parent chunks).

  • hierarchy_level (int) – Level in hierarchy (0=parent, 1+=children).

  • metadata (dict[str, Any] | None) – Optional metadata dictionary.

Return type:

None

get_chunk(chunk_id)[source]

Get a chunk by ID.

Parameters:

chunk_id (str) – Chunk identifier.

Returns:

Chunk data dictionary or None if not found.

Return type:

dict[str, Any] | None

get_parent(chunk_id)[source]

Get the parent of a chunk.

Parameters:

chunk_id (str) – Child chunk identifier.

Returns:

Parent chunk data or None if not found or no parent.

Return type:

dict[str, Any] | None

get_children(parent_id)[source]

Get all children of a parent chunk.

Parameters:

parent_id (str) – Parent chunk identifier.

Returns:

List of child chunk data dictionaries.

Return type:

list[dict[str, Any]]

get_parent_chunks()[source]

Get all parent chunks (level 0).

Returns:

List of parent chunk data dictionaries.

Return type:

list[dict[str, Any]]

get_child_chunks()[source]

Get all child chunks (level > 0).

Returns:

List of child chunk data dictionaries.

Return type:

list[dict[str, Any]]

get_child_embeddings()[source]

Get IDs and embeddings for all child chunks.

Returns:

Tuple of (chunk_ids, embeddings).

Return type:

tuple[list[str], list[list[float]]]

enhance_results_with_context(results, include_parent=None)[source]

Enhance search results with parent context.

Parameters:

  • results (list[dict[str, Any]]) – Search results with chunk IDs.

  • include_parent (bool | None) – Override default include_parent_context.

Returns:

Enhanced results with parent content added.

Return type:

list[dict[str, Any]]

delete_chunk(chunk_id)[source]

Delete a chunk and update relationships.

Parameters:

chunk_id (str) – Chunk identifier to delete.

Returns:

True if deleted, False if not found.

Return type:

bool

delete_by_reference(reference_name)[source]

Delete all chunks for a document reference.

Parameters:

reference_name (str) – Document reference name in metadata.

Returns:

Number of chunks deleted.

Return type:

int

clear()[source]

Clear all chunks from the index.

Return type:

None

count()[source]

Get total number of chunks.

Return type:

int

count_parents()[source]

Get number of parent chunks.

Return type:

int

count_children()[source]

Get number of child chunks.

Return type:

int

get_stats()[source]

Get index statistics.

Return type:

dict[str, Any]

class HierarchicalTextChunk(content, chunk_index, start_char, end_char, metadata=None, id='', parent_id=None, child_ids=None, hierarchy_level=0)[source]

Bases: TextChunk

A text chunk that participates in a parent-child hierarchy.

Used by hierarchical chunking strategies that maintain relationships between larger parent chunks and smaller child chunks.

Parameters:
id

Unique identifier for this chunk.

Type:

str

parent_id

Identifier of the parent chunk, or None if this is a root.

Type:

str | None

child_ids

Identifiers of child chunks, or an empty list.

Type:

list[str] | None

hierarchy_level

Depth in the hierarchy (0 = root).

Type:

int

id: str = ''
parent_id: str | None = None
child_ids: list[str] | None = None
hierarchy_level: int = 0
metadata: dict[str, Any] | None = None
content: str
chunk_index: int
start_char: int
end_char: int
class HybridSearch(vector_search_fn=None, bm25_search_fn=None, alpha=0.6, rrf_k=60)[source]

Bases: object

Hybrid search combining vector similarity and BM25 keyword matching.

Uses reciprocal rank fusion (RRF) to combine results from multiple retrieval methods.

Parameters:

  • vector_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Function that takes (query, n_results) and returns vector results.

  • bm25_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Function that takes (query, n_results) and returns BM25 results.

  • alpha (float) – Weight for vector search (1-alpha for BM25). Default 0.6.

  • rrf_k (int) – RRF constant for rank fusion. Default 60.

search(query, n_results=5, mode='hybrid')[source]

Perform search using the specified mode.

Parameters:

  • query (str) – Search query.

  • n_results (int) – Maximum number of results.

  • mode (str) – Search mode - “vector”, “bm25”, or “hybrid”.

Returns:

List of search results with merged scores.

Return type:

list[dict[str, Any]]

set_alpha(alpha)[source]

Set the alpha parameter (vector weight).

Parameters:

alpha (float) – Weight for vector search (0.0 to 1.0).

Return type:

None

set_search_functions(vector_search_fn=None, bm25_search_fn=None)[source]

Set or update search functions.

Parameters:
Return type:

None

class MarkdownChunker(chunk_size=1000, chunk_overlap=200, metadata=None, split_headers=None, preserve_structure=True, max_chunk_fallback=True)[source]

Bases: ChunkerBase

Markdown-aware chunking that splits by headers.

Parameters:

  • chunk_size (int) – Maximum characters per chunk.

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • split_headers (list[str] | None) – List of header levels to split on (e.g., [“h1”, “h2”]).

  • preserve_structure (bool) – Whether to include header in chunk content.

  • max_chunk_fallback (bool) – Whether to further split large sections.

HEADER_PATTERNS = {'h1': '^# .+', 'h2': '^## .+', 'h3': '^### .+', 'h4': '^#### .+', 'h5': '^##### .+', 'h6': '^###### .+'}
chunk(text, metadata=None)[source]

Split markdown text by headers while respecting size limits.

Parameters:

  • text (str) – The markdown text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

property name: str

Return the chunker strategy name.

class NoopReranker[source]

Bases: RerankerBase

Pass-through re-ranker that does nothing.

Used when re-ranking is disabled but a reranker interface is expected.

rerank(query, results, n_results=None)[source]

Return results unchanged.

Parameters:

  • query (str) – Original search query (ignored).

  • results (list[dict[str, Any]]) – Search results.

  • n_results (int | None) – Number of results to return.

Returns:

Original results, optionally truncated.

Return type:

list[dict[str, Any]]

class QueryExpander(llm_generate_fn=None, n_variations=3, include_original=True)[source]

Bases: object

Expand queries using LLM for multi-query retrieval.

Generates multiple reformulations of a search query to improve recall by catching different phrasings and aspects of the topic.

Parameters:

  • llm_generate_fn (Callable[[str], Any] | None) – Function that takes a prompt string and returns LLM response.

  • n_variations (int) – Number of query variations to generate (default: 3).

  • include_original (bool) – Whether to include original query in results (default: True).

Example

>>> def mock_llm(prompt): return "1. What is auth?\n2. How to authenticate?"
>>> expander = QueryExpander(llm_generate_fn=mock_llm, n_variations=2)
>>> queries = expander.expand("authentication methods")
>>> # Returns: ["authentication methods", "What is auth?", "How to authenticate?"]
expand(query)[source]

Generate query variations for improved retrieval.

Parameters:

query (str) – Original search query.

Returns:

List of query variations including original (if include_original=True).

Return type:

list[str]

set_llm_function(fn)[source]

Set or update the LLM generate function.

Parameters:

fn (Callable[[str], Any]) – Function that takes a prompt and returns LLM response.

Return type:

None

class QueryResult(answer, hits, sources, prompt, usage=None, cost_usd=0.0, duration_ms=None)[source]

Bases: object

Result of a RAG query combining search and generation.

Parameters:
answer

Generated answer text.

Type:

str

hits

Search hits used as context.

Type:

list[ffai.rag.types.SearchHit]

sources

Deduplicated source identifiers.

Type:

list[str]

prompt

Full prompt sent to the generation function.

Type:

str

usage

Provider-specific token usage object.

Type:

Any | None

cost_usd

Estimated generation cost in USD.

Type:

float

duration_ms

Generation duration in milliseconds.

Type:

float | None

answer: str
hits: list[SearchHit]
sources: list[str]
prompt: str
usage: Any | None = None
cost_usd: float = 0.0
duration_ms: float | None = None
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:

int

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:

int

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:

list[SearchHit]

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:

list[SearchHit]

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:
Return type:

QueryResult

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.

  1. Searches for relevant chunks via asearch.

  2. Formats hits into a context string (format_hits).

  3. Fills the prompt template with {context} and {question}.

  4. 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:

QueryResult

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

class RecursiveChunker(chunk_size=1000, chunk_overlap=200, metadata=None, separators=None, keep_separator=True)[source]

Bases: ChunkerBase

Recursive chunking that splits hierarchically by separators.

Parameters:

  • chunk_size (int) – Maximum characters per chunk.

  • chunk_overlap (int) – Overlap between consecutive chunks.

  • metadata (dict[str, Any] | None) – Default metadata for all chunks.

  • separators (list[str] | None) – List of separators in order of preference.

  • keep_separator (bool) – Whether to keep the separator with chunks.

DEFAULT_SEPARATORS = ['\n\n\n', '\n\n', '\n', '. ', '! ', '? ', '; ', ', ', ' ', '']
chunk(text, metadata=None)[source]

Split text recursively using hierarchical separators.

Parameters:

  • text (str) – The text to split.

  • metadata (dict[str, Any] | None) – Optional metadata to attach to each chunk.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

property name: str

Return the chunker strategy name.

class RerankerBase[source]

Bases: object

Base class for re-rankers.

rerank(query, results, n_results=None)[source]

Re-rank search results.

Parameters:

  • query (str) – Original search query.

  • results (list[dict[str, Any]]) – Search results to re-rank.

  • n_results (int | None) – Number of results to return (None = all).

Returns:

Re-ranked results.

Return type:

list[dict[str, Any]]

class SearchHit(content, score, source='', metadata=<factory>, parent_content=None, id='')[source]

Bases: object

Single search result from a RAG query.

Parameters:
content

Matched chunk text.

Type:

str

score

Relevance score (higher is better).

Type:

float

source

Source document identifier.

Type:

str

metadata

Additional metadata from indexing.

Type:

dict[str, Any]

parent_content

Parent chunk text for hierarchical context.

Type:

str | None

id

Unique chunk identifier.

Type:

str

content: str
score: float
source: str = ''
metadata: dict[str, Any]
parent_content: str | None = None
id: str = ''
class TextChunk(content, chunk_index, start_char, end_char, metadata=None)[source]

Bases: object

A single chunk of text produced by a chunker.

Parameters:
content

The text content of the chunk.

Type:

str

chunk_index

Zero-based index of this chunk in the sequence.

Type:

int

start_char

Character offset where this chunk begins in the source.

Type:

int

end_char

Character offset where this chunk ends in the source.

Type:

int

metadata

Optional metadata attached to this chunk.

Type:

dict[str, Any] | None

content: str
chunk_index: int
start_char: int
end_char: int
metadata: dict[str, Any] | None = None
VectorStore

alias of ChromaVectorStore

class VectorStoreBase[source]

Bases: ABC

Abstract base class for vector store backends.

All backends must implement the core CRUD + search methods. The where parameter in asearch accepts a backend-neutral filter dict with string keys and values. Each backend translates this to its native filter format internally.

Example backend-neutral filters:

{"source": "doc1"}
{"chunking_strategy": "recursive", "source": "doc1"}

For compound filters, backends should support an $and key:

{"$and": [{"source": "doc1"}, {"chunking_strategy": "recursive"}]}
abstract property name: str

Backend identifier (e.g. "chroma", "pgvector").

abstractmethod async aadd(ids, texts, embeddings, metadatas)[source]

Add documents with pre-computed embeddings to the store.

Parameters:

  • ids (list[str]) – Unique identifiers for each document.

  • texts (list[str]) – Document text content.

  • embeddings (list[list[float]]) – Pre-computed embedding vectors.

  • metadatas (list[dict[str, Any]]) – Metadata dicts (must include source key).

Returns:

Number of documents added.

Return type:

int

abstractmethod async asearch(query_embedding, top_k=5, where=None)[source]

Search for documents by vector similarity.

Parameters:

  • query_embedding (list[float]) – Query vector.

  • top_k (int) – Maximum number of results.

  • where (dict[str, Any] | None) – Metadata filter dict (e.g. {"source": "doc1"}).

Returns:

Ranked search hits.

Return type:

list[SearchHit]

abstractmethod delete_by_source(source)[source]

Delete all chunks matching source.

Parameters:

source (str)

Return type:

None

abstractmethod delete_by_source_and_strategy(source, strategy)[source]

Delete chunks matching both source and chunking_strategy.

Parameters:
Return type:

None

abstractmethod count()[source]

Return the total number of stored chunks.

Return type:

int

abstractmethod clear()[source]

Delete all stored data and recreate the collection.

Return type:

None

abstractmethod list_sources()[source]

Return a sorted list of indexed source names.

Return type:

list[str]

abstractmethod get_all()[source]

Return all stored documents as dicts with id, content, metadata keys.

Return type:

list[dict[str, Any]]

abstractmethod needs_reindex(source, checksum, strategy='default')[source]

Check whether source needs re-indexing.

Parameters:

  • source (str) – Source identifier.

  • checksum (str) – Expected document checksum.

  • strategy (str) – Chunking strategy name.

Returns:

True if the source has changed or is not indexed.

Return type:

bool

Functions

chunk_text(text, strategy='recursive', chunk_size=1000, chunk_overlap=200, metadata=None, **kwargs)[source]

Convenience function to chunk text with a single call.

Parameters:

  • text (str) – Text to chunk.

  • strategy (str) – Chunking strategy name.

  • chunk_size (int) – Maximum chunk size.

  • chunk_overlap (int) – Overlap between chunks.

  • metadata (dict[str, Any] | None) – Metadata to attach to chunks.

  • **kwargs (Any) – Additional strategy-specific parameters.

Returns:

List of TextChunk objects.

Return type:

list[TextChunk]

format_hits(hits, max_chars=None, include_parent=True)[source]

Format search hits into a numbered, human-readable text string.

Each hit is rendered with its index, source, relevance score, content, and optionally a snippet of its parent context.

Parameters:

  • hits (list[SearchHit]) – List of SearchHit objects to format.

  • max_chars (int | None) – Optional character budget; hits beyond the budget are omitted.

  • include_parent (bool) – Whether to include parent-context snippets.

Returns:

A formatted string, or an empty string if hits is empty.

Return type:

str

fuse_search_results(result_lists, n_results=5, dedupe_by='id')[source]

Fuse results from multiple searches with deduplication.

Combines results from multiple query variations, removing duplicates while preserving relevance ordering.

Parameters:

  • result_lists (list[list[dict[str, Any]]]) – List of search result lists from different queries.

  • n_results (int) – Maximum number of final results to return.

  • dedupe_by (str) – Field to use for deduplication (default: “id”).

Returns:

Fused and deduplicated list of results.

Return type:

list[dict[str, Any]]

get_chunker(strategy='recursive', chunk_size=1000, chunk_overlap=200, metadata=None, **kwargs)[source]

Get a chunker instance based on strategy name.

Parameters:

  • strategy (str) – Chunking strategy name (character, recursive, markdown, code, hierarchical).

  • chunk_size (int) – Maximum chunk size.

  • chunk_overlap (int) – Overlap between chunks.

  • metadata (dict[str, Any] | None) – Default metadata for chunks.

  • **kwargs (Any) – Additional strategy-specific parameters.

Returns:

Configured chunker instance.

Raises:

ValueError – If strategy name is not recognized.

Return type:

ChunkerBase

get_reranker(reranker_type='none', **kwargs)[source]

Get a reranker by type name.

Parameters:

  • reranker_type (str) – Type of reranker (“cross_encoder”, “diversity”, “none”).

  • **kwargs (Any) – Additional arguments for the reranker.

Returns:

Configured reranker instance.

Return type:

RerankerBase

get_store(backend='chroma', **kwargs)[source]

Get a vector store instance by backend name.

Parameters:

  • backend (str) – Backend name ("chroma", "pgvector", "qdrant", "sqlite_vss").

  • **kwargs (Any) – Backend-specific constructor arguments.

Returns:

Configured vector store instance.

Raises:

  • ValueError – If backend name is not recognized.

  • ImportError – If the backend’s dependency is not installed.

Return type:

VectorStoreBase

is_store_available(name)[source]

Check if a backend’s dependencies are installed.

Parameters:

name (str)

Return type:

bool

list_available_stores()[source]

List backend names whose dependencies are installed.

Return type:

list[str]

list_chunkers()[source]

List available chunking strategies.

Returns:

List of strategy names.

Return type:

list[str]

list_stores()[source]

List all known backend names (regardless of availability).

Return type:

list[str]

litellm_generate_fn(model, api_key=None, temperature=0.5, max_tokens=1024, **kwargs)[source]
Parameters:
Return type:

Callable[[str], GenerationResult]

reciprocal_rank_fusion(result_lists, k=60, weights=None)[source]

Combine multiple result lists using reciprocal rank fusion.

Parameters:

  • result_lists (list[list[dict[str, Any]]]) – List of result lists to fuse.

  • k (int) – RRF constant.

  • weights (list[float] | None) – Optional weights for each result list.

Returns:

Fused and sorted results.

Return type:

list[dict[str, Any]]