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 |
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
|
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'
|
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 |
None
|
Returns:
| Type | Description |
|---|---|
EmbeddingArray
|
Float32 array of shape |
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]
|
|
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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
( |
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
|
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. |
to_dict
to_dict() -> dict[str, Any]
Serialise this config to a plain dict suitable for JSON storage.
Returns:
| Type | Description |
|---|---|
dict[str, Any]
|
|
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 |
{}
|
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: |
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 |
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 |
metric_type |
MetricType
|
Distance metric locked to this index. Required; cannot be |
num_neighbors |
int
|
Number of bi-directional links per node ( |
ef_search |
int
|
Size of the dynamic candidate list during search. Default |
ef_construction |
int
|
Size of the dynamic candidate list during index construction.
Default |
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. |
provider_type |
str
|
omop-llm provider key that serves the model (e.g. |
index_config |
IndexConfig
|
Active index configuration. |
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 |
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 |
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 |
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. |
required |
Returns:
| Type | Description |
|---|---|
str
|
Normalised name (e.g. |
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 |
required |
model_name
|
str
|
Canonical provider model ID. |
required |
Returns:
| Type | Description |
|---|---|
str
|
Store-scoped name |
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 |
require_active |
bool
|
When |
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 |
concept_name |
(str, optional)
|
Human-readable concept name. |
domain_id |
(str, optional)
|
OMOP domain (e.g. |
vocabulary_id |
(str, optional)
|
Source vocabulary (e.g. |
is_standard |
(bool, optional)
|
|
is_active |
(bool, optional)
|
|
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. |
vocabulary_id |
str
|
Source vocabulary (e.g. |
is_standard |
bool
|
Value derived from |
is_valid |
bool
|
Value derived from |
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 |
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 |
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 |
None
|
metadata
|
Mapping[str, object]
|
Free-form operational metadata. Must not contain reserved keys
(see |
None
|
registered_at
|
datetime
|
Backdate |
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 |
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 |
required |
vacuum
|
bool
|
Run the backend's maintenance step after the rebuild. pgvector
runs |
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 |
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 |
required |
embeddings
|
EmbeddingArray
|
Array of shape |
required |
defer_index_update
|
bool
|
Skip the index update for this call. Call |
False
|
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
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 |
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 |
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 |
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 |
DEFAULT_K_NEAREST
|
Returns:
| Type | Description |
|---|---|
tuple[tuple[NearestConceptMatch, ...], ...]
|
Shape |
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
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]
|
|
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
FLATindex 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_indexrecreates 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
VACUUMalone does not return.compact_storagerecreates 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_cachetherefore stores a Core :class:~sqlalchemy.Tableused for every query except the initialCREATE 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 |
parse_index_type
parse_index_type(value: str | IndexType) -> IndexType
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 |
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 |
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
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:
- @writes methods raise ReadOnlyStoreError at the Python level.
- 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.