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(*, description: str = '', vector_stores: list[VectorStoreInfo] = [], bm25_stores: list[BM25StoreInfo] = [], filter_fields: list[FilterFieldInfo] = [])[source]¶

Bases: BaseModel

Configuration for a single domain.

description: str¶: Human-readable description of what the domain covers.

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

Ingestion pre-checks¶

Pre-check source directories to decide whether OCR is needed per PDF.

File: klea_utils/stores/precheck.py

klea_utils.stores.precheck.IMAGE_PAGE_CHAR_THRESHOLD = 50¶: A page whose extracted text has fewer than this many non-whitespace characters is treated as an image page (a scan or a figure-only page). Text-based PDFs easily produce thousands of characters per page, while a scanned page yields near-zero (just embedded text stamps/watermarks).

klea_utils.stores.precheck.IMAGE_PAGE_FRACTION = 0.5¶: A PDF is classified as needing OCR when at least this fraction of its pages are image pages.

klea_utils.stores.precheck.NO_OCR_DIR_NAME = 'no-ocr'¶: Output directory for text-based PDFs (and all non-PDF files).

klea_utils.stores.precheck.NO_PAGES_NEEDS_OCR = True¶: A PDF with no readable pages (e.g. zero pages) is conservatively classified as needing OCR – it is almost always a bad scan, and the safe default is to OCR it rather than silently store empty chunks.

klea_utils.stores.precheck.OCR_DIR_NAME = 'ocr'¶: Output directory for scanned/image PDFs (OCR required).

class klea_utils.stores.precheck.PageStats(pages: int = 0, image_pages: int = 0, total_chars: int = 0)[source]¶

Bases: object

Per-PDF page text statistics for the pre-check report.

image_pages: int = 0¶: Number of pages that fell below IMAGE_PAGE_CHAR_THRESHOLD.

pages: int = 0¶: Number of pages read from the PDF.

total_chars: int = 0¶: Total non-whitespace characters extracted across all pages.

class klea_utils.stores.precheck.PdfClassification(path: Path, needs_ocr: bool, stats: PageStats)[source]¶

Bases: object

The OCR decision and statistics for a single PDF.

klea_utils.stores.precheck.classify_directory(source_dir: Path) → dict[Path, PdfClassification][source]¶

Classify every PDF under source_dir.

Recursively walks **/*.pdf, skipping the chunk cache and the ocr/no-ocr organise output directories (generated artifacts, not source documents – see _SKIP_DIR_NAMES).

Logs progress at INFO roughly every 10% of PDFs (every file for small corpora) so a long pre-check does not sit silently before the final report prints.

Parameters:: source_dir – Directory to walk recursively
Returns:: Mapping of PDF path to its PdfClassification

klea_utils.stores.precheck.classify_pdf(path: Path) → PdfClassification[source]¶

Decide whether a PDF needs OCR, based on its embedded text layer.

Reads the PDF with pypdfium2 (already a Docling dependency) and counts non-whitespace characters per page. A page under IMAGE_PAGE_CHAR_THRESHOLD characters is an image page; a PDF with at least IMAGE_PAGE_FRACTION image pages needs OCR. Text-based (born-digital) PDFs carry a full text layer and classify as no-OCR; scanned/image PDFs have almost none and classify as needs-OCR.

A PDF that cannot be read (missing, corrupt, or pypdfium2 unavailable) is conservatively classified as needing OCR.

Parameters:: path – Path to the PDF file
Returns:: A PdfClassification with the decision and stats

klea_utils.stores.precheck.format_precheck_report(classifications: dict[Path, PdfClassification]) → str[source]¶

Render a classify_directory() result as compact text.

Parameters:: classifications – Mapping from classify_directory()
Returns:: Multi-line summary of the classification outcome

klea_utils.stores.precheck.organise_directory(source_dir: Path, classifications: dict[Path, PdfClassification] | None = None) → tuple[Path, Path][source]¶

Copy classified files into ocr/ and no-ocr/ subdirectories.

Copies (never moves) files so the researcher’s original bibliography directory is left untouched. Scanned/image PDFs go to ocr/; text-based PDFs go to no-ocr/. Non-PDF files (which never need OCR) are also copied to no-ocr/ so a single chunk no-ocr/ pass ingests them alongside the text PDFs.

Re-running is idempotent: existing output directories are reused and files are overwritten, and classify_directory() skips the output dirs so copies are never re-classified.

Parameters:

source_dir – Directory containing the source documents
classifications – Optional precomputed classifications (from classify_directory()); computed if not given

Returns:

(ocr_dir, no_ocr_dir) the two output directory paths

Ingestion post-checks¶

LLM-free linting of a stored BM25/vector corpus (post-store check).

File: klea_utils/stores/postcheck.py

klea_utils.stores.postcheck.MAX_LISTED = 20¶: How many short chunks / missing-metadata files to list before truncating the report (the rest are counted in the summary line).

klea_utils.stores.postcheck.MIN_CHUNK_CHARS = 20¶: A chunk whose page_content has fewer than this many characters is flagged as suspicious – typically a failed conversion or an OCR miss on a scanned PDF.

klea_utils.stores.postcheck.PER_ANCHOR_SAMPLES = 3¶: Contiguous chunks shown at each sampling location. The user reviews a small run of consecutive chunks per location rather than one isolated chunk, to check that chunking stays coherent across boundaries.

klea_utils.stores.postcheck.REQUIRED_METADATA_KEYS = ('file_name', 'title', 'year', 'doi')¶: Metadata keys that, when absent from a chunk, suggest the file’s DEFAULT entry did not apply or extraction failed. file_name is machine-set so its presence is expected; the others come from the map.

klea_utils.stores.postcheck.SAMPLE_TEXT_WIDTH = 200¶: Maximum characters of page_content shown per sample chunk.

klea_utils.stores.postcheck.format_chunk_sample(doc: Document, width: int = 200) → str[source]¶

Render one chunk for human review: truncated text + metadata line.

Parameters:

doc – A stored chunk document
width – Maximum characters of page_content to show

Returns:

Multi-line sample text

klea_utils.stores.postcheck.format_store_lint_report(report: dict[str, Any], samples: list[list[Document]]) → str[source]¶

Render a lint_store() report as compact text.

Parameters:

report – Report dict from lint_store()
samples – List of sample windows (each a list of Document) to print, from select_sample_windows()

Returns:

Multi-line report, including the given sample windows

klea_utils.stores.postcheck.lint_store(docs: list[Document]) → dict[str, Any][source]¶

Return a structured report of issues in a stored chunk corpus.

Runs only deterministic checks (no LLM): a corpus summary, suspicious chunks (near-empty text, missing bibliographic metadata), and structural problems (chunks without a file_name, invalid types). docs is the unpickled BM25 corpus (a list of langchain_core.documents.Document).

Parameters:: docs – List of stored chunk documents
Returns:: A dict with total, files, chunks_per_file, total_chars, empty (list of (file_name, chars)), missing_metadata ({file_name: [missing keys]}), no_file_name (count of chunks lacking file_name), invalid_content (count of non-string page_content), and invalid_year (count of non-int year)

klea_utils.stores.postcheck.select_sample_windows(docs: list[Document], anchors: int, per_anchor: int = 3) → list[list[Document]][source]¶

Pick anchors evenly-spaced windows of contiguous chunks to review.

Each window is per_anchor consecutive chunks starting at an anchor position, so the user sees a small coherent run at several different places in the corpus (early/middle/late) rather than the first few chunks. Anchor starts are evenly spaced across the whole corpus and avoid the very first/last chunk, giving roughly 25% / 50% / 75% for three anchors.

For a corpus smaller than one window, a single window of the whole corpus is returned. The selection is deterministic – no RNG – so re-runs show the same samples for easy comparison.

Parameters:

docs – List of stored chunk documents
anchors – Number of sampling locations (--samples)
per_anchor – Contiguous chunks to show at each location

Returns:

List of windows (each a list of Document); empty when anchors <= 0 or the corpus is empty

Chunk workers¶

Subprocess chunking workers – isolate Docling’s per-conversion memory leak

See devdocs/system/store-create.md; ADRs 0001 and 0002.

File: klea_utils/stores/chunk_worker.py

class klea_utils.stores.chunk_worker.ChunkItemResult(file_name: str, file_hash: str, status: str, file_headings_entry: dict[str, Any] | None = None, error: str | None = None, file_path: str | None = None)[source]¶

Bases: object

Outcome of converting one file inside a worker.

status is one of:

ok – converted and cached; file_headings_entry holds the file’s metadata-map template entry
zero_chunks – converted but produced no chunks (scanned PDF)
failed – conversion raised; error has the message
deferred – not processed because the worker hit its RSS cap; file_path lets the parent re-dispatch it to a fresh worker

class klea_utils.stores.chunk_worker.ChunkWorkerConfig(source_path: str, max_tokens: int, tokenizer_model: str, do_ocr: bool, mem_limit_bytes: int | None, log_level: int)[source]¶

Bases: object

Picklable configuration handed to one conversion worker batch.

The parent’s StoresBuilder is not picklable (it holds the Docling converter, chunker, resolver, etc.), so the worker rebuilds its own from these plain fields.

klea_utils.stores.chunk_worker.convert_batch_worker(config: ChunkWorkerConfig, items: list[tuple[str, str]], queue: Any | None = None, current_file: Any | None = None) → list[ChunkItemResult] | None[source]¶

Convert and cache items in this process, one per result.

Runs in a spawned child for real chunking (with queue and current_file for observability), or directly in the parent under tests (queue is None and results are returned). Rebuilds its own StoresBuilder, so Docling’s per-conversion memory leak (unreclaimable in-process; see docling-project/docling#2788) is confined to this process and reclaimed by the OS when it exits.

Before each file, if the process RSS exceeds config.mem_limit_bytes the remaining items are returned as deferred (the parent re-runs them in a fresh worker). On platforms where RSS is unavailable the cap is skipped and the parent’s fixed batch size bounds the worker instead.

When queue is given (spawned worker), each ChunkItemResult is put on the queue incrementally so the parent can log progress and, on worker death, knows which file was being processed via current_file and which results were already completed. In that mode this function returns None. Without a queue it returns the list – this path exists only for unit tests (in-process direct calls) so they stay simple and hermetic.

Parameters:

config – Worker configuration
items – (absolute_file_path, file_hash) tuples to process
queue – Queue for incremental puts (spawned worker) or None (in-process unit tests only)
current_file – Shared Value for the file currently being processed, or None (unit tests omit it)

Returns:

List of results when queue is None (tests), else None (spawned worker)

klea_utils.stores.chunk_worker.dispatch_conversion_batches(parent_logger: Logger, config: ChunkWorkerConfig, items: list[tuple[str, str]], batch_size: int) → list[ChunkItemResult][source]¶

Convert items by spawning a fresh worker process per batch.

A fresh process per batch is deliberate: Docling leaks memory per conversion that cannot be freed in-process, so each worker’s leak is reclaimed when it exits. A worker stops at whichever comes first – its RSS cap (config.mem_limit_bytes) or batch_size – marking the unprocessed tail deferred; the parent re-dispatches those to a fresh worker. On platforms where RSS is unavailable only batch_size bounds a worker. A worker that dies without returning results (e.g. OOM-killed) has its whole batch marked failed and the run continues – those files are simply re-converted on a later run.

Parameters:

parent_logger – Logger for dispatch progress/errors
config – Worker configuration
items – (absolute_file_path, file_hash) items to convert
batch_size – Max items handed to any single worker

Returns:

Flattened ChunkItemResult for every item

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) limits with graph-wide fallbacks, and the retrieval contract retrieve(domain, query, k_values=...) -> list[tuple[Document, float]].

One manager is shared by all graph threads, so it does not keep the current k. The graph keeps it in its state and passes it in.

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(k_values: dict[str, dict[str, int]]) → 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.

Parameters:: k_values – Current k for each store
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(k_values: dict[str, dict[str, int]]) → tuple[dict[str, dict[str, int]], 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.

Parameters:: k_values – Current k for each store; left unchanged
Returns:: The new k for each store, and 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.

retrieve(domain_name: str, query: str, metadata_filter: dict[str, Any] | None = None, k_values: dict[str, dict[str, int]] | None = None) → list[tuple[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)
k_values – Current k for each store; stores not listed fall back to their per-store configured default

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[Document], f: dict[str, Any]) → list[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)

Retrieval query schema¶

Strict-safe retrieval-query schema builder.

Builds the RAG query generator’s structured-output schema from the domain’s configured filter fields, so each allowed filter is a real, typed, optional field instead of a dynamic-key filters object. A dynamic-key object becomes a JSON object with additionalProperties, which providers’ strict structured-output modes close to {} (ADR-0044) - the model then cannot express any constraint at all. One typed field per configured filter avoids that and restricts the model to the deployment’s declared field names.

Field type follows FilterFieldInfo.value_type:

string -> str or list[str] (a list is an $in constraint);
int -> int, list[int], or an IntRange ({"gte": ..., "lte": ...}; equality is a bare value);
float -> float, list[float], or a FloatRange;
list -> list[str] (element membership).

Every filter field is optional (omitted when the question states no constraint). A bare value is equality; a *Range value carries eq/gte/lte bounds.

The built schema is: {search_query: str, <filter field>: <type> | None, ...}.

File: klea_utils/stores/query_schema.py

class klea_utils.stores.query_schema.FloatRange(*, eq: float | None = None, gte: float | None = None, lte: float | None = None)[source]¶

Bases: BaseModel

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

class klea_utils.stores.query_schema.IntRange(*, eq: int | None = None, gte: int | None = None, lte: int | None = None)[source]¶

Bases: BaseModel

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

klea_utils.stores.query_schema.build_retrieval_query_schema(fields: list[FilterFieldInfo]) → type[BaseModel][source]¶

Build the query generator’s strict-safe structured-output schema.

Parameters:: fields – The allowed filter fields for the run’s domains (already resolved by the node).
Returns:: A pydantic model with search_query and one optional typed field per valid, uniquely-named configured filter.

BM25 index¶

class klea_utils.stores.langchain_bm25.BM25Retriever(*args: ~typing.Any, name: str | None = None, tags: list[str] | None = None, metadata: dict[str, ~typing.Any] | None = None, vectorizer: ~typing.Any = None, docs: list[~langchain_core.documents.base.Document], k: int = 4, preprocess_func: ~collections.abc.Callable[[str], list[str]] = <function default_preprocessing_func>)[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.base.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.

model_config = {'arbitrary_types_allowed': True, 'extra': 'ignore', 'protected_namespaces': ()}¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None¶

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

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.DEFAULT_CROSS_ENCODER_MODEL = 'cross-encoder/ms-marco-MiniLM-L-6-v2'¶: Default cross-encoder model for reranking. Small MS MARCO model; runs locally via sentence-transformers.

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.STORES_CONFIG_ENV_VAR = 'KLEA_STORES_CONFIG_FILE'¶: Environment variable naming the deployment-default stores configuration file. Read once from the server process environment and used as the fallback when a request carries no STORES_CONFIG_META_KEY value.

klea_utils.stores.utils.STORES_CONFIG_META_KEY = 'klea/stores-config-path/v1'¶: Versioned request-metadata key carrying the stores-configuration JSON path to use for this request. Its value is exactly one absolute path to a stores JSON file; its absence in the request metadata means “use the host default” (the path configured by the agent app). Versioned so the payload format can evolve without colliding with an older reader.

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.cross_encoder_rerank(query: str, docs: list[tuple[Document, float]], *, model: Any) → list[tuple[Document, float]][source]¶

Re-rank fused retrieval results with a cross-encoder instance.

Intended to run after rrf_merge() and before downstream steps such as rerank_by_recency() and truncate_reference_material().

The caller owns the model: load it once (see load_cross_encoder()) and pass the same instance here. This function neither caches nor skips, so “reranking disabled” is the caller’s decision.

Each document is scored by a query–passage cross-encoder (local/offline via sentence-transformers). Original per-source scores in SOURCE_SCORES_KEY metadata are preserved; only the tuple’s relevance score is replaced by the cross-encoder score.

Parameters:

query – Normalized retrieval query (same form as passed to retrievers)
docs – (document, score) tuples, typically from rrf_merge()
model – A loaded cross-encoder exposing predict(pairs, show_progress_bar=...) -> scores

Returns:

Documents ordered by cross-encoder score (descending)

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.load_cross_encoder(model_name: str) → Any[source]¶

Load a CrossEncoder instance.

A plain factory: it imports sentence-transformers, downloads the weights on first use (into the Hugging Face cache), and returns a new instance. It performs no caching and holds no shared state, so callers that need one shared model must load it once and reuse the instance themselves. It is likewise not safe for concurrent loads (each call constructs a new model), so load once at startup rather than from several worker threads.

Parameters:

model_name – Hugging Face cross-encoder model id. Must be a non-empty string.

Returns:

A newly constructed cross-encoder for model_name

Raises:

ValueError – If model_name is empty or None
ImportError – If sentence-transformers is not installed

klea_utils.stores.utils.load_retriever_config(path: str | Path) → RetrieverConfig[source]¶

Load (and cache) the RetrieverConfig stored in path.

The cache is keyed by the resolved path and invalidated whenever the file’s mtime or size changes, so an edited stores file is re-read on the next call while an unchanged one is served from memory.

Parameters:: path – JSON file holding the stores configuration
Returns:: The parsed configuration
Raises:: FileNotFoundError – When path 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 results blending in document recency.

Keeps rrf_merge() pure (relevance only) and applies recency as a separate post-fusion re-rank. Each document’s relevance 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_relevance + 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 relevance value maps to 1.0 and a single distinct year maps to 1.0.

Parameters:

merged – (document, score) tuples, typically from rrf_merge() (or the cross-encoder rerank when enabled)
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 input score in each tuple

klea_utils.stores.utils.resolve_stores_config_path(configured: str | Path | None = None) → Path | None[source]¶

Resolve the path of the stores configuration JSON file.

Precedence:

The deployment default named by the STORES_CONFIG_ENV_VAR environment variable (read once by _host_stores_config_path() and validated the same way as configured). This is a process-wide default, so it is consulted only when the request carries no path.
Fail closed: when neither source is configured, None is returned and a warning naming both the request-metadata key and the env var is logged rather than guessing a path.

Parameters:: configured – Stores-configuration path supplied for this request, if any
Returns:: The resolved path, or None when nothing is configured
Raises:: ValueError – When configured is unusable: empty, missing, not a regular file, or not a valid stores configuration

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