Skip to content

API Reference

EmbeddingReaderInterface

Backend-neutral read interface for embedding search and retrieval.

Generic over the store type S it holds, so backend keeps the type of the store passed in.

Parameters:

Name Type Description Default
backend S

Store from open_vector_store_reader() or open_vector_store_writer().

required
metric_type MetricType

Distance metric used for KNN queries and validated against the registry.

required
cdm_session_factory CDMSessionFactory

Opens sessions on the user's OMOP CDM, e.g. from omop_alchemy.cross_database.cdm_sessionmaker. When provided, KNN results are enriched with concept_name from the CDM. When absent, concept_name is None. domain_id, vocabulary_id, is_standard, and is_active are populated directly from the embedding table by the backend.

None
model str

Model name that is expected to be canonicalized by the constructor. A possibly-canonical name is accepted as we cannot know if canonical or not.

required
provider_type str

omop-llm provider key. Defaults to 'ollama'.

'ollama'
k int

Default number of nearest neighbors to return.

DEFAULT_K_NEAREST
faiss_cache_dir str

Optional directory for FAISS index caching. If provided, the interface will attempt to use FAISS for faster KNN search. Only supported if the 'faiss' package is installed.

None

generate_embeddings staticmethod

generate_embeddings(
    model_backend: ModelBackend,
    text: Union[str, List[str], Tuple[str, ...]],
    *,
    role: EmbeddingRole,
    batch_size: Optional[int] = None,
) -> EmbeddingArray

Embed text against model_backend, validating the resulting shape.

Callable without an interface instance for callers that hold a ModelBackend directly (e.g. on-the-fly query embedding against a registry-resolved model, decoupled from any specific write session). Role-prefix application is model_backend's own job (see ModelBackend.embed_texts's role parameter); this method only forwards role and validates the returned shape.

Parameters:

Name Type Description Default
model_backend ModelBackend

Backend to generate embeddings with.

required
text str | list[str] | tuple[str, ...]

Input text(s) to embed.

required
role EmbeddingRole

Role of the input text(s), used by model_backend to apply a configured prefix.

required
batch_size int

Forwarded to ModelBackend.embed_texts.

None

Returns:

Type Description
EmbeddingArray

Float32 array of shape (n_texts, embedding_dim).

Raises:

Type Description
RuntimeError

If the model returns something other than a finite 2-D array, or a different number of vectors than texts.

get_embedding_count_by_vocabulary

get_embedding_count_by_vocabulary() -> Mapping[str, int]

Return stored embedding counts grouped by vocabulary_id.

Returns:

Type Description
Mapping[str, int]

vocabulary_id to embedding count, for every vocabulary with at least one stored embedding.

get_nearest_concepts

get_nearest_concepts(
    query_embedding: EmbeddingArray,
    *,
    concept_filter: Optional[EmbeddingConceptFilter] = None,
    k: Optional[int] = None,
    faiss_index_config: Optional[IndexConfig] = None,
) -> Tuple[Tuple[NearestConceptMatch, ...], ...]

Return nearest stored concepts for each query embedding row.

Parameters:

Name Type Description Default
query_embedding ndarray

Shape (Q, D), Q query vectors of dimension D. D must match the embedding dimension for the registered model.

required
concept_filter EmbeddingConceptFilter

In-DB pre-filter applied during KNN (domain, vocabulary, standard).

None
k int

Number of nearest neighbors (defaults to interface-level k).

None

Returns:

Type Description
Tuple[Tuple[NearestConceptMatch, ...], ...]

Shape (Q, ≤k). A row has fewer than k entries only when fewer than k stored concepts exist that match concept_filter (or exist at all). domain_id, vocabulary_id, is_standard, and is_active are always populated from the embedding table. concept_name is None if no CDM engine was provided to the interface.

get_nearest_concepts_from_query_texts

get_nearest_concepts_from_query_texts(
    query_texts: Union[str, Tuple[str, ...], List[str]],
    model_backend: Optional[ModelBackend] = None,
    *,
    concept_filter: Optional[EmbeddingConceptFilter] = None,
    batch_size: Optional[int] = None,
    k: Optional[int] = None,
    faiss_index_config: Optional[IndexConfig] = None,
) -> Tuple[Tuple[NearestConceptMatch, ...], ...]

Embed query_texts then search for nearest concepts.

get_similar_concepts

get_similar_concepts(
    concept_ids: Union[int, Sequence[int]],
    k: Optional[int] = None,
    *,
    concept_filter: Optional[EmbeddingConceptFilter] = None,
    faiss_index_config: Optional[IndexConfig] = None,
) -> Tuple[Tuple[NearestConceptMatch, ...], ...]

Return nearest stored concepts for one or more already-embedded concepts.

Convenience wrapper around :meth:get_nearest_concepts: resolves each of concept_ids to its own stored embedding (via :meth:get_embeddings_by_concept_ids) instead of requiring the caller to fetch and pass a raw vector. Each concept is excluded from its own row of results.

Parameters:

Name Type Description Default
concept_ids int or sequence of int

One or more concept IDs to search neighbours for. Every ID must already have a stored embedding. A bare int is treated as a single-element sequence.

required
k int

Number of nearest neighbours to return per query concept (defaults to interface-level k).

None
concept_filter EmbeddingConceptFilter

In-DB pre-filter applied during KNN (domain, vocabulary, standard).

None
faiss_index_config IndexConfig

Required only if a FAISS cache is configured on this interface.

None

Returns:

Type Description
Tuple[Tuple[NearestConceptMatch, ...], ...]

Shape (Q, ≤k) where Q == len(concept_ids) (1 for a bare int), in the same order as concept_ids. Each row excludes its own query concept.

Raises:

Type Description
ValueError

If concept_ids is empty, or any entry has no stored embedding.

get_joint_embedding

get_joint_embedding(
    concept_ids: Tuple[int, ...],
    weights: Optional[Tuple[float, ...]] = None,
) -> np.ndarray

Return the (optionally weighted) centroid of stored concept embeddings.

Parameters:

Name Type Description Default
concept_ids tuple of int

Concept IDs whose stored embeddings should be combined. Must be non-empty, and every ID must already have a stored embedding for the interface's model.

required
weights tuple of float

Per-concept weight, same length as concept_ids. Defaults to an unweighted mean.

None

Returns:

Type Description
ndarray

Shape (D,) centroid vector, suitable as a single query row for :meth:get_nearest_concepts. Not normalised: backends compute cosine distance directly from raw vectors (and the FAISS cache normalises query vectors internally for COSINE), so this method has no normalisation to do regardless of metric_type.

Raises:

Type Description
ValueError

If concept_ids is empty, weights has a mismatched length, weights sums to zero, or any concept_ids entry has no stored embedding.

get_concepts_without_embedding

get_concepts_without_embedding(
    *, concept_filter: Optional[CDMConceptFilter] = None
) -> Mapping[int, Row]

Return CDM rows for concepts lacking embeddings, keyed by concept_id.

Each row contains concept name and filter metadata, including canonical standardness and validity derived by omop-alchemy.

count_concepts_without_embedding

count_concepts_without_embedding(
    *, concept_filter: Optional[CDMConceptFilter] = None
) -> int

Return how many CDM concepts match concept_filter but lack an embedding.

get_concepts_without_embedding_batched

get_concepts_without_embedding_batched(
    *,
    batch_size: int,
    concept_filter: Optional[CDMConceptFilter] = None,
    limit: Optional[int] = None,
) -> Iterable[Mapping[int, Row]]

Yield {concept_id: Row} batches for concepts lacking embeddings.

Streams CDM rows and filters against already-embedded IDs on-the-fly, so only one batch of CDM rows is in memory at a time.

count_concepts_requiring_embedding

count_concepts_requiring_embedding(
    *,
    concept_filter: Optional[CDMConceptFilter] = None,
    batch_size: int = 10000,
) -> int

Count missing concepts and concepts with changed filter metadata.

get_concepts_requiring_embedding_batched

get_concepts_requiring_embedding_batched(
    *,
    batch_size: int,
    concept_filter: Optional[CDMConceptFilter] = None,
    limit: Optional[int] = None,
) -> Iterable[Mapping[int, Row]]

Yield concepts whose embedding is missing or has stale metadata.

CDM rows are streamed and stored metadata is fetched in bounded batches. Existing embeddings are reprocessed only when their domain, vocabulary, standardness, or validity differs from the CDM.

EmbeddingWriterInterface

Bases: EmbeddingReaderInterface[EmbeddingBackend]

Reader interface extended with embedding generation and write operations.

Builds and owns an omop_llm.ModelBackend directly: there is no separate client object between this interface and the model backend.

Parameters:

Name Type Description Default
backend EmbeddingBackend

Pre-constructed backend.

required
metric_type MetricType

Distance metric for the table.

required
resolved_model ResolvedModel

A model resolved via oa_configurator.Resolver.resolve_model(), e.g. Resolver.from_active_config().resolve_model(cfg.embedding_model_name). Provider, connection details, embedding_dim, and document_prefix/query_prefix all come from this, not from omop-emb's own config.

required
embedding_batch_size int

Default number of texts per API call. Default is 32.

32
cdm_session_factory CDMSessionFactory

Opens CDM sessions, for result enrichment and for the methods that read concepts to decide what to embed.

None

embedding_dim property

embedding_dim: int

Embedding vector dimension, resolved on first access and cached.

Delegates entirely to ModelBackend.dimensions()'s three-tier lookup (configured override, provider fast path, live probe); the configured override was already passed in as configuration at construction time.

register_model

register_model(
    *,
    index_config: Optional[IndexConfig] = None,
    metadata: Optional[Mapping[str, object]] = None,
) -> EmbeddingModelRecord

Register the embedding model in the backend registry.

Parameters:

Name Type Description Default
index_config IndexConfig

Index configuration. Defaults to FlatIndexConfig() when not provided.

None
metadata Mapping[str, object]

Free-form operational metadata. Must not contain reserved keys.

None

Returns:

Type Description
EmbeddingModelRecord

delete_model

delete_model() -> None

Irreversibly delete the model and all associated embeddings.

rebuild_index

rebuild_index(
    index_config: IndexConfig, vacuum: bool = True
) -> EmbeddingModelRecord

Build or rebuild the index on the embedding table.

Parameters:

Name Type Description Default
index_config IndexConfig

New index configuration.

required
vacuum bool

Run the backend's maintenance step after the rebuild (VACUUM (ANALYZE) on pgvector, VACUUM after a table recreation on sqlite-vec). Default True.

True

Returns:

Type Description
EmbeddingModelRecord

Updated registry record.

upsert_concept_embeddings

upsert_concept_embeddings(
    *,
    records: Sequence[ConceptEmbeddingRecord],
    embeddings: EmbeddingArray,
) -> None

Upsert pre-built ConceptEmbeddingRecords with their embeddings.

delete_concept_embeddings

delete_concept_embeddings(
    concept_ids: Sequence[int],
) -> int

Delete the stored embeddings of the given concepts and mark the model as updated.

Parameters:

Name Type Description Default
concept_ids Sequence[int]

OMOP concept IDs whose embeddings are removed.

required

Returns:

Type Description
int

Number of embedding rows deleted.

compact_storage

compact_storage() -> CompactionResult

Return the disk space held by deleted or rewritten rows, see EmbeddingBackend.compact_storage.

bulk_upsert_concept_embeddings

bulk_upsert_concept_embeddings(
    batches: Iterable[
        Tuple[
            Sequence[ConceptEmbeddingRecord], EmbeddingArray
        ]
    ],
    total_n_batches: Optional[int] = None,
    disable_progress: bool = False,
) -> None

Upsert from a lazy (records, embeddings) iterable.

embed_and_upsert_concepts

embed_and_upsert_concepts(
    *,
    concept_ids: Sequence[int],
    concept_texts: Sequence[str],
    concept_meta: Mapping[int, Row],
    batch_size: Optional[int] = None,
) -> EmbeddingArray

Generate and store concept embeddings with their filter metadata.

Parameters are the same as embed_concepts. Returns the generated embedding array after writing the batch and refreshing its timestamp.

embed_concepts

embed_concepts(
    *,
    concept_ids: Sequence[int],
    concept_texts: Sequence[str],
    concept_meta: Mapping[int, Row],
    batch_size: Optional[int] = None,
) -> Tuple[
    Sequence[ConceptEmbeddingRecord], EmbeddingArray
]

Prepare records and embeddings without writing them to the store.

The result can be yielded to bulk_upsert_concept_embeddings so index maintenance happens once across multiple generated batches. Concept IDs, texts and metadata must be aligned in the same order. fetch_cdm_concepts_for_filter can be used to get aligned concept_meta for a set of concept_ids.

Parameters:

Name Type Description Default
concept_ids Sequence[int]

OMOP concept IDs to embed.

required
concept_texts Sequence[str]

Text strings to embed (aligned with concept_ids).

required
concept_meta Mapping[int, Row]

CDM rows keyed by concept_id, as returned by get_concepts_without_embedding. Used to populate domain_id, vocabulary_id, is_standard, and is_valid.

required

Returns:

Type Description
tuple

Filter records in input order and their float32 embedding array.

BackendType

Bases: StrEnum

Embedding storage backend.

Members

SQLITEVEC Default backend. Requires no external database server. PGVECTOR Optional backend. Requires a PostgreSQL instance with the pgvector extension (pip install omop-emb[pgvector]). DUCKDBVSS Optional backend. Requires a DuckDB database and its VSS extension (pip install omop-emb[duckdbvss]).

IndexType

Bases: StrEnum

Index structure built on an embedding table for nearest-neighbor search.

Members

FLAT Exact sequential scan. Always correct, slower at scale. HNSW Approximate nearest-neighbor index. Supported by pgvector, DuckDB VSS, and FAISS, not by sqlite-vec.

MetricType

Bases: StrEnum

Distance metric used for nearest-neighbor search.

Members

L2 Euclidean (L2) distance. COSINE Cosine distance. L1 Manhattan (L1) distance. HAMMING Hamming distance (bit vectors). JACCARD Jaccard distance (bit vectors).

IndexConfig dataclass

Bases: ABC

Abstract base class for index configurations.

Attributes:

Name Type Description
index_type IndexType

Index structure (FLAT or HNSW).

metric_type MetricType or None

Distance metric the model is locked to. None for FLAT without a metric, which leaves the metric to the query.

to_dict

to_dict() -> dict[str, Any]

Serialise this config to a plain dict suitable for JSON storage.

Returns:

Type Description
dict[str, Any]

dataclasses.asdict output. Enum values are stored as their Python objects; the ORM JSON encoder converts them to strings.

Raises:

Type Description
TypeError

If called on a non-dataclass instance.

from_kwargs classmethod

from_kwargs(**kwargs: Any) -> Self

Instantiate from keyword arguments, silently ignoring unknown keys.

Parameters:

Name Type Description Default
**kwargs Any

Keyword arguments matching the dataclass fields. Unknown keys and explicit None values are discarded, so unset fields keep the dataclass's own default instead of being overridden with None.

{}

Returns:

Type Description
Self

A new instance of the concrete subclass.

Raises:

Type Description
TypeError

If called on the abstract base class or a non-dataclass.

from_dict classmethod

from_dict(config_dict: dict[str, Any]) -> Self

Reconstruct a config from the raw dict stored in the ORM column.

Parameters:

Name Type Description Default
config_dict dict[str, Any]

The dict previously produced by :meth:to_dict. Enum fields may be stored as string values and are coerced back to their types.

required

Returns:

Type Description
Self

A new instance of the concrete subclass.

Raises:

Type Description
TypeError

If called on the abstract base class or a non-dataclass.

ValueError

If a required field is missing or an enum value is unrecognised.

Notes

Use this method to reconstruct an IndexConfig from the ORM index_config JSON column.

FlatIndexConfig dataclass

Bases: IndexConfig

Configuration for a flat (exact scan) embedding table.

Attributes:

Name Type Description
index_type IndexType

Always IndexType.FLAT.

metric_type MetricType or None

Backends that fix the metric on the table (sqlite-vec) lock the model to it and default to cosine. Backends that compute distances from the raw vectors (pgvector) drop it, so queries choose the metric.

HNSWIndexConfig dataclass

Bases: IndexConfig

Configuration for an HNSW approximate nearest-neighbour index.

Attributes:

Name Type Description
index_type IndexType

Always IndexType.HNSW.

metric_type MetricType

Distance metric locked to this index. Required; cannot be None.

num_neighbors int

Number of bi-directional links per node (M in the HNSW paper). Default 32.

ef_search int

Size of the dynamic candidate list during search. Default 16.

ef_construction int

Size of the dynamic candidate list during index construction. Default 64.

Notes

Supported by pgvector, DuckDB VSS, and FAISS backends. sqlite-vec does not support HNSW.

EmbeddingModelRecord dataclass

Immutable snapshot of one registered embedding model.

Each record maps to exactly one row in the model_registry table, identified by model_name.

Attributes:

Name Type Description
model_name str

Canonical model name including tag (e.g. 'nomic-embed-text:v1.5').

provider_type str

omop-llm provider key that serves the model (e.g. 'ollama').

index_config IndexConfig

Active index configuration. index_type and metric_type are derived from this field via properties.

dimensions int

Embedding vector dimensionality.

storage_identifier str

Physical table name stored at registration. New names are store-scoped; existing rows keep the identifier already recorded here.

metadata Mapping[str, object]

Free-form operational data.

created_at (datetime, optional)

Row creation timestamp (UTC).

updated_at (datetime, optional)

Row last-updated timestamp (UTC).

Notes

metric_type is None for a FLAT model on a backend that computes distances from the raw vectors (pgvector). In that case any backend-supported metric is valid at query time. When metric_type is set (HNSW, or FLAT on a backend that fixes the metric on the table), queries must use that metric exactly.

index_type property

index_type: IndexType

Index type derived from index_config.

metric_type property

metric_type: Optional[MetricType]

Distance metric the model is locked to, or None if queries choose it.

RegistryManager

Registry of embedding models co-located with the embedding store.

For sqlite-vec the engine points at the same .db file as the vec0 tables. For pgvector it points at the same Postgres database.

Parameters:

Name Type Description Default
embedding_engine Engine

SQLAlchemy embedding engine connected to the embedding store. Must have been built via registry_reader_engine()/registry_writer_engine(), which is checked here to ensure the registry schema claim is present.

required
Notes

The registry table is created under the registry schema tag (a schema_translate_map key, resolved to the physical schema named by MODEL_REGISTRY_SCHEMA), independent of whichever schema the embedding store itself resolves to.

Raises:

Type Description
UnregisteredSchemaTagError

If embedding_engine wasn't built with the registry schema claim.

registry_available property

registry_available: bool

Whether the registry table currently exists on this engine.

True for any writable-bootstrapped engine (ensure_registry_table already ran). May be False for a peek-constructed engine on a store that's never been configured -- callers must not assume this is always True.

embedding_engine property

embedding_engine: Engine

SQLAlchemy engine connected to the embedding store.

emb_session_factory property

emb_session_factory: sessionmaker

Session factory bound to embedding_engine.

get_registered_models

get_registered_models(
    *,
    model_name: str | None = None,
    provider_type: str | None = None,
) -> tuple[EmbeddingModelRecord, ...]

Return all registered models matching the given filters.

Parameters:

Name Type Description Default
model_name str

Filter by canonical model name.

None
provider_type str

omop-llm provider key that serves the model.

None

Returns:

Type Description
tuple[EmbeddingModelRecord, ...]

register_model

register_model(
    *,
    model_name: str,
    index_config: IndexConfig,
    dimensions: int,
    provider_type: str,
    metadata: Mapping[str, object] | None = None,
    registered_at: datetime | None = None,
) -> EmbeddingModelRecord

Register a model or return the existing record if already registered.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
index_config IndexConfig

Initial index configuration.

required
dimensions int

Embedding vector dimensionality.

required
provider_type str

omop-llm provider key that serves the model.

required
metadata Mapping[str, object]

Free-form operational metadata.

None
registered_at datetime

Backdate created_at/updated_at to this timestamp instead of "now" (e.g. to match the source data's true snapshot time when importing an export bundle on a different machine). Only applies to a brand-new registration; ignored when the model is already registered.

None

Returns:

Type Description
EmbeddingModelRecord

The newly created or already-existing record.

Raises:

Type Description
ModelRegistrationConflictError

If the model is already registered with a different configuration.

ValueError

If metadata contains a reserved key.

delete_model

delete_model(*, model_name: str) -> None

Delete a registry row. No-op if the row does not exist.

Parameters:

Name Type Description Default
model_name str
required

update_index_config

update_index_config(
    *, model_name: str, index_config: IndexConfig
) -> EmbeddingModelRecord

Replace the index configuration of an existing registry row.

Parameters:

Name Type Description Default
model_name str
required
index_config IndexConfig

New index configuration. The @validates hook syncs index_type and metric_type columns automatically.

required

Returns:

Type Description
EmbeddingModelRecord

Updated record.

Raises:

Type Description
ValueError

If the model is not registered.

update_metadata

update_metadata(
    *, model_name: str, metadata: Mapping[str, object]
) -> EmbeddingModelRecord

Replace the free-form metadata of an existing registry row.

Parameters:

Name Type Description Default
model_name str
required
metadata Mapping[str, object]

New metadata dict. Replaces the existing value entirely.

required

Returns:

Type Description
EmbeddingModelRecord

Updated record.

Raises:

Type Description
ValueError

If the model is not registered.

refresh_model_updated_at_timestamp

refresh_model_updated_at_timestamp(
    *, model_name: str
) -> None

Bump a registry row's updated_at to now, with no other change.

Notes

Uses a Python-side timestamp rather than func.now() so the value carries microsecond precision on every backend. SQLite's CURRENT_TIMESTAMPonly has whole-second resolution, which made staleness checks racy when an upsert and its preceding FAISS export land in the same second.

Parameters:

Name Type Description Default
model_name str
required

safe_model_name staticmethod

safe_model_name(model_name: str) -> str

Normalise a model name for use in table identifiers.

Lowercases and replaces any run of non-word characters with a single underscore, then strips leading/trailing underscores.

Parameters:

Name Type Description Default
model_name str

Raw model name (e.g. 'nomic-embed-text:v1.5').

required

Returns:

Type Description
str

Normalised name (e.g. 'nomic_embed_text_v1_5').

storage_name staticmethod

storage_name(
    database_config_name: str | None, model_name: str
) -> str

Build the deterministic physical table name for a new model registration.

Parameters:

Name Type Description Default
database_config_name str or None

The owning [databases.*] entry name, if available.

required
model_name str

Canonical provider model ID.

required

Returns:

Type Description
str

Store-scoped name emb_<readable>_<8 hex SHA-256 characters>.

EmbeddingConceptFilter dataclass

Search constraints applied during KNN retrieval.

All fields are optional. Unset fields impose no constraint. To limit the number of KNN results returned, pass k to get_nearest_concepts/get_similar_concepts.

Notes

Mirrors OMOP grounding needs without importing omop_graph or its search-constraint types into omop_emb. Field shape duplicates omop_graph.graph.constraints.SearchConstraintConcept for the same reason; tracked at https://github.com/AustralianCancerDataNetwork/OMOP_Alchemy/issues/11, remove the duplication once that's resolved.

Attributes:

Name Type Description
concept_ids (tuple[int, ...], optional)

Restrict results to this set of concept IDs.

domains (tuple[str, ...], optional)

Restrict results to concepts in these OMOP domains.

vocabularies (tuple[str, ...], optional)

Restrict results to concepts from these vocabularies.

require_standard bool

When True, only concepts satisfying omop-alchemy's canonical Concept.is_standard rule are returned. Default False.

require_active bool

When True, only concepts satisfying omop-alchemy's canonical Concept.is_valid rule are returned. Default False.

is_empty

is_empty() -> bool

Return True if no constraints are set.

NearestConceptMatch dataclass

Single nearest-neighbour result as returned to callers.

concept_id and similarity are always populated by the backend. domain_id, vocabulary_id, is_standard, and is_active are also populated by the backend directly from the embedding table's filter columns (see :class:~omop_emb.backends.embedding_table.ConceptEmbeddingMixin), so they are available regardless of whether a CDM engine is configured. concept_name is the only field enriched by the interface layer from the OMOP CDM, and only when a cdm_session_factory is provided.

Attributes:

Name Type Description
concept_id int

OMOP concept ID of the matched concept.

similarity float

Similarity score in [0.0, 1.0]. Higher is more similar.

concept_name (str, optional)

Human-readable concept name. None when no CDM engine is provided.

domain_id (str, optional)

OMOP domain (e.g. 'Condition', 'Drug'), taken from the embedding table. None only if the backend could not resolve it.

vocabulary_id (str, optional)

Source vocabulary (e.g. 'SNOMED', 'RxNorm'), taken from the embedding table. None only if the backend could not resolve it.

is_standard (bool, optional)

True if standard_concept is 'S' or 'C', taken from the embedding table. None only if the backend could not resolve it.

is_active (bool, optional)

True if invalid_reason is not 'D' or 'U', taken from the embedding table. None only if the backend could not resolve it.

to_dict

to_dict() -> dict

Convert to a dictionary for serialization.

ConceptEmbeddingRecord dataclass

Concept metadata for a single embedding upsert row.

Populated from the OMOP CDM by the caller (interface layer) before being passed to the backend.

Attributes:

Name Type Description
concept_id int

OMOP concept ID.

domain_id str

OMOP domain (e.g. 'Condition', 'Drug').

vocabulary_id str

Source vocabulary (e.g. 'SNOMED', 'RxNorm').

is_standard bool

Value derived from Concept.is_standard_expr().

is_valid bool

Value derived from Concept.is_valid_expr().

matches_cdm_row

matches_cdm_row(row: Any) -> bool

Whether row's domain, vocabulary, standardness and validity equal this record's.

Parameters:

Name Type Description Default
row Any

CDM concept row with domain_id, vocabulary_id, is_standard and is_valid.

required

Returns:

Type Description
bool

EmbeddingBackend

Bases: EmbeddingStoreReader, ABC

Abstract base class for embedding storage and retrieval backends; implements EmbeddingStoreReader.

Parameters:

Name Type Description Default
emb_engine Engine

SQLAlchemy engine pointing at the embedding store. For sqlite-vec this is the same .db file as the registry; for pgvector it is the same Postgres database.

required
writable bool

False makes every @writes method raise ReadOnlyStoreError. Set by open_vector_store_reader().

True
Properties

backend_type : BackendType Backend type identifier (e.g. BackendType.PGVECTOR). backend_name : str String value of backend_type (e.g. "pgvector"). dialect : str SQL dialect this backend requires (e.g. "postgresql", "sqlite"). Matched against emb_engine.dialect.name at construction time. emb_engine : Engine SQLAlchemy engine for the embedding store. Obtained through the registry manager. emb_session_factory : sessionmaker Session factory bound to emb_engine. Obtained through the registry manager.

Notes

The SQLAlchemy engine is obtained from the RegistryManager to prevent duplication in code but also to have a single source of truth for the database connection.

backend_name property

backend_name: str

String value of backend_type.

dialect abstractmethod property

dialect: str

SQL dialect this backend requires, matching emb_engine.dialect.name.

default_flat_metric abstractmethod property

default_flat_metric: MetricType | None

Metric a FLAT model takes when none is given.

None for backends that compute distances from the raw vectors, whose FLAT models carry no metric. A metric for backends that lock it on the table.

emb_engine property

emb_engine: Engine

SQLAlchemy engine for the embedding store.

emb_session_factory property

emb_session_factory: sessionmaker

Session factory bound to emb_engine.

initialized property

initialized: bool

Whether the model registry table exists.

close

close() -> None

Dispose the underlying engine.

iter_stored_embeddings

iter_stored_embeddings(
    model_name: str, *, batch_size: int = 10000
) -> Iterator[ConceptEmbeddingRecord]

Stream concept metadata; yields nothing if the model is not registered.

physical_indexes abstractmethod

physical_indexes(model_name: str) -> tuple[str, ...]

Return existing physical indexes for this model's table, without creating or changing them.

drop_index_sql

drop_index_sql(model_name: str) -> tuple[str, ...]

Return reviewed index-removal statements without executing them.

register_model

register_model(
    *,
    model_name: str,
    dimensions: int,
    provider_type: str,
    index_config: IndexConfig | None = None,
    metadata: Mapping[str, object] | None = None,
    registered_at: datetime | None = None,
) -> EmbeddingModelRecord

Register a model and create its physical storage table.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
dimensions int

Embedding vector dimensionality.

required
provider_type str

omop-llm provider key that serves the model.

required
index_config IndexConfig

Index configuration. Defaults to FlatIndexConfig() when not provided.

None
metadata Mapping[str, object]

Free-form operational metadata. Must not contain reserved keys (see RESERVED_METADATA_KEYS in index_config.py).

None
registered_at datetime

Backdate created_at/updated_at to this timestamp instead of "now". Only applies to a brand-new registration.

None

Returns:

Type Description
EmbeddingModelRecord

The new or existing registry record.

Raises:

Type Description
ModelRegistrationConflictError

If the model is already registered with a different dimensionality.

ValueError

If metadata contains a reserved key, or if index_config is not a FLAT config (non-FLAT indexes may only be built after registration, not at registration time), or its metric is unsupported by the backend.

delete_model

delete_model(*, model_name: str) -> None

Delete the registry row and drop the physical embedding table.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required

Raises:

Type Description
ValueError

If the model is not registered.

Notes

One row per model means deletion always drops the physical table. There is no shared-table check.

The physical table is dropped before the registry row is removed. If the DDL step fails the registry entry remains intact and the call is re-runnable. If the DDL succeeds but the registry delete fails the only failure mode is a registry entry pointing at a table that no longer exists -- the next read or write against that model raises MissingStorageTableError instead of silently recreating it.

rebuild_index

rebuild_index(
    *,
    model_name: str,
    index_config: IndexConfig,
    vacuum: bool = True,
) -> EmbeddingModelRecord

Build or rebuild the index on an existing embedding table.

This is the unified entry point for both rebuilding an existing index and switching index type (e.g. FLAT → HNSW or HNSW → FLAT). Pass the desired target index_config regardless of the current state.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
index_config IndexConfig

Target index configuration. Pass HNSWIndexConfig (with metric_type set) to build an HNSW index, or FlatIndexConfig() to revert to an exact scan. On sqlite-vec a FlatIndexConfig locks the model to its metric (cosine if none is given); pgvector FLAT models take the metric per query, so any metric is dropped.

required
vacuum bool

Run the backend's maintenance step after the rebuild. pgvector runs VACUUM (ANALYZE) on the table to refresh planner statistics. sqlite-vec runs VACUUM on the database file after recreating a table for a metric change. Neither returns space freed by deleted rows, see compact_storage. Default True.

True

Returns:

Type Description
EmbeddingModelRecord

Updated registry record reflecting the new index type and metric.

Raises:

Type Description
ValueError

If the model is not registered, or the index type or metric is unsupported by the backend.

Notes

Afterwards the table carries exactly the index index_config describes: any other physical index on it is dropped. Queries then must use index_config.metric_type when it is set.

get_registered_model

get_registered_model(
    *, model_name: str
) -> EmbeddingModelRecord | None

Return the registry record for model_name, or None if not registered.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required

Returns:

Type Description
EmbeddingModelRecord or None

get_registered_models

get_registered_models(
    *,
    model_name: str | None = None,
    provider_type: str | None = None,
) -> tuple[EmbeddingModelRecord, ...]

Return all registry entries for this backend, with optional filters.

Parameters:

Name Type Description Default
model_name str

Canonical model name to filter by.

None
provider_type str

omop-llm provider key that serves the model.

None

Returns:

Type Description
tuple[EmbeddingModelRecord, ...]

patch_model_metadata

patch_model_metadata(
    *, model_name: str, key: str, value: object
) -> None

Merge a single key-value pair into the registry metadata.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
key str

Metadata key to set or overwrite. Must not be a reserved key.

required
value object

JSON-serialisable value.

required

Raises:

Type Description
ValueError

If the model is not registered or key is a reserved metadata key.

refresh_model_updated_at_timestamp

refresh_model_updated_at_timestamp(
    *, model_name: str
) -> None

Bump a registry row's updated_at to now. Required for faiss-cache freshness validation after live upserts. No-op if the model is not registered.

Parameters:

Name Type Description Default
model_name str
required

sync_index

sync_index(*, model_name: str) -> None

Bring model_name's physical index back in line with the table. No-op on a backend that doesn't need it.

upsert_embeddings

upsert_embeddings(
    *,
    model_name: str,
    records: Sequence[ConceptEmbeddingRecord],
    embeddings: EmbeddingArray,
    defer_index_update: bool = False,
) -> None

Insert or update embeddings for a set of concepts.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
records Sequence[ConceptEmbeddingRecord]

Concept metadata rows aligned with embeddings.

required
embeddings EmbeddingArray

Array of shape (N, D) where N = len(records) and D is the registered dimensionality. Converted to float32.

required
defer_index_update bool

Skip the index update for this call. Call sync_index yourself afterwards.

False

Raises:

Type Description
ValueError

If embeddings is not 2-D, has the wrong dimensionality or contains NaN or infinite values, or its row count differs from len(records).

delete_embeddings

delete_embeddings(
    *,
    model_name: str,
    concept_ids: Sequence[int],
    defer_index_update: bool = False,
) -> int

Delete the stored embeddings of the given concepts.

Concept IDs without a stored embedding are ignored. The deletion runs in chunks of _DELETE_CHUNK_SIZE IDs, one transaction each.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
concept_ids Sequence[int]

OMOP concept IDs whose embeddings are removed.

required
defer_index_update bool

Skip the index update for this call. Call sync_index yourself afterwards.

False

Returns:

Type Description
int

Number of embedding rows deleted.

compact_storage

compact_storage(*, model_name: str) -> CompactionResult

Return the disk space held by deleted or rewritten rows of a model.

Deleting rows, and upserting existing concepts, leaves space that the backend does not hand back on its own. This is the only operation that does. It rewrites the model's storage, so it needs free disk space for a second copy and, on pgvector, an exclusive lock on the table.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required

Returns:

Type Description
CompactionResult

Storage size before and after.

bulk_upsert_embeddings

bulk_upsert_embeddings(
    *,
    model_name: str,
    batches: Iterable[
        tuple[
            Sequence[ConceptEmbeddingRecord], EmbeddingArray
        ]
    ],
    total_n_batches: int | None = None,
    disable_progress: bool = False,
) -> None

Upsert embeddings in multiple batches, delegating to upsert_embeddings.

Updates the index once, after all batches, instead of once per batch.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
batches Iterable[tuple[Sequence[ConceptEmbeddingRecord], EmbeddingArray]]

Iterable of (records, embeddings) pairs.

required
total_n_batches Optional[int]

Total number of batches for the progress bar.

None
disable_progress bool

Suppress the progress bar, e.g. in a non-interactive or logged context.

False

get_embeddings_by_concept_ids

get_embeddings_by_concept_ids(
    *, model_name: str, concept_ids: Sequence[int]
) -> Mapping[int, ndarray]

Retrieve stored embeddings for the given concept IDs.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
concept_ids Sequence[int]

OMOP concept IDs to look up.

required

Returns:

Type Description
Mapping[int, ndarray]

Mapping from concept ID to its float32 embedding vector.

Raises:

Type Description
ValueError

If any requested concept ID is not found in the table.

get_nearest_concepts

get_nearest_concepts(
    *,
    model_name: str,
    metric_type: MetricType,
    query_embeddings: EmbeddingArray,
    concept_filter: EmbeddingConceptFilter | None = None,
    k: int = DEFAULT_K_NEAREST,
) -> tuple[tuple[NearestConceptMatch, ...], ...]

Find the nearest stored concepts for one or more query vectors.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
metric_type MetricType

Distance metric for the KNN query. Must match the HNSW metric exactly, or be any backend-supported metric for FLAT.

required
query_embeddings EmbeddingArray

Array of shape (Q, D) where Q is the number of queries and D is the registered dimensionality. Converted to float32.

required
concept_filter EmbeddingConceptFilter

Constraints applied during retrieval (domain, vocabulary, standard flag, concept ID allowlist, result limit).

None
k int

Maximum number of results per query. Default 10.

DEFAULT_K_NEAREST

Returns:

Type Description
tuple[tuple[NearestConceptMatch, ...], ...]

Shape (Q, <=K). Outer tuple is one entry per query vector; inner tuple contains up to k matches ordered by similarity descending.

Raises:

Type Description
ValueError

If query_embeddings is not 2-D, has the wrong dimensionality or contains NaN or infinite values, or the metric is incompatible with the model.

has_any_embeddings

has_any_embeddings(*, model_name: str) -> bool

Return True if at least one embedding row exists in the table.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required

Returns:

Type Description
bool

get_stored_concept_ids

get_stored_concept_ids(
    *,
    model_name: str,
    concept_filter: EmbeddingConceptFilter | None = None,
) -> set[int]

Return the stored concept IDs satisfying concept_filter, or all of them without one.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
concept_filter EmbeddingConceptFilter

Filter constraints to evaluate.

None

Returns:

Type Description
set[int]

get_concept_filter_metadata

get_concept_filter_metadata(
    *, model_name: str, concept_ids: Sequence[int]
) -> Mapping[int, ConceptEmbeddingRecord]

Fetch the stored concept metadata for the given concept IDs.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required
concept_ids Sequence[int]

OMOP concept IDs to look up. IDs that are not stored are omitted.

required

Returns:

Type Description
Mapping[int, ConceptEmbeddingRecord]

get_embedding_count

get_embedding_count(*, model_name: str) -> int

Return the number of embeddings stored in the table.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required

Returns:

Type Description
int

get_embedding_count_by_vocabulary

get_embedding_count_by_vocabulary(
    *, model_name: str
) -> Mapping[str, int]

Return the number of stored embeddings, grouped by vocabulary_id.

Parameters:

Name Type Description Default
model_name str

Canonical model name including tag.

required

Returns:

Type Description
Mapping[str, int]

vocabulary_id to embedding count, for every vocabulary with at least one stored embedding.

EmbeddingStoreReader

Bases: Protocol

Read-only view of an embedding store, as returned by open_vector_store_reader().

SQLiteVecEmbeddingBackend

Bases: EmbeddingBackend[Table]

sqlite-vec embedding backend.

All data lives in a single .db file. vec0 virtual tables (embeddings) and the registry table share the same file and SQLAlchemy engine.

Notes
  • Only the FLAT index type is supported. Each model is locked to one metric (L2, cosine or L1, cosine by default), declared on its vec0 table and recorded in the registry. rebuild_index recreates the table from the stored vectors when the metric differs or the table was created without a declared metric.
  • vec0 stores rows in fixed-size chunks and does not fill the free slots of an earlier chunk, so deleted or rewritten rows leave space that VACUUM alone does not return. compact_storage recreates the table to return it.
  • vec0 is a SQLite virtual table type the SQLAlchemy ORM cannot map, but it accepts ordinary SQL once created. _table_cache therefore stores a Core :class:~sqlalchemy.Table used for every query except the initial CREATE VIRTUAL TABLE.

physical_indexes

physical_indexes(model_name: str) -> tuple[str, ...]

vec0 is always a flat scan; there is no secondary physical index.

EmbeddingPopulationPlan dataclass

A repeatable identity-aware plan shared by coverage and execution.

prunable_ids property

prunable_ids: frozenset[int]

Stored IDs that are stale or outside the scope.

PopulationScope dataclass

The exact CDM predicate used by a population plan.

concept_filter

concept_filter() -> ConceptFilter

Return the canonical omop-alchemy filter for this scope.

VocabularyPopulationPlan dataclass

Identity comparison for one vocabulary under a population scope.

parse_backend_type

parse_backend_type(value: str | BackendType) -> BackendType

Parse a string or BackendType into a BackendType.

Parameters:

Name Type Description Default
value str | BackendType

Backend identifier string or enum member.

required

Returns:

Type Description
BackendType

Raises:

Type Description
ValueError

If value is not a recognised backend type.

parse_index_type

parse_index_type(value: str | IndexType) -> IndexType

Parse a string or IndexType into an IndexType.

Parameters:

Name Type Description Default
value str | IndexType

Index type identifier string or enum member.

required

Returns:

Type Description
IndexType

Raises:

Type Description
ValueError

If value is not a recognised index type.

parse_metric_type

parse_metric_type(value: str | MetricType) -> MetricType

Parse a string or MetricType into a MetricType.

Parameters:

Name Type Description Default
value str | MetricType

Metric type identifier string or enum member.

required

Returns:

Type Description
MetricType

Raises:

Type Description
ValueError

If value is not a recognised metric type.

index_config_from_index_type

index_config_from_index_type(
    index_type: IndexType, **kwargs: Any
) -> IndexConfig

Build an IndexConfig for a given index type from raw kwargs.

Parameters:

Name Type Description Default
index_type IndexType

Target index type.

required
**kwargs Any

Parameters forwarded to the concrete config class. Unknown keys are silently ignored.

{}

Returns:

Type Description
IndexConfig

A concrete FlatIndexConfig or HNSWIndexConfig.

Raises:

Type Description
ValueError

If index_type has no registered IndexConfig subclass.

open_vector_store_reader

open_vector_store_reader(
    resolved: ResolvedVectorStore,
) -> EmbeddingStoreReader

Open a resolved vector store read-only: no registration, no DDL, no writes.

Notes
  • Double protection against writes:
    1. @writes methods raise ReadOnlyStoreError at the Python level.
    2. The database connection itself is opened in read-only mode.
  • A store that's never been set up reports an empty registry instead of raising.

open_vector_store_writer

open_vector_store_writer(
    resolved: ResolvedVectorStore,
) -> EmbeddingBackend

Open a resolved vector store for reading and writing.

Registers the registry's schema claim and ensures its table. Reserve this for a CLI/entry-point boundary; library code should receive an EmbeddingBackend as a parameter instead.

plan_population

plan_population(
    cdm_session_factory: CDMSessionFactory,
    store: EmbeddingStoreReader,
    *,
    model_name: str,
    scope: PopulationScope = PopulationScope(),
    batch_size: int = 10000,
) -> EmbeddingPopulationPlan

Compare current CDM concept identities with stored vector identities.

CDM and vector rows are streamed. Only the stored identity map and the ID sets returned in the final plan are retained. Stored source text is not retained by the current embedding schema, so source-text changes are explicitly reported as unavailable.

omop_llm.EmbeddingRole

Bases: StrEnum

Role of text being embedded, for models with asymmetric prefixes.