Vector stores

Configuration

Retriever store configuration models

File: klea_utils/stores/config.py

Copyright 2026 Ankur Sinha Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>

class klea_utils.stores.config.BM25StoreInfo(*, name: str, path: str, default_k: int | None = None, k_max: int | None = None, k_inc: int | None = None, loaded_object: Any | None = None)

Bases: StoreInfo

Information about a single BM25 store.

path points to the pickled document corpus that the BM25RetrieverManager loads to build its keyword index.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class klea_utils.stores.config.PerDomainConfig(*, vector_stores: list[VectorStoreInfo] = [], bm25_stores: list[BM25StoreInfo] = [])

Bases: BaseModel

Configuration for a single domain.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class klea_utils.stores.config.RetrieverConfig(*, domains: dict[str, PerDomainConfig])

Bases: BaseModel

Top-level retriever configuration.

Holds the per-domain store configuration for all retriever managers (vector stores and BM25 stores).

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class klea_utils.stores.config.StoreInfo(*, name: str, path: str, default_k: int | None = None, k_max: int | None = None, k_inc: int | None = None, loaded_object: Any | None = None)

Bases: BaseModel

Information about a single store used by a retriever manager.

default_k, k_max, and k_inc configure retrieval depth per store. When left None they fall back to the global values set on the retriever manager, so stores that do not need tuning inherit the graph-wide defaults.

loaded_object holds the lazily-instantiated retriever object for the store (e.g. a LangChain VectorStore or BM25Retriever).

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class klea_utils.stores.config.VectorStoreInfo(*, name: str, path: str, default_k: int | None = None, k_max: int | None = None, k_inc: int | None = None, loaded_object: Any | None = None)

Bases: StoreInfo

Information about a single vector store.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

Ingestion

class klea_utils.stores.ingestion.StoresBuilder(embedding_model: str, logger: Logger, max_tokens: int = 450, merge_peers: bool = True, tokenizer_model: str = 'BAAI/bge-m3', do_ocr: bool = True, embed_batch_size: int = 256)

Bases: object

Build stores from a directory of source documents.

Uses Docling for document conversion and token-aware chunking, then embeds chunks and writes them to a vector store backend. Optionally also writes the combined chunked corpus for BM25 retrieval.

DEFAULT_EMBED_BATCH_SIZE = 256

Number of chunks embedded per add_documents call in store_all(). Batching gives the embedding phase (which can take minutes for large corpora) a progress signal between calls; embedding backends like Ollama send all texts in a single request otherwise. The value is mostly a progress-granularity knob, not a throughput one.

build(source_dir: str, store_uri: str, collection_name: str, force: bool = False, metadata_map_path: str | None = None, bm25_path: str | None = None) → None

Full pipeline: chunk documents and write them to a vector store.

Convenience wrapper around chunk_all() + store_all().

Parameters:

• source_dir – Path to a directory containing source documents

• store_uri – Vector store URI (e.g. chroma:/path)

• collection_name – Collection name for the store

• force – Re-process all files even if unchanged

• metadata_map_path – Optional path to a metadata map JSON file

• bm25_path – Optional path to write the combined BM25 corpus to

chunk_all(source_path: Path, metadata_map: dict[str, dict[str, Any]] | None = None, force: bool = False) → tuple[list[tuple[str, list[langchain_core.documents.Document], Path]], dict[str, dict[str, Any]]]

Convert, chunk, cache, and enrich metadata for all files.

Skips converting files whose cache entry exists (unless force is True). Always caches newly-converted chunks. Heading chains are collected per file for template generation. The per-file DEFAULT template entry is pre-filled with the automatically-extracted bibliographic metadata (see extract_metadata()).

Parameters:

• source_path – Resolved source directory path

• metadata_map – Metadata map for heading-based enrichment, or None

• force – Re-process all files even if cached

Returns:

(results, file_headings) where results is a list of (file_hash, docs, file_path) tuples and file_headings is a {file_name: {"DEFAULT": {extracted metadata}, "heading > heading": {}, ...}} dict

store_all(results: list[tuple[str, list[langchain_core.documents.Document], Path]], store_uri: str, collection_name: str, force: bool = False, bm25_path: str | None = None) → None

Write chunked documents to a vector store.

Initialises the embedding model on first call if not already done. Skips files whose hash is already present in the store (unless force is True). Optionally also writes the combined document corpus for BM25 retrieval.

Parameters:

• results – List of (file_hash, docs, file_path) tuples from chunk_all()

• store_uri – Vector store URI

• collection_name – Collection name for the store

• force – Re-store all files even if already indexed

• bm25_path – Optional path to write the combined BM25 corpus to

write_bm25_store(results: list[tuple[str, list[langchain_core.documents.Document], Path]], bm25_path: str) → None

Write the combined chunked documents to a BM25 corpus.

Flattens the per-file chunked documents from chunk_all() into a single list and pickles it to bm25_path. This file is the BM25 store: a BM25RetrieverManager loads it at runtime to build its keyword index. It is independent of the per-file .klea-cache, so the cache can be removed once the corpus has been written.

The corpus holds the same chunk units (and metadata) that are stored in the vector store, so BM25 and vector retrieval return consistent results.

Parameters:

• results – List of (file_hash, docs, file_path) tuples from chunk_all()

• bm25_path – Path to write the combined corpus pickle to

write_heading_template(file_headings: dict[str, dict[str, Any]], source_dir: Path) → None

Write a metadata-map template JSON file organised per source file.

Each file gets a "DEFAULT" placeholder and one entry per unique heading chain found in that file. The user fills in the {} with their metadata key-value pairs.

Refuses to write when file_headings is empty (no files were chunked): an existing template is preserved rather than clobbered with an empty one.

Parameters:

• file_headings – {file_name: {"DEFAULT": {}, "heading > heading": {}, ...}, ...} from chunk_all()

• source_dir – Resolved source directory path (template is written alongside it)

Retrieval

class klea_utils.stores.retrieval.base.BaseKleaRetriever(config: RetrieverConfig, logger: Logger, default_k: int = 5, k_max: int = 10, k_inc: int = 1)

Bases: ABC

Base class for domain-aware retriever managers.

Holds the machinery common to all retrievers: lazy per-domain store loading, per-store retrieval depth (k) tracking with graph-wide fallbacks, and the retrieval contract retrieve(domain, query) -> list[tuple[Document, float]].

Subclasses must implement:

• _stores_of(): the list of stores configured for a domain

• _instantiate_store(): build the underlying retriever object for a store

• _retrieve_from_store(): run a single store against a query

Subclasses should set source_label to a human-readable name for the retriever type (e.g. "vector store", "BM25"), used to label the original per-source scores preserved during fusion.

property domains: list[str]

Get a list of all configured domains.

inc_k() → bool

Increase k for all loaded stores by their per-store increment.

Each store’s k is capped by its own k_max, so stores with a smaller cap stop being incremented sooner. Stores that are not yet loaded keep their default k until they are loaded.

Returns:

True if at least one store’s k was increased

load(domain_name: str) → None

Load stores for a domain (lazy loading).

Parameters:

domain_name – Name of the domain to load stores for

load_all_stores() → None

Load all stores for all domains.

reset_k() → None

Reset k for all loaded stores to their per-store default value.

retrieve(domain_name: str, query: str) → list[tuple[langchain_core.documents.Document, float]]

Retrieve documents from all stores for a domain.

Parameters:

• domain_name – Name of the domain to search in

• query – User query string

Returns:

List of (document, relevance_score) tuples

setup() → None

Hook for subclasses to initialise shared resources.

Called once by the orchestrator before retrieval. Subclasses that need no shared setup leave this as a no-op.

source_label: str = 'retriever'

Human-readable name for this retriever type, used to label scores.

class klea_utils.stores.retrieval.vs.VSRetriever(config: RetrieverConfig, logger: Logger, embedding_model: str, default_k: int = 5, k_max: int = 10, k_inc: int = 1)

Bases: BaseKleaRetriever

Manages domain-specific vector stores.

Loads vector stores on demand per domain and provides similarity search retrieval across multiple stores within a domain.

Store paths use a URI-style scheme prefix to identify the backend:

• chroma:/path/to/dir — ChromaDB (persistent, local disk)

• qdrant:http://host:port — Qdrant (remote HTTP)

• pgvector:postgresql://host/db — PGVector (PostgreSQL)

setup() → None

Initialise embedding model.

source_label: str = 'vector store'

Human-readable name for this retriever type, used to label scores.

class klea_utils.stores.retrieval.bm25.BM25RetrieverManager(config: RetrieverConfig, logger: Logger, default_k: int = 5, k_max: int = 10, k_inc: int = 1)

Bases: BaseKleaRetriever

Manages domain-specific BM25 keyword stores.

Each BM25 store is a pickled corpus of chunked documents written by write_bm25_store. Stores are loaded lazily per domain: the corpus is unpickled and used to build a langchain_community.retrievers.BM25Retriever, which is queried with BM25 keyword scoring.

A store whose corpus file is missing is skipped with a warning, so a misconfigured domain degrades gracefully instead of failing retrieval.

Scalability note: this is a pure-Python in-memory index (rank_bm25). Building and querying stay fast to well over ~100k chunks, but memory grows roughly with the total number of unique terms (one Python dict entry per term per chunk, ~100-150 bytes each), so a single collection becomes heavy in the ~50-100k chunk range. That is far beyond current corpora, but if large platformed deployments are ever planned, consider a proper keyword backend (e.g. Elasticsearch, Qdrant sparse vectors, or Postgres full-text search) instead.

source_label: str = 'BM25'

Human-readable name for this retriever type, used to label scores.

Utilities

Vector store utilities

File: klea_utils/stores/utils.py

Copyright 2026 Ankur Sinha Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>

klea_utils.stores.utils.CHROMA_HNSW_SPACE = 'cosine'

HNSW distance space used for Chroma collections created by Klea. Cosine makes the vector-store relevance scores true cosine similarities (1 - cosine_distance), so a score_threshold reads as a minimum cosine similarity.

klea_utils.stores.utils.RRF_K = 60

Rank offset for Reciprocal Rank Fusion. A document at rank r within a source’s result list contributes 1 / (RRF_K + r) to its fused score.

klea_utils.stores.utils.SOURCE_SCORES_KEY = '_source_scores'

Metadata key holding each document’s original per-source scores (e.g. {"vector store": 0.87, "BM25": 3.21}), set by rrf_merge().

klea_utils.stores.utils.format_source_scores(doc: langchain_core.documents.Document, precision: int = 2) → str | None

Format a document’s original per-source scores for display.

Parameters:

• doc – Document with SOURCE_SCORES_KEY metadata

• precision – Number of decimal places per score

Returns:

Joined string like "vector store 0.87, BM25 3.21", or None if the document has no per-source scores

klea_utils.stores.utils.instantiate_vector_store(path: str, name: str, embeddings, logger: Logger, create: bool = False)

Instantiate a vector store based on the URI scheme in path.

Expected format: "scheme:location".

If create is True, the store is created if it does not exist (relevant for ChromaDB which requires a local directory). For Qdrant and PGVector the flag is a no-op — collections are created on first write.

For ChromaDB, location must point at the store folder. Chroma always stores its database as <folder>/chroma.sqlite3 and the filename is not configurable, so a path pointing at an existing file (even the chroma.sqlite3 itself) is rejected. The collection name selects which collection within the store file is addressed: a single ChromaDB store file can hold multiple collections, so reusing an existing folder with a new collection name creates a new collection in it.

New Chroma collections are created with the CHROMA_HNSW_SPACE HNSW distance space (cosine). The configuration is only applied at collection creation; loading an existing collection keeps its own distance space.

Parameters:

• path – URI-style string with scheme prefix (e.g. "chroma:/path/to/dir", "qdrant:http://localhost:6333", "pgvector:postgresql://localhost/db")

• name – Collection name for the vector store

• embeddings – Embedding function to use

• logger – Logger instance

• create – If True, allow creating a new store

Returns:

Instantiated LangChain VectorStore

Raises:

• ValueError – If the scheme is missing or unknown

• FileNotFoundError – If create is False and a local ChromaDB store does not exist

klea_utils.stores.utils.normalize_text(text: str) → str

Normalise free text for consistent indexing and retrieval.

Document conversion (e.g. Docling’s PDF extraction) embeds typographic artifacts that hurt search: soft hyphens (\u00ad) split words mid-token, no-break / zero-width characters distort embeddings and BM25 keyword matching, and ligatures / full-width forms / superscripts / typographic spaces tokenise differently from their plain equivalents. This strips or maps them so that indexed chunks and retrieval queries share the same plain-text form.

The final pass uses NFKC compatibility composition, which (unlike NFC) also folds ligatures (\ufb01 -> “fi”), full-width forms (\uff21 -> “A”), superscripts (\u00b2 -> “2”), typographic spaces (en/em/thin/ideographic), and the non-breaking hyphen (\u2011 -> \u2010). Typographic em/en dashes are kept unchanged.

Parameters:

text – Raw text, possibly containing typographic artifacts

Returns:

Normalised plain text

klea_utils.stores.utils.rrf_merge(result_sets: list[tuple[str, list[tuple[langchain_core.documents.Document, float]]]], num_refs_max: int) → list[tuple[langchain_core.documents.Document, float]]

Fuse per-source retrieval results with Reciprocal Rank Fusion.

Scores from different retrievers (e.g. cosine similarity vs BM25) are not comparable, so each document is scored purely by its rank within each source’s result list. The original per-source scores are preserved in each document’s SOURCE_SCORES_KEY metadata for display.

Parameters:

• result_sets – List of (source_label, results) pairs, where each results is a list of (document, score) tuples already ranked by its source

• num_refs_max – Maximum number of documents to return

Returns:

Documents ordered by RRF score, deduplicated by content, capped at num_refs_max

klea_utils.stores.utils.serialize_vs_retrieval(reference_material: dict[str, list[tuple[langchain_core.documents.Document, float]]]) → str

Serialize vector store retrieval results into text for use in prompt context.

Documents are sorted by relevance score within each group. Uses Docling HybridChunker metadata format:

• headings: list of heading hierarchy (most specific last)

• file_name: source filename

• _source_scores: optional per-retriever scores (from the RRF merge)

• Optional custom keys from the --metadata-map (e.g., url)

Parameters:

reference_material – Dict mapping query/domain to list of (doc, score) tuples

Returns:

Formatted string representation of references