rag.indexing

Classes

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

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)

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)

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')

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)

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)

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()

Clear all documents from the index.

Return type:

None

delete_by_metadata(key, value)

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()

Get the number of documents in the index.

Return type:

int

get_stats()

Get index statistics.

Return type:

dict[str, Any]

class ChunkDeduplicator(mode='exact', similarity_threshold=0.95)

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)

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)

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)

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()

Clear seen hashes and embeddings.

Return type:

None

get_stats()

Get deduplicator statistics.

Returns:

Dict with mode, threshold, and counts.

Return type:

dict[str, Any]

class ContextualEmbeddings(context_template=None, max_context_length=200)

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)

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)

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 HierarchicalIndex(include_parent_context=True)

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)

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)

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)

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)

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()

Get all parent chunks (level 0).

Returns:

List of parent chunk data dictionaries.

Return type:

list[dict[str, Any]]

get_child_chunks()

Get all child chunks (level > 0).

Returns:

List of child chunk data dictionaries.

Return type:

list[dict[str, Any]]

get_child_embeddings()

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)

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)

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)

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()

Clear all chunks from the index.

Return type:

None

count()

Get total number of chunks.

Return type:

int

count_parents()

Get number of parent chunks.

Return type:

int

count_children()

Get number of child chunks.

Return type:

int

get_stats()

Get index statistics.

Return type:

dict[str, Any]

class LateChunkingEmbeddings(embedding_model=None)

Bases: object

Late chunking strategy for token-level embeddings.

Instead of embedding chunks separately, embeds the full document and extracts chunk-level representations from the token embeddings.

This is a placeholder for ColBERT-style late interaction embeddings. Full implementation would require a model that outputs token-level embeddings and late interaction scoring.

Parameters:

embedding_model (Any | None) – The underlying embedding model.

embed_document_with_tokens(document_id, content, chunk_boundaries)

Embed document and extract chunk representations.

Note: This is a simplified implementation. A full ColBERT-style implementation would use a model that outputs token embeddings.

Parameters:

• document_id (str) – Document identifier.

• content (str) – Full document content.

• chunk_boundaries (list[tuple[int, int]]) – List of (start, end) tuples for chunks.

Returns:

List of chunk embeddings.

Return type:

list[list[float]]