Vector stores¶

Configuration¶

Retriever store configuration models

File: klea_utils/stores/config.py

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)[source]¶

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.FilterFieldInfo(*, name: str, description: str, value_type: Literal['string', 'int', 'float', 'list'] = 'string')[source]¶

Bases: BaseModel

Configuration for a single retrievable metadata filter field.

A deployment declares, per domain, the metadata fields the retrieval query generator may filter on. Each entry describes one field: its name (the metadata key stored on the documents), its semantics for the LLM, and the operand type it accepts.

value_type controls how a bare operand from the LLM is mapped to the filter DSL (see klea_utils.stores.filters.normalize_config_filters()):

"string" / "int" / "float" — scalar fields. A bare value becomes $eq; a list of values becomes $in.
"list" — element-membership fields (e.g. tags). A bare value becomes $contains; several values combine with $and (every value must be present).

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] = [], filter_fields: list[FilterFieldInfo] = [])[source]¶

Bases: BaseModel

Configuration for a single domain.

filter_fields: list[FilterFieldInfo]¶: Retrieval filter fields the query generator may use for this 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])[source]¶

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)[source]¶

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)[source]¶

Bases: StoreInfo

Information about a single vector store.

model_config = {}¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

Stored metadata schema¶

Stored-metadata schema for vector store ingestion.

File: klea_utils/stores/metadata.py

klea_utils.stores.metadata.ALWAYS_STORED_METADATA_KEYS = frozenset({'authors', 'doi', 'file_hash', 'file_name', 'headings', 'journal', 'keywords', 'title', 'year'})¶: Metadata keys that are always stored in the vector store, together with the bibliographic fields produced by the extraction cascade (title, authors, keywords, year, journal, doi). Any url* key (url, url_1, url_doi, …) is also always stored. This whitelist documents the guaranteed stored schema; presence of the cascade fields is determined by the metadata map (whose researcher-curated keys pass through unmodified). See klea_utils.stores.ingestion._apply_store_metadata_policy().

klea_utils.stores.metadata.MACHINE_SET_METADATA_KEYS = frozenset({'file_hash', 'file_name', 'headings'})¶: Metadata keys set on every chunk by the ingestion pipeline itself, rather than supplied by the metadata map. They are always stored.

klea_utils.stores.metadata.PERSON_NAME_FILTER_FIELDS = frozenset({'authors'})¶: Metadata fields whose values are lists of human full names. At store time such fields are expanded with per-word variants (surname, every whitespace token, and lowercase forms) so an exact-membership retrieval filter matches the partial name a user is likely to use – e.g. “find papers by Sinha” matches an author stored as “Ankur Sinha”. This is a deliberate, internally-fixed set (no user-facing configuration): the authors bibliographic field is the canonical person-name list today. Other name-bearing fields (titles, venues, repository names, usernames) are not expanded – users refer to those by exact value, or by partial text that the content retrievers (vector/BM25) already match. See klea_utils.stores.utils.expand_person_names() and klea_utils.stores.utils.display_person_names().

klea_utils.stores.metadata.SHARED_DOC_METADATA_KEYS = frozenset({'authors', 'doi', 'journal', 'keywords', 'title', 'year'})¶: Bibliographic fields that every chunk of a source file is expected to share (they come from the file’s DEFAULT metadata map entry). When serializing reference material, these are emitted once per source file rather than repeated on every chunk.

klea_utils.stores.metadata.STORE_DROPPED_METADATA_KEYS = frozenset({'source_path', 'source_type', 'source_url'})¶: Metadata keys that are never stored. Provenance keys from the biblio cascade; keys starting with _ (e.g. _metadata_complete, _sources, _source_scores) are also always dropped. These guide the researcher reviewing metadata-map.template.json but carry no meaning in a store. See klea_utils.stores.ingestion._apply_store_metadata_policy().

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, store_dir: Path | None = None)[source]¶

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, worker_mem_limit: int | None = 4294967296, worker_batch_size: int = 200) → None[source]¶

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

Composes the two memory-bounded paths: chunk_all() runs in worker-isolated, cache-only mode (uncached files are converted in short-lived subprocesses), then _load_and_fold_results() streams the cached chunks into store_all(). No phase holds the whole corpus in memory, so very large corpora stay bounded.

When no metadata_map_path is given, the map is generated in the chunk phase from the extracted bibliographic metadata (and written to metadata-map.template.json, exactly as write_heading_template() does) and consumed in the store phase – so build works without a prior chunk, at the cost of no review step. For the review-driven flow (chunk, edit the template, store), pass the map explicitly or let store auto-fall back to the template.

Files whose conversion failed have no cache entry and are skipped with an error; everything else is stored.

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
worker_mem_limit – Maximum RSS per conversion worker in bytes; None (default) keeps the chunk phase in-process
worker_batch_size – Maximum files handed to a single worker

chunk_all(source_path: Path, force: bool = False, worker_mem_limit: int | None = 4294967296, worker_batch_size: int = 200) → dict[str, dict[str, Any]][source]¶

Convert, chunk, and cache all files in memory-bounded workers.

Uncached files are converted in short-lived subprocess workers (klea_utils.stores.chunk_worker) so Docling’s per-conversion memory leak is reclaimed on worker exit; cache hits are handled in-process. The chunks live in the on-disk cache – they are never accumulated in this process – so the run stays memory-bounded on corpora of any size. Cache entries whose source file no longer exists are pruned at the end.

Callers that need the chunked documents read them back from the cache with _load_and_fold_results() (the store command and build()), which streams them one file at a time.

Parameters:

source_path – Resolved source directory path
force – Re-process all files even if cached
worker_mem_limit – Maximum RSS per conversion worker in bytes; the limit includes the ~1-1.5 GiB Docling’s models occupy at worker startup, so headroom for its per-conversion growth is roughly the limit minus that. None disables the cap (workers are then bounded only by worker_batch_size).
worker_batch_size – Maximum files handed to any single conversion worker before it is restarted. A worker stops at whichever comes first: its memory cap or this batch size.

Returns:

file_headings – a {file_name: {"DEFAULT": {extracted metadata}, "heading > heading": {}, ...}} dict for template generation, pre-filled with the automatically-extracted bibliographic metadata (see extract_metadata())

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

Write chunked documents to a vector store.

Incremental by default: a store manifest (<source_dir>/.klea-cache/<collection>.manifest.json) records which files are in the collection and how many chunks each has, so unchanged files are skipped, changed files have their old chunk IDs deleted and are re-added, and new files are added. Files absent from the source directory are left untouched (never pruned).

With force the whole collection is dropped and rebuilt from scratch (see klea_utils.stores.utils.drop_collection()), then the manifest is rewritten. This is the portable way to update a collection, since documents within a collection cannot be updated in place across all backends.

Chunk IDs are deterministic (<file_name>:<chunk_index>) so deletion by ID works on every backend.

results may be any iterable, including the lazy generator from _load_and_fold_results(); each file’s chunks are released once they are stored, keeping memory bounded per file. When a BM25 corpus is requested it is written inline from the same chunks during the store loop, one pickled batch per embed_batch_size chunk (read back by looping pickle.load until EOFError).

Parameters:

results – Iterable of (file_hash, docs, file_path) tuples from chunk_all() or _load_and_fold_results()
store_uri – Vector store URI
collection_name – Collection name for the store
source_dir – Resolved source directory (for the manifest)
force – Drop the collection and re-store everything
bm25_path – Optional path to write the combined BM25 corpus to

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

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.

The template is written into the source directory’s cache folder (<source_dir>/.klea-cache/metadata-map.template.json), the same place the chunk cache and doi-cache.json live. To review it, copy it out (e.g. to metadata-map.json), edit, and pass the copy to klea-stores-create store --metadata-map <path>.

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 into its cache folder)

Metadata map linting¶

LLM-free linting of a metadata map.

File: klea_utils/stores/map_lint.py

klea_utils.stores.map_lint.CORE_FIELDS = ('authors', 'doi', 'journal', 'keywords', 'title', 'year')¶: Core bibliographic fields a complete DEFAULT entry should carry – the shared document fields every chunk inherits. Derived (not duplicated) so a change to the stored schema is picked up here.

klea_utils.stores.map_lint.URL_WARN_THRESHOLD = 5¶: DEFAULT entries with more than this many url* keys almost certainly picked up reference URLs during extraction (a paper typically carries a DOI page, a journal page, and a couple of extras).

klea_utils.stores.map_lint.YEAR_MAX = 2028¶

the current year plus a small margin for in-press/early-access items (a hardcoded far-future cap would go stale as the current year advances).

Type:: Latest plausible publication year

klea_utils.stores.map_lint.YEAR_MIN = 1800¶: Earliest plausible publication year. Anything before this is flagged.

klea_utils.stores.map_lint.format_metadata_lint_report(report: dict[str, Any]) → str[source]¶: Render a lint_metadata_map() report as compact text.

klea_utils.stores.map_lint.lint_file_metadata(file_name: str, entry: dict[str, Any]) → list[str][source]¶

Return human-readable issues for one metadata-map file entry.

entry is the per-file dict from a metadata map: a "DEFAULT" metadata dict plus one dict per heading chain (empty {} placeholders in a generated template). A flat metadata dict (the old heading-keyed format) is reported as a single structural issue instead of a misleading list of missing fields.

Parameters:

file_name – Source filename (used for the year-vs-stem check)
entry – Per-file metadata-map entry

Returns:

Sorted list of issue strings; empty when the entry is clean

klea_utils.stores.map_lint.lint_metadata_map(data: dict[str, dict[str, Any]], source_files: Iterable[str] | None = None) → dict[str, Any][source]¶

Lint a whole metadata map and return a structured report.

When source_files is given (the basenames of the files the store will ingest, e.g. from find_source_files()), the top-level keys are checked against it: a source file with no entry is a fatal error (the store step raises ValueError), while a key that is not a source file is a stale or heading-keyed leftover that store simply ignores. With None (library callers, the auto-print after chunk) these key checks are skipped.

Parameters:

data – Parsed metadata-map JSON ({file_name: entry})
source_files – Optional basenames of the source files the store will ingest; when given, top-level keys are validated against them

Returns:

dict with files (total), complete (count of _metadata_complete DEFAULTs), issues ({file_name: [issue, ...]} for files with at least one issue), placeholders ({file_name: int} count of empty heading placeholders per file), missing_keys (source files with no map entry – store will fail), and unknown_keys (map keys that are not source files – stale/heading-keyed)

Retrieval¶

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

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.

can_inc_k() → bool[source]¶

Return whether any loaded store still has room to grow k.

Non-mutating counterpart of inc_k(): reports whether an inc_k() call would increase at least one store’s k, without changing any k values. Routers use this to decide between retrieving more information and re-querying, so the actual increment happens once, at the point of retrieval.

Returns:: True if at least one loaded store’s k is below its cap

property domains: list[str]¶: Get a list of all configured domains.

inc_k() → bool[source]¶

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[source]¶

Load stores for a domain (lazy loading).

Parameters:: domain_name – Name of the domain to load stores for

load_all_stores() → None[source]¶: Load all stores for all domains.

reset_k() → None[source]¶: Reset k for all loaded stores to their per-store default value.

retrieve(domain_name: str, query: str, metadata_filter: dict[str, Any] | None = None) → list[tuple[langchain_core.documents.Document, float]][source]¶

Retrieve documents from all stores for a domain.

Parameters:

domain_name – Name of the domain to search in
query – User query string
metadata_filter – Optional metadata filter in the DSL (see klea_utils.stores.filters.validate_metadata_filter()). Applied natively by stores that support a backend filter and post-filtered for stores that do not (BM25)

Returns:

List of (document, relevance_score) tuples

setup() → None[source]¶

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)[source]¶

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[source]¶: 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)[source]¶

Bases: BaseKleaRetriever

Manages domain-specific BM25 keyword stores.

Each BM25 store is a pickled corpus of chunked documents written alongside the vector store by store_all() (one pickled list per batch; read back by looping pickle.load until EOFError). Stores are loaded lazily per domain: the corpus is unpickled and used to build a klea_utils.stores.langchain_bm25, 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.

Metadata filters¶

Metadata filter translation for store backends

File: klea_utils/stores/filters.py

klea_utils.stores.filters.FIELD_OPERATORS = frozenset({'$contains', '$eq', '$gt', '$gte', '$in', '$lt', '$lte', '$ne', '$nin'})¶: Operators valid inside a single field clause of the normalized filter DSL. $contains is the element-membership operator for list-valued fields (e.g. authors, keywords).

klea_utils.stores.filters.LOGICAL_OPERATORS = frozenset({'$and', '$or'})¶: Logical combinators valid at any level of the normalized filter DSL.

klea_utils.stores.filters.filter_docs_by_metadata(docs: list[langchain_core.documents.Document], f: dict[str, Any]) → list[langchain_core.documents.Document][source]¶

Return the subset of docs whose metadata matches the filter.

Python-side matcher for backends without native filter support: the BM25 store post-filters its results through this (the rank_bm25 index has no filter). Documents without the filter’s metadata field never match.

Example:

docs = [Document(page_content="a", metadata={"authors": ["Magee"]}),
        Document(page_content="b", metadata={"authors": ["Jones"]})]
filter_docs_by_metadata(docs, {"authors": {"$contains": "Magee"}})
-> [Document(page_content="a", ...)]

Parameters:

docs – Documents to filter, each carrying metadata
f – Metadata filter in the DSL (see validate_metadata_filter())

Returns:

Documents whose metadata satisfies the filter

Raises:

ValueError – When the filter is not well-formed

klea_utils.stores.filters.normalize_config_filters(filters: dict[str, Any], allowed_fields: list[FilterFieldInfo]) → list[dict[str, Any]][source]¶

Validate configured-domain filters into canonical DSL clauses.

filters maps a metadata field name (from a deployment’s filter_fields configuration) to an operand produced by the retrieval query generator: a bare scalar, a list of scalars, or an operator expression dict ({op: value}). Each configured field’s value_type decides how a bare operand is interpreted:

scalar fields (string/int/float): a bare value is $eq; a list of values is $in.
list fields (e.g. tags): a bare value requires element membership ($contains); several values combine with $and (every value must be present), mirroring the bibliographic authors/keywords handling in RetrievalQueryOutput.to_metadata_filter().

An operator expression dict is validated and normalized through validate_metadata_filter(); an unsupported operator or malformed operand raises ValueError.

Field names not declared in allowed_fields are ignored with a warning and never reach a backend (the generator must not be able to emit a key the deployment did not configure). An empty operand list is likewise ignored.

The result is a list of single-clause DSL dicts, each directly consumable by validate_metadata_filter() (and hence by every backend translator and the in-memory matcher). An empty input returns [].

Example:

fields = [
    FilterFieldInfo(name="repository_type", description="...",
                    value_type="string"),
    FilterFieldInfo(name="tags", description="...", value_type="list"),
]
normalize_config_filters(
    {"repository_type": ["github", "dandi"], "tags": "moose"},
    fields,
)
-> [
    {"repository_type": {"$in": ["github", "dandi"]}},
    {"tags": {"$contains": "moose"}},
]

Parameters:

filters – Field-name -> operand mapping from the query generator
allowed_fields – Configured filter fields for the domain

Returns:

Canonical DSL single-clause dicts

Raises:

ValueError – When an operator expression uses an unsupported operator or malformed operand

klea_utils.stores.filters.restrict_metadata_filter(metadata_filter: dict[str, Any] | None, allowed_field_names: set[str]) → dict[str, Any] | None[source]¶

Restrict a metadata filter to clauses on the allowed fields.

metadata_filter is a combined DSL filter (as produced by RetrievalQueryOutput.to_metadata_filter()): a single field clause, a top-level $and/$or of clauses, or None. This keeps only the clauses whose referenced metadata field(s) are in allowed_field_names and recombines the survivors (a single clause returned as-is, several wrapped in $and), so a domain that declares only e.g. journal never has a username clause applied to its retrievers. A filter with nothing left returns None.

A top-level $or is kept whole only when every field it references is allowed (partially splitting an $or would change its semantics); a top-level $and is decomposed and its sub-clauses filtered independently, so the multi-value $and clauses emitted by normalize_config_filters() are kept or dropped as a unit.

Example:

restrict_metadata_filter(
    {"$and": [{"journal": {"$eq": "nature"}},
              {"username": {"$eq": "padraig"}}]},
    {"journal"},
)
-> {"journal": {"$eq": "nature"}}

Parameters:

metadata_filter – Combined metadata filter, or None
allowed_field_names – Metadata field names the caller accepts

Returns:

The restricted filter, or None

klea_utils.stores.filters.to_chroma_filter(f: dict[str, Any]) → dict[str, Any][source]¶

Translate a filter to a Chroma where dict.

The canonical normalized form is Chroma’s native where syntax (single-operator field clauses, $and/$or combinators, and the $contains array-membership operator), so this validates the filter and returns the normalized form unchanged.

Example:

to_chroma_filter({"authors": {"$contains": "Magee"}})
-> {"authors": {"$contains": "Magee"}}

Parameters:: f – Metadata filter in the DSL (see validate_metadata_filter())
Returns:: Chroma where dict, passable as filter= to a langchain_chroma store’s similarity search
Raises:: ValueError – When the filter is not well-formed

klea_utils.stores.filters.to_pgvector_filter(f: dict[str, Any]) → dict[str, Any][source]¶

Translate a filter to a langchain_postgres filter dict.

langchain_postgres accepts the canonical normalized form directly. The one gap is $contains: the backend has no array containment operator for metadata fields, so it is approximated with $like over the serialized JSON array text (e.g. a substring match against ["Magee","Smith"]). This is a documented approximation for the Postgres backend only.

Example:

to_pgvector_filter({"authors": {"$contains": "Magee"}})
-> {"authors": {"$like": "%Magee%"}}

Parameters:: f – Metadata filter in the DSL (see validate_metadata_filter())
Returns:: langchain_postgres filter dict, passable as filter= to a langchain_postgres store’s similarity search
Raises:: ValueError – When the filter is not well-formed

klea_utils.stores.filters.to_qdrant_filter(f: dict[str, Any]) → Any[source]¶

Translate a filter to a Qdrant models.Filter object.

Scalar equality and $contains become MatchValue (an array element match), $in/$nin become MatchAny/MatchExcept, and range operators become a Range condition. Compound clauses are grouped as nested Filter objects.

Example:

to_qdrant_filter({"year": {"$gte": 2020, "$lte": 2025}})
-> Filter(must=[Filter(must=[FieldCondition(key='year',
                                            range=Range(gte=2020.0))]),
                Filter(must=[FieldCondition(key='year',
                                            range=Range(lte=2025.0))])])

Parameters:: f – Metadata filter in the DSL (see validate_metadata_filter())
Returns:: Qdrant models.Filter object, passable as filter= to a langchain_qdrant store’s similarity search
Raises:: ValueError – When the filter is not well-formed

klea_utils.stores.filters.translate_metadata_filter(path: str, f: dict[str, Any]) → Any[source]¶

Translate a filter for the backend named by a store path.

Dispatch helper used by the retrievers: reads the URI scheme from a store path (e.g. chroma:/data/store) and returns the matching backend-native filter. filter_docs_by_metadata is not returned here – callers that need the in-memory matcher (the BM25 store) call it directly.

Example:

translate_metadata_filter("chroma:/data/store",
                          {"authors": {"$contains": "Magee"}})
-> {"authors": {"$contains": "Magee"}}

Parameters:

path – Store URI (scheme:location)
f – Metadata filter in the DSL (see validate_metadata_filter())

Returns:

Backend-native filter for the store’s scheme

Raises:

ValueError – When the scheme is missing or unknown

klea_utils.stores.filters.validate_metadata_filter(f: dict[str, Any]) → dict[str, Any][source]¶

Validate a filter and normalize it to the canonical form.

Filters use a small backend-agnostic DSL: a dict mapping a metadata field name to either a bare value (implicit $eq) or an operator expression {op: value}, combined with the $and/$or combinators. Supported operators are $eq, $ne, $gt, $gte, $lt, $lte, $in, $nin and $contains (the last matches a list-valued field, e.g. authors, that contains the value).

The canonical form is the one every backend translator accepts:

every field value is an operator expression with exactly one operator (bare values are wrapped as {"$eq": value})
several constraints (multiple top-level fields, or several operators on one field) are combined with $and
the top level has exactly one key (a field, $and or $or)

Chroma and langchain_postgres reject multi-operator field expressions and multi-key top levels, so this normalization is what makes a single filter usable across all backends.

Example:

validate_metadata_filter({"year": {"$gte": 2020, "$lte": 2025}})
-> {"$and": [{"year": {"$gte": 2020}}, {"year": {"$lte": 2025}}]}

Parameters:: f – Metadata filter in the DSL described above
Returns:: Canonical normalized filter
Raises:: ValueError – When the filter is not well-formed (empty, unknown operator, malformed operand)

BM25 index¶

class klea_utils.stores.langchain_bm25.BM25Retriever(*args: Any, **kwargs: Any)[source]¶

Bases: BaseRetriever

BM25 retriever using rank_bm25 for keyword-based document scoring.

This retriever processes a collection of documents into an in-memory Okapi BM25 index using the rank_bm25 package. It tokenizes page contents and ranks documents based on term frequency and inverse document frequency (TF-IDF derivative).

vectorizer¶

The underlying rank_bm25 search index object (BM25Okapi).

Type:: Any

docs¶

List of underlying LangChain Document objects.

Type:: List[Document]

k¶

Default number of documents to return per query.

Type:: int

preprocess_func¶

Function used to tokenize text.

Type:: Callable[[str], List[str]]

classmethod from_documents(documents: ~collections.abc.Iterable[langchain_core.documents.Document], *, bm25_params: dict[str, ~typing.Any] | None = None, preprocess_func: ~collections.abc.Callable[[str], list[str]] = <function default_preprocessing_func>, **kwargs: ~typing.Any) → BM25Retriever[source]¶

Create a BM25Retriever from a list of Documents. :param documents: A list of Documents to vectorize. :param bm25_params: Parameters to pass to the BM25 vectorizer. :param preprocess_func: A function to preprocess each text before vectorization. :param **kwargs: Any other arguments to pass to the retriever.

Returns:: A BM25Retriever instance.

classmethod from_texts(texts: ~collections.abc.Iterable[str], metadatas: ~collections.abc.Iterable[dict] | None = None, ids: ~collections.abc.Iterable[str] | None = None, bm25_params: dict[str, ~typing.Any] | None = None, preprocess_func: ~collections.abc.Callable[[str], list[str]] = <function default_preprocessing_func>, **kwargs: ~typing.Any) → BM25Retriever[source]¶

Create a BM25Retriever from a list of texts. :param texts: A list of texts to vectorize. :param metadatas: A list of metadata dicts to associate with each text. :param ids: A list of ids to associate with each text. :param bm25_params: Parameters to pass to the BM25 vectorizer. :param preprocess_func: A function to preprocess each text before vectorization. :param **kwargs: Any other arguments to pass to the retriever.

Returns:: A BM25Retriever instance.

Utilities¶

Vector store utilities

File: klea_utils/stores/utils.py

klea_utils.stores.utils.CACHE_DIR_NAME = '.klea-cache'¶: Name of the chunk-cache directory created inside a source directory by the ingestion pipeline (holds the per-file pickled chunks, the generated metadata-map template, the DOI cache, and store manifests). Excluded from ingestion by find_source_files().

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.RECENCY_MISSING_YEAR_SCORE = 0.5¶: Recency score assigned to documents without a usable year metadata value. A fixed midpoint (not a derived statistic) so it is immune to distribution skew: it ranks such documents below known-recent papers but above the oldest retrieved document.

klea_utils.stores.utils.RECENCY_WEIGHT_RELEVANCE = 0.9¶: Weight given to the normalized relevance (RRF) component of the final blended score in rerank_by_recency().

klea_utils.stores.utils.RECENCY_WEIGHT_TIME = 0.1¶: Weight given to the recency (time) component of the final blended score. Newer documents are boosted because academic work builds on – and often corrects – earlier results, so recent information is more authoritative.

klea_utils.stores.utils.REF_DOC_OVERHEAD = 200¶: Per-document character overhead attributed by truncate_reference_material() for the serialized markup that wraps each reference in the LLM context (### Source document N/M: [file] plus the optional metadata line). Approximate; the page content dominates in practice.

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.TEMPLATE_FILE_NAME = 'metadata-map.template.json'¶: Name of the metadata-map template chunk writes into the cache directory, organised per source file.

klea_utils.stores.utils.display_person_names(names: list[str]) → list[str][source]¶

Return the display-safe subset of an expanded person-name list.

expand_person_names() appends per-word and lowercase variants to a person-name field for filtering, but those variants must not appear in the reference material shown to the answer LLM (citations would otherwise read Ankur Sinha; Sinha; ankur; sinha). This keeps only the real names: whole-lowercase entries and single-word entries that make up a longer entry are dropped, so a genuine single-name author (not a token of another entry) still displays.

Example:

display_person_names(
    ["Ankur Sinha", "Padraig Gleeson", "Sinha", "gleeson"]
)
-> ["Ankur Sinha", "Padraig Gleeson"]

Parameters:: names – Person-name list, possibly expanded
Returns:: The original full names, order-preserving

klea_utils.stores.utils.drop_collection(store, store_path: str, collection_name: str) → None[source]¶

Drop a vector store collection (portable across backends).

Used by store --force to replace a collection wholesale, since documents within a collection cannot be updated in place portably. Chroma and PGVector expose delete_collection on their LangChain wrapper; Qdrant’s wrapper does not, so its raw client is used instead.

Parameters:

store – Instantiated LangChain vector store
store_path – Store URI (scheme:location)
collection_name – Collection name to drop

Raises:

ValueError – When the scheme is missing or unknown

klea_utils.stores.utils.expand_person_names(names: list[str]) → list[str][source]¶

Return names with per-word variants appended, order-preserving.

Humans are referred to by parts of their full name (“find papers by Sinha” for an author stored as “Ankur Sinha”), so person-name list fields (klea_utils.stores.metadata.PERSON_NAME_FILTER_FIELDS) are expanded at store time: each full name is kept, every whitespace token plus its lowercase form is added, and the lowercased full name is added too. This makes an exact-membership retrieval filter ($contains) match the partial name in any form and case, uniformly on every store backend.

Example:

expand_person_names(["Ankur Sinha"])
-> ["Ankur Sinha", "ankur sinha", "Ankur", "Sinha",
    "ankur", "sinha"]

The expansion is idempotent (expanding an already-expanded list is a no-op), so re-applying the store metadata policy is safe.

Parameters:: names – List of author display names
Returns:: Names plus per-word and lowercase variants, deduplicated in order of first appearance. Non-string elements are skipped.

klea_utils.stores.utils.find_source_files(source_dir: Path, *, metadata_map_path: Path | None = None, store_dir: Path | None = None, logger: Logger | None = None) → list[Path][source]¶

Walk source_dir and return files whose extensions are in docling’s FormatToExtensions.

This is the canonical “what will the store ingest” enumeration: the ingestion pipeline and map-lint both use it, so a metadata map that lints clean against its output is guaranteed to resolve at store time. Files with unsupported extensions are logged as a warning (when a logger is given) and skipped.

Generated artifacts are excluded: the cache directory (CACHE_DIR_NAME, e.g. .klea-cache), the metadata map passed via metadata_map_path (when it lives inside source_dir), and the vector store directory – either the configured store_dir when it lies under source_dir, or any directory inside source_dir that contains a chroma.sqlite3 (so a store created without setting store_dir is still not ingested).

Parameters:

source_dir – Directory to walk recursively
metadata_map_path – The metadata-map file to exclude from ingestion (mirrors how the ingestion pipeline remembers the loaded map)
store_dir – Configured vector store directory that may live inside the source directory; None for remote backends with no local folder
logger – Optional logger for the unsupported-extension warning

Returns:

Sorted list of files with supported extensions

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

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[source]¶

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.rerank_by_recency(merged: list[tuple[Document, float]], relevance_weight: float = 0.9, time_weight: float = 0.1) → list[tuple[Document, float]][source]¶

Re-rank RRF results blending in document recency.

Keeps rrf_merge() pure (relevance only) and applies recency as a separate post-fusion re-rank. Each document’s pure RRF score is min-max normalized to [0, 1] across the result set, a time score is computed from its year metadata, and the final score is a weighted combination:

final = relevance_weight * norm_rrf + time_weight * time_score

The time score is (year - year_min) / (year_max - year_min) where year_min/year_max are the min and max year across the retrieved set (relative normalization). Documents without a usable year (missing or non-int) get RECENCY_MISSING_YEAR_SCORE.

Division-by-zero cases are guarded: a single distinct RRF value maps to 1.0 and a single distinct year maps to 1.0.

Parameters:

merged – (doc, rrf_score) tuples from rrf_merge()
relevance_weight – Weight for the normalized relevance component
time_weight – Weight for the recency component

Returns:

The same documents, re-sorted descending by the blended score, with the blended score replacing the pure RRF score in each tuple

klea_utils.stores.utils.rrf_merge(result_sets: list[tuple[str, list[tuple[Document, float]]]], num_refs_max: int | None = None) → list[tuple[Document, float]][source]¶

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 debugging and introspection (they are not shown to the answer LLM).

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, or None to return every fused document. Callers that want to bound the context fed to an LLM should cap by characters via truncate_reference_material() instead of by document count.

Returns:

Documents ordered by RRF score, deduplicated by file+content, capped at num_refs_max when set

klea_utils.stores.utils.serialize_reference_material(reference_material: dict[str, list[tuple[Document, float]]]) → str[source]¶

Serialize reference material into text for use in prompt context.

Documents are grouped by their source file (file_name metadata) and each source file’s document-level metadata is emitted once, with the file’s chunks listed underneath. The shared bibliographic fields (authors, year, journal, …) are identical on every chunk of a file, so they are listed once on the source header; a url* key is also hoisted to the header when the whole file shares the same value. Per-chunk metadata that differs (e.g. a heading-specific url) is emitted inline so no chunk is misattributed. Files are ordered by their best chunk’s score; chunks within a file by score. Relevance scores are not included in the prompt – the ranked order is what matters to the answer LLM.

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

klea_utils.stores.utils.truncate_reference_material(reference_material: dict[str, list[tuple[Document, float]]], max_chars: int) → dict[str, list[tuple[Document, float]]][source]¶

Truncate reference material to a global character budget.

The RRF merge orders documents by fused rank but does not bound how much context the answer LLM receives; that is what this function does. Documents are consumed in RRF order per domain (domains in their dict order), counting len(page_content) plus the per-document serialization overhead (REF_DOC_OVERHEAD), until the budget is exhausted. The first document that crosses the budget is still included, so a single large chunk never silently yields empty context.

Parameters:

reference_material – {domain: [(doc, score), ...]} in RRF order
max_chars – Total character budget across all domains

Returns:

New mapping with the same domain keys, lists truncated to the budget