From 1c9dd3a8534bb97330d8fdfdad528290246deb5a Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 10 Jul 2026 03:20:46 +0000 Subject: [PATCH 1/6] Centroid embedding and similar concepts function --- src/omop_emb/interface.py | 125 ++++++++++++++-- tests/test_interface_concept_similarity.py | 162 +++++++++++++++++++++ 2 files changed, 272 insertions(+), 15 deletions(-) create mode 100644 tests/test_interface_concept_similarity.py diff --git a/src/omop_emb/interface.py b/src/omop_emb/interface.py index 6ff2579..42a0ad1 100644 --- a/src/omop_emb/interface.py +++ b/src/omop_emb/interface.py @@ -262,7 +262,10 @@ def get_nearest_concepts( Returns ------- Tuple[Tuple[NearestConceptMatch, ...], ...] - Shape ``(Q, ≤k)``. Enrichment fields are ``None`` if no CDM engine. + 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). Enrichment fields of each NearestConceptMatch are + ``None`` if no CDM engine was provided to the interface. """ effective_k = k or (concept_filter.limit if concept_filter else None) or self._k @@ -303,6 +306,71 @@ def get_nearest_concepts( ) return self._enrich(raw) + def get_similar_concepts( + self, + 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 + ---------- + 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. + k : int, optional + Number of nearest neighbours to return per query concept (defaults + to interface-level *k*). + concept_filter : EmbeddingConceptFilter, optional + In-DB pre-filter applied during KNN (domain, vocabulary, standard). + faiss_index_config : IndexConfig, optional + Required only if a FAISS cache is configured on this interface. + + Returns + ------- + 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 + ------ + ValueError + If *concept_ids* is empty, or any entry has no stored embedding. + """ + ids = (concept_ids,) if isinstance(concept_ids, int) else tuple(concept_ids) + if not ids: + raise ValueError("concept_ids must be non-empty.") + + stored = self.get_embeddings_by_concept_ids(ids) + missing = [cid for cid in ids if cid not in stored] + if missing: + raise ValueError(f"No stored embedding for concept_ids: {missing}") + + vectors = np.asarray([stored[cid] for cid in ids], dtype=np.float64) + effective_k = k or (concept_filter.limit if concept_filter else None) or self._k + + raw = self.get_nearest_concepts( + vectors, + concept_filter=concept_filter, + k=effective_k + 1, # +1 because we will filter out the query concept itself from results + faiss_index_config=faiss_index_config, + ) + return tuple( + tuple(m for m in matches if m.concept_id != cid)[:effective_k] + for cid, matches in zip(ids, raw) + ) + def get_nearest_concepts_from_query_texts( self, query_texts: Union[str, Tuple[str, ...], List[str]], @@ -322,7 +390,7 @@ def get_nearest_concepts_from_query_texts( embedding_role=EmbeddingRole.QUERY, ) return self.get_nearest_concepts( - query_embedding=query_embeddings, + query_embeddings, concept_filter=concept_filter, k=k, faiss_index_config=faiss_index_config, @@ -338,28 +406,55 @@ def get_embeddings_by_concept_ids( concept_ids=concept_ids, ) - def get_indexed_concept_ids( + def get_joint_embedding( self, - concept_filter: Optional[EmbeddingConceptFilter] = None, - ) -> set[int]: - """Return every stored concept_id matching *concept_filter*. + concept_ids: Tuple[int, ...], + weights: Optional[Tuple[float, ...]] = None, + ) -> np.ndarray: + """Return the (optionally weighted) centroid of stored concept embeddings. Parameters ---------- - concept_filter : EmbeddingConceptFilter, optional - Filter constraints to evaluate (domain, vocabulary, standard, - concept ID allowlist). When omitted, every stored concept_id is - returned. + 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. + weights : tuple of float, optional + Per-concept weight, same length as *concept_ids*. Defaults to an + unweighted mean. Returns ------- - set[int] + 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 + ------ + ValueError + If *concept_ids* is empty, *weights* has a mismatched length, or + any *concept_ids* entry has no stored embedding. """ - return self._backend.get_concept_ids_matching_filter( - model_name=self.canonical_model_name, - metric_type=self._metric_type, - concept_filter=concept_filter or EmbeddingConceptFilter(), + if not concept_ids: + raise ValueError("concept_ids must be non-empty.") + if weights is not None and len(weights) != len(concept_ids): + raise ValueError( + f"weights must have the same length as concept_ids " + f"({len(weights)} != {len(concept_ids)})." + ) + + vectors_by_id = self.get_embeddings_by_concept_ids(concept_ids) + missing = [cid for cid in concept_ids if cid not in vectors_by_id] + if missing: + raise ValueError(f"No stored embedding for concept_ids: {missing}") + + vectors = np.asarray( + [vectors_by_id[cid] for cid in concept_ids], dtype=np.float64 ) + return np.average(vectors, axis=0, weights=weights) # ------------------------------------------------------------------ # Concepts without embedding (requires CDM) diff --git a/tests/test_interface_concept_similarity.py b/tests/test_interface_concept_similarity.py new file mode 100644 index 0000000..a16b93d --- /dev/null +++ b/tests/test_interface_concept_similarity.py @@ -0,0 +1,162 @@ +"""Tests for EmbeddingReaderInterface.get_joint_embedding / get_similar_concepts.""" + +from unittest.mock import Mock + +import numpy as np +import pytest + +from omop_emb.config import MetricType +from omop_emb.interface import EmbeddingReaderInterface +from omop_emb.utils.embedding_utils import NearestConceptMatch + + +def _make_backend() -> Mock: + backend = Mock() + backend.backend_name = "pgvector" + backend.get_registered_model.return_value = None + return backend + + +def _make_interface(backend: Mock, metric_type: MetricType) -> EmbeddingReaderInterface: + return EmbeddingReaderInterface( + model="test-model:v1", + backend=backend, + metric_type=metric_type, + ) + + +@pytest.mark.unit +class TestGetJointEmbedding: + def test_unweighted_mean(self): + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = { + 1: [1.0, 0.0], + 2: [0.0, 1.0], + } + interface = _make_interface(backend, MetricType.L2) + + centroid = interface.get_joint_embedding((1, 2)) + + np.testing.assert_allclose(centroid, [0.5, 0.5]) + + def test_weighted_mean(self): + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = { + 1: [1.0, 0.0], + 2: [0.0, 1.0], + } + interface = _make_interface(backend, MetricType.L2) + + centroid = interface.get_joint_embedding((1, 2), weights=(3.0, 1.0)) + + np.testing.assert_allclose(centroid, [0.75, 0.25]) + + def test_not_normalised_for_cosine_metric(self): + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = { + 1: [3.0, 0.0], + 2: [1.0, 0.0], + } + interface = _make_interface(backend, MetricType.COSINE) + + centroid = interface.get_joint_embedding((1, 2)) + + # Backends compute cosine distance directly from raw vectors, so the + # centroid is left as the plain mean regardless of metric_type. + np.testing.assert_allclose(centroid, [2.0, 0.0]) + + def test_empty_concept_ids_raises(self): + interface = _make_interface(_make_backend(), MetricType.L2) + + with pytest.raises(ValueError, match="non-empty"): + interface.get_joint_embedding(()) + + def test_mismatched_weights_length_raises(self): + interface = _make_interface(_make_backend(), MetricType.L2) + + with pytest.raises(ValueError, match="same length"): + interface.get_joint_embedding((1, 2), weights=(1.0,)) + + def test_missing_embedding_raises(self): + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = {1: [1.0, 0.0]} + interface = _make_interface(backend, MetricType.L2) + + with pytest.raises(ValueError, match="No stored embedding"): + interface.get_joint_embedding((1, 2)) + + +@pytest.mark.unit +class TestGetSimilarConcepts: + def test_bare_int_returns_single_row_excluding_self(self): + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = {1: [1.0, 0.0]} + backend.get_nearest_concepts.return_value = ( + ( + NearestConceptMatch(concept_id=1, similarity=1.0), + NearestConceptMatch(concept_id=2, similarity=0.9), + NearestConceptMatch(concept_id=3, similarity=0.8), + ), + ) + interface = _make_interface(backend, MetricType.COSINE) + + result = interface.get_similar_concepts(1, k=2) + + assert len(result) == 1 + assert [m.concept_id for m in result[0]] == [2, 3] + # requests k+1 so the self-match can be dropped without losing a slot + assert backend.get_nearest_concepts.call_args.kwargs["k"] == 3 + + def test_truncates_to_k_when_no_self_match_present(self): + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = {1: [1.0, 0.0]} + backend.get_nearest_concepts.return_value = ( + ( + NearestConceptMatch(concept_id=2, similarity=0.9), + NearestConceptMatch(concept_id=3, similarity=0.8), + NearestConceptMatch(concept_id=4, similarity=0.7), + ), + ) + interface = _make_interface(backend, MetricType.COSINE) + + result = interface.get_similar_concepts(1, k=2) + + assert [m.concept_id for m in result[0]] == [2, 3] + + def test_sequence_of_ids_returns_one_row_per_id_in_order(self): + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = { + 1: [1.0, 0.0], + 2: [0.0, 1.0], + } + backend.get_nearest_concepts.return_value = ( + ( + NearestConceptMatch(concept_id=1, similarity=1.0), + NearestConceptMatch(concept_id=3, similarity=0.5), + ), + ( + NearestConceptMatch(concept_id=2, similarity=1.0), + NearestConceptMatch(concept_id=4, similarity=0.4), + ), + ) + interface = _make_interface(backend, MetricType.COSINE) + + result = interface.get_similar_concepts((1, 2), k=1) + + assert len(result) == 2 + assert [m.concept_id for m in result[0]] == [3] + assert [m.concept_id for m in result[1]] == [4] + + def test_empty_sequence_raises(self): + interface = _make_interface(_make_backend(), MetricType.COSINE) + + with pytest.raises(ValueError, match="non-empty"): + interface.get_similar_concepts(()) + + def test_missing_embedding_raises(self): + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = {1: [1.0, 0.0]} + interface = _make_interface(backend, MetricType.COSINE) + + with pytest.raises(ValueError, match="No stored embedding"): + interface.get_similar_concepts((1, 2), k=1) From 20c5d8027dc13ed4b214d8491bb7059d6f13bc55 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 16 Jul 2026 02:07:06 +0000 Subject: [PATCH 2/6] Fix weights devision by 0 --- src/omop_emb/interface.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/omop_emb/interface.py b/src/omop_emb/interface.py index 42a0ad1..e6998a3 100644 --- a/src/omop_emb/interface.py +++ b/src/omop_emb/interface.py @@ -435,8 +435,9 @@ def get_joint_embedding( Raises ------ ValueError - If *concept_ids* is empty, *weights* has a mismatched length, or - any *concept_ids* entry has no stored embedding. + If *concept_ids* is empty, *weights* has a mismatched length, + *weights* sums to zero, or any *concept_ids* entry has no stored + embedding. """ if not concept_ids: raise ValueError("concept_ids must be non-empty.") @@ -445,6 +446,8 @@ def get_joint_embedding( f"weights must have the same length as concept_ids " f"({len(weights)} != {len(concept_ids)})." ) + if weights is not None and sum(weights) == 0: + raise ValueError("weights sum to zero; cannot compute a weighted average.") vectors_by_id = self.get_embeddings_by_concept_ids(concept_ids) missing = [cid for cid in concept_ids if cid not in vectors_by_id] From d39a72386c0a2c8dd7f1fa6f2f0c8ad1fec47d3e Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 16 Jul 2026 02:08:12 +0000 Subject: [PATCH 3/6] Docs and adapt the limit vs k obfuscation using new CDMConceptFilter --- docs/usage/interface-guide.md | 61 ++++++++++++++++++- src/omop_emb/__init__.py | 2 + src/omop_emb/backends/pgvector/pg_sql.py | 2 - .../backends/sqlitevec/sqlitevec_sql.py | 2 - src/omop_emb/cli/cli_embeddings.py | 9 ++- src/omop_emb/interface.py | 22 ++++--- src/omop_emb/utils/cdm.py | 8 +-- src/omop_emb/utils/embedding_utils.py | 56 +++++++++++++++-- tests/shared_backend_tests.py | 21 +++++-- tests/test_cdm_concept_filter.py | 39 ++++++++++++ 10 files changed, 193 insertions(+), 29 deletions(-) create mode 100644 tests/test_cdm_concept_filter.py diff --git a/docs/usage/interface-guide.md b/docs/usage/interface-guide.md index e5391f9..cbe06d3 100644 --- a/docs/usage/interface-guide.md +++ b/docs/usage/interface-guide.md @@ -151,6 +151,35 @@ results = reader.get_nearest_concepts( # results: tuple[tuple[NearestConceptMatch, ...], ...] — one inner tuple per query row ``` +### Query similar concepts + +Find neighbours of concepts that are already embedded, without fetching and +passing a raw vector yourself. The query concept is excluded from its own +result row. + +```python +results = reader.get_similar_concepts( + concept_ids=(201826, 320128), + k=5, + concept_filter=EmbeddingConceptFilter(require_standard=True), +) +# results: tuple[tuple[NearestConceptMatch, ...], ...] — one row per concept_id, same order +``` + +### Combine embeddings (joint/centroid queries) + +Build a single query vector from several stored concept embeddings — e.g. to +search for concepts similar to a *combination* of conditions — then feed it +into `get_nearest_concepts`. + +```python +joint_vec = reader.get_joint_embedding( + concept_ids=(201826, 320128), + weights=(0.7, 0.3), # optional; defaults to an unweighted mean +) +results = reader.get_nearest_concepts(query_embedding=joint_vec[None, :], k=10) +``` + ### Query by text ```python @@ -197,7 +226,10 @@ when `faiss_cache_dir` is not passed directly. `EmbeddingConceptFilter` is an in-database pre-filter applied during KNN search. All filtering happens before the nearest-neighbour step — only matching concepts -are candidates. +are candidates. It controls only *which* concepts are eligible; it never controls +*how many* results come back, as that's always controlled by `k` (see +[Query nearest concepts](#query-nearest-concepts)). For plain CDM queries (no KNN +involved), use [`CDMConceptFilter`](#cdmconceptfilter) instead. ```python from omop_emb.utils.embedding_utils import EmbeddingConceptFilter @@ -208,7 +240,6 @@ concept_filter = EmbeddingConceptFilter( concept_ids=(313217, 4329847), # restrict to specific concept IDs require_standard=True, # standard_concept = 'S' or 'C' require_active=True, # invalid_reason NOT IN ('D', 'U') - limit=20, # cap on results returned ) ``` @@ -218,6 +249,32 @@ primary backend — no CDM round-trip at query time. --- +## CDMConceptFilter + +`CDMConceptFilter` is the equivalent filter for plain queries against the OMOP +CDM `concept` table. It should not be used for embeddings or KNN search. Used by +`get_concepts_without_embedding`, `count_concepts_without_embedding`, and +`get_concepts_without_embedding_batched` on `EmbeddingWriterInterface`. Unlike +`EmbeddingConceptFilter`, it carries its own `limit`, since there's no separate +`k`-style parameter for these CDM-only queries. + +```python +from omop_emb.utils.embedding_utils import CDMConceptFilter + +concept_filter = CDMConceptFilter( + domains=("Condition", "Observation"), + require_standard=True, + limit=1000, # cap on CDM rows returned; unrelated to any KNN k +) + +n_missing = embedding_writer.count_concepts_without_embedding( + omop_cdm_engine=cdm_engine, + concept_filter=concept_filter, +) +``` + +--- + ## EmbeddingClient and providers `EmbeddingClient` wraps any OpenAI-compatible endpoint. It canonicalises the diff --git a/src/omop_emb/__init__.py b/src/omop_emb/__init__.py index 1ef2deb..7ff71ce 100644 --- a/src/omop_emb/__init__.py +++ b/src/omop_emb/__init__.py @@ -28,6 +28,7 @@ RegistryManager, ) from omop_emb.utils.embedding_utils import ( + CDMConceptFilter, EmbeddingConceptFilter, NearestConceptMatch, ) @@ -59,6 +60,7 @@ "EmbeddingModelRecord", "RegistryManager", "EmbeddingConceptFilter", + "CDMConceptFilter", "NearestConceptMatch", "ConceptEmbeddingRecord", "EmbeddingBackend", diff --git a/src/omop_emb/backends/pgvector/pg_sql.py b/src/omop_emb/backends/pgvector/pg_sql.py index 12cbcf2..7d64cb9 100644 --- a/src/omop_emb/backends/pgvector/pg_sql.py +++ b/src/omop_emb/backends/pgvector/pg_sql.py @@ -224,8 +224,6 @@ def q_nearest_concept_ids( inner_stmt = apply_concept_filter_where( inner_stmt, sa_inspect(embedding_table).columns, concept_filter ) - if concept_filter.limit is not None: - inner_stmt = inner_stmt.limit(concept_filter.limit) lateral_subq = inner_stmt.lateral("top_k") diff --git a/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py b/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py index d85f765..73865e5 100644 --- a/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py +++ b/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py @@ -258,8 +258,6 @@ def query_knn( if concept_filter is not None: setup_concept_filter_temps(session, concept_filter, "sqlite") stmt = apply_concept_filter_where(stmt, table.c, concept_filter) - if concept_filter.limit is not None: - stmt = stmt.limit(concept_filter.limit) rows = session.execute(stmt).all() return [(int(row[0]), float(row[1]), int(row[2])) for row in rows] diff --git a/src/omop_emb/cli/cli_embeddings.py b/src/omop_emb/cli/cli_embeddings.py index 47e2eea..5bb666e 100644 --- a/src/omop_emb/cli/cli_embeddings.py +++ b/src/omop_emb/cli/cli_embeddings.py @@ -20,7 +20,11 @@ ) from omop_emb.embeddings import EmbeddingClient from omop_emb.interface import EmbeddingReaderInterface, EmbeddingWriterInterface -from omop_emb.utils.embedding_utils import EmbeddingConceptFilter, NearestConceptMatch +from omop_emb.utils.embedding_utils import ( + CDMConceptFilter, + EmbeddingConceptFilter, + NearestConceptMatch, +) logger = logging.getLogger(__name__) app = typer.Typer( @@ -196,7 +200,7 @@ def add_embeddings( embedding_writer.register_model() # Filter concepts - concept_filter = EmbeddingConceptFilter( + concept_filter = CDMConceptFilter( require_standard=standard_only, domains=tuple(domains) if domains else None, vocabularies=tuple(vocabularies) if vocabularies else None, @@ -650,7 +654,6 @@ def search( require_standard=standard_only, domains=tuple(domains) if domains else None, vocabularies=tuple(vocabularies) if vocabularies else None, - limit=k, ) for batch_id, batched_queries in enumerate( diff --git a/src/omop_emb/interface.py b/src/omop_emb/interface.py index e6998a3..fc80c36 100644 --- a/src/omop_emb/interface.py +++ b/src/omop_emb/interface.py @@ -50,7 +50,11 @@ ) from omop_emb.backends.index_config import IndexConfig from omop_emb.config import BackendType, MetricType, ProviderType -from omop_emb.utils.embedding_utils import EmbeddingConceptFilter, NearestConceptMatch +from omop_emb.utils.embedding_utils import ( + CDMConceptFilter, + EmbeddingConceptFilter, + NearestConceptMatch, +) if TYPE_CHECKING: from omop_emb.storage.faiss import FAISSCache @@ -240,6 +244,10 @@ def get_embedding_count_by_vocabulary(self) -> Mapping[str, int]: # Search # ------------------------------------------------------------------ + def _resolve_effective_k(self, k: Optional[int]) -> int: + """Resolve the number of nearest neighbours to request.""" + return k or self._k + def get_nearest_concepts( self, query_embedding: np.ndarray, @@ -267,7 +275,7 @@ def get_nearest_concepts( exist at all). Enrichment fields of each NearestConceptMatch are ``None`` if no CDM engine was provided to the interface. """ - effective_k = k or (concept_filter.limit if concept_filter else None) or self._k + effective_k = self._resolve_effective_k(k) if self._faiss_cache is not None: if faiss_index_config is None: @@ -358,7 +366,7 @@ def get_similar_concepts( raise ValueError(f"No stored embedding for concept_ids: {missing}") vectors = np.asarray([stored[cid] for cid in ids], dtype=np.float64) - effective_k = k or (concept_filter.limit if concept_filter else None) or self._k + effective_k = self._resolve_effective_k(k) raw = self.get_nearest_concepts( vectors, @@ -467,7 +475,7 @@ def get_concepts_without_embedding( self, omop_cdm_engine: Engine, *, - concept_filter: Optional[EmbeddingConceptFilter] = None, + concept_filter: Optional[CDMConceptFilter] = None, ) -> Mapping[int, Row]: """Return CDM rows for concepts lacking embeddings, keyed by concept_id. @@ -491,7 +499,7 @@ def count_concepts_without_embedding( self, omop_cdm_engine: Engine, *, - concept_filter: Optional[EmbeddingConceptFilter] = None, + concept_filter: Optional[CDMConceptFilter] = None, ) -> int: """Return how many CDM concepts match *concept_filter* but lack an embedding.""" embedded_ids = self._backend.get_all_stored_concept_ids( @@ -505,7 +513,7 @@ def get_concepts_without_embedding_batched( omop_cdm_engine: Engine, *, batch_size: int, - concept_filter: Optional[EmbeddingConceptFilter] = None, + concept_filter: Optional[CDMConceptFilter] = None, limit: Optional[int] = None, ) -> Iterable[Mapping[int, Row]]: """Yield ``{concept_id: Row}`` batches for concepts lacking embeddings. @@ -554,7 +562,7 @@ def _enrich( return raw unique_ids = {r.concept_id for results in raw for r in results} - concept_filter = EmbeddingConceptFilter(concept_ids=tuple(unique_ids)) + concept_filter = CDMConceptFilter(concept_ids=tuple(unique_ids)) rows = fetch_cdm_concepts_for_filter( concept_filter=concept_filter, cdm_engine=self._cdm_engine ) diff --git a/src/omop_emb/utils/cdm.py b/src/omop_emb/utils/cdm.py index f9c4168..9bf1e5b 100644 --- a/src/omop_emb/utils/cdm.py +++ b/src/omop_emb/utils/cdm.py @@ -11,7 +11,7 @@ from sqlalchemy.orm import Session, sessionmaker from omop_alchemy.cdm.model.vocabulary import Concept -from omop_emb.utils.embedding_utils import EmbeddingConceptFilter +from omop_emb.utils.embedding_utils import CDMConceptFilter logger = logging.getLogger(__name__) @@ -44,7 +44,7 @@ def check_concept_cdm(cdm_engine: Engine) -> None: def fetch_cdm_concepts_for_filter( - concept_filter: Optional[EmbeddingConceptFilter], + concept_filter: Optional[CDMConceptFilter], cdm_engine: Engine, ) -> dict[int, Row]: """Return CDM rows matching *concept_filter*, keyed by concept_id. @@ -68,7 +68,7 @@ def fetch_cdm_concepts_for_filter( def iter_cdm_concepts_for_filter( - concept_filter: Optional[EmbeddingConceptFilter], + concept_filter: Optional[CDMConceptFilter], cdm_engine: Engine, chunk_size: int = 5_000, ) -> Iterator[Row]: @@ -95,7 +95,7 @@ def iter_cdm_concepts_for_filter( def count_missing_concepts( - concept_filter: Optional[EmbeddingConceptFilter], + concept_filter: Optional[CDMConceptFilter], cdm_engine: Engine, embedded_ids: set[int], chunk_size: int = 10_000, diff --git a/src/omop_emb/utils/embedding_utils.py b/src/omop_emb/utils/embedding_utils.py index 06c1d67..d3ef3f0 100644 --- a/src/omop_emb/utils/embedding_utils.py +++ b/src/omop_emb/utils/embedding_utils.py @@ -20,14 +20,60 @@ class EmbeddingConceptFilter: """Search constraints applied during KNN retrieval. - All fields are optional. Unset fields impose no constraint. ``limit`` - maps directly to the ``k`` nearest neighbours returned. + All fields are optional. Unset fields impose no constraint. This filter + controls only *which* concepts are eligible candidates; it never controls + *how many* results come back. Pass ``k`` to + :meth:`EmbeddingReaderInterface.get_nearest_concepts` (or similar) for that. + For CDM-only queries (e.g. ``get_concepts_without_embedding``), use + :class:`CDMConceptFilter` instead. Notes ----- Mirrors OMOP grounding needs without importing ``omop_graph`` or its search-constraint types into ``omop_emb``. + Attributes + ---------- + 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 standard concepts (``standard_concept`` in + ``('S', 'C')``) are returned. Default ``False``. + require_active : bool + When ``True``, only active concepts (``invalid_reason`` not in + ``('D', 'U')``) are returned. Default ``False``. + """ + + concept_ids: Optional[tuple[int, ...]] = None + domains: Optional[tuple[str, ...]] = None + vocabularies: Optional[tuple[str, ...]] = None + require_standard: bool = False + require_active: bool = False + + def is_empty(self) -> bool: + """Return ``True`` if no constraints are set.""" + return ( + self.concept_ids is None + and self.domains is None + and self.vocabularies is None + and not self.require_standard + and not self.require_active + ) + + +@dataclass(frozen=True) +class CDMConceptFilter: + """Search constraints applied to plain CDM ``concept``-table queries. + + All fields are optional. Unset fields impose no constraint. Distinct from + :class:`EmbeddingConceptFilter`: this filter is for CDM-only queries (e.g. + ``get_concepts_without_embedding``, ``count_concepts_without_embedding``), + not KNN search, and ``limit`` caps the number of CDM rows returned. + Attributes ---------- concept_ids : tuple[int, ...], optional @@ -43,8 +89,8 @@ class EmbeddingConceptFilter: When ``True``, only active concepts (``invalid_reason`` not in ``('D', 'U')``) are returned. Default ``False``. limit : int, optional - Maximum number of nearest neighbours to return. If not set, the - backend default is used. + Maximum number of CDM rows to return. If not set, all matching rows + are returned. """ concept_ids: Optional[tuple[int, ...]] = None @@ -57,7 +103,7 @@ class EmbeddingConceptFilter: def __post_init__(self) -> None: if self.limit is not None and self.limit <= 0: raise ValueError( - f"EmbeddingConceptFilter.limit must be a positive integer, got {self.limit}." + f"CDMConceptFilter.limit must be a positive integer, got {self.limit}." ) def apply(self, query: Select, table: type) -> Select: diff --git a/tests/shared_backend_tests.py b/tests/shared_backend_tests.py index 1a2ce4b..a8f0056 100644 --- a/tests/shared_backend_tests.py +++ b/tests/shared_backend_tests.py @@ -240,7 +240,7 @@ def test_knn_domain_filter(self, backend: EmbeddingBackend): model_name=MODEL_NAME, metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, - concept_filter=EmbeddingConceptFilter(domains=("Drug",), limit=10), + concept_filter=EmbeddingConceptFilter(domains=("Drug",)), ) returned_ids = {r.concept_id for r in results[0]} assert ASPIRIN_ID in returned_ids @@ -253,7 +253,7 @@ def test_knn_vocabulary_filter(self, backend: EmbeddingBackend): model_name=MODEL_NAME, metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, - concept_filter=EmbeddingConceptFilter(vocabularies=("SNOMED",), limit=10), + concept_filter=EmbeddingConceptFilter(vocabularies=("SNOMED",)), ) returned_ids = {r.concept_id for r in results[0]} assert HYPERTENSION_ID in returned_ids @@ -267,7 +267,7 @@ def test_knn_concept_id_filter(self, backend: EmbeddingBackend): metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, concept_filter=EmbeddingConceptFilter( - concept_ids=(HYPERTENSION_ID, ASPIRIN_ID), limit=10 + concept_ids=(HYPERTENSION_ID, ASPIRIN_ID) ), ) returned_ids = {r.concept_id for r in results[0]} @@ -279,12 +279,25 @@ def test_knn_require_standard_filter(self, backend: EmbeddingBackend): model_name=MODEL_NAME, metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, - concept_filter=EmbeddingConceptFilter(require_standard=True, limit=10), + concept_filter=EmbeddingConceptFilter(require_standard=True), ) returned_ids = {r.concept_id for r in results[0]} assert NON_STANDARD_ID not in returned_ids assert HYPERTENSION_ID in returned_ids + def test_knn_k_controls_result_count_with_filter(self, backend: EmbeddingBackend): + """k is the sole result-count control; EmbeddingConceptFilter has no such field.""" + self._upsert_all(backend) + results = backend.get_nearest_concepts( + model_name=MODEL_NAME, + metric_type=MetricType.L2, + query_embeddings=QUERY_EMBEDDING, + k=1, + concept_filter=EmbeddingConceptFilter(vocabularies=("SNOMED",)), + ) + assert len(results[0]) == 1 + assert results[0][0].concept_id == DIABETES_ID # nearest SNOMED concept by L2 + # ------------------------------------------------------------------ # KNN — similarity math # ------------------------------------------------------------------ diff --git a/tests/test_cdm_concept_filter.py b/tests/test_cdm_concept_filter.py new file mode 100644 index 0000000..a7dfb96 --- /dev/null +++ b/tests/test_cdm_concept_filter.py @@ -0,0 +1,39 @@ +"""Tests for CDMConceptFilter.apply() — the CDM-only WHERE/LIMIT builder.""" + +import pytest +from sqlalchemy import select + +from omop_alchemy.cdm.model.vocabulary import Concept +from omop_emb.utils.embedding_utils import CDMConceptFilter + + +@pytest.mark.unit +class TestCDMConceptFilterApply: + def test_empty_filter_adds_no_clauses(self): + query = select(Concept.concept_id) + result = CDMConceptFilter().apply(query, Concept) + + assert str(result) == str(query) + + def test_concept_ids_adds_in_clause(self): + query = select(Concept.concept_id) + result = CDMConceptFilter(concept_ids=(1, 2, 3)).apply(query, Concept) + + compiled = str(result) + assert "WHERE" in compiled + assert "concept_id IN" in compiled + + def test_limit_is_applied(self): + query = select(Concept.concept_id) + result = CDMConceptFilter(limit=5).apply(query, Concept) + + assert "LIMIT" in str(result) + + def test_negative_limit_raises(self): + with pytest.raises(ValueError, match="positive integer"): + CDMConceptFilter(limit=0) + + def test_is_empty(self): + assert CDMConceptFilter().is_empty() + assert not CDMConceptFilter(limit=5).is_empty() + assert not CDMConceptFilter(domains=("Drug",)).is_empty() From c7cb1b9f18d7c77502d0f313ce8311002bed2da2 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 16 Jul 2026 02:13:49 +0000 Subject: [PATCH 4/6] Leave explicit note in the CDMConceptFilter --- src/omop_emb/utils/embedding_utils.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/omop_emb/utils/embedding_utils.py b/src/omop_emb/utils/embedding_utils.py index d3ef3f0..16f8683 100644 --- a/src/omop_emb/utils/embedding_utils.py +++ b/src/omop_emb/utils/embedding_utils.py @@ -74,6 +74,14 @@ class CDMConceptFilter: ``get_concepts_without_embedding``, ``count_concepts_without_embedding``), not KNN search, and ``limit`` caps the number of CDM rows returned. + Notes + ----- + Mirrors omop_graph.graph.constraints.SearchConstraintConcept as we cannot + import omop-graph into omop_emb. This issues is being noted here: + https://github.com/AustralianCancerDataNetwork/OMOP_Alchemy/issues/11 + Once that is solved, this can be removed again and imported from omop_alchemy + + Attributes ---------- concept_ids : tuple[int, ...], optional From 9865dd0dbb73b80fa7cfef68816cebd309e5c014 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 21 Jul 2026 04:12:24 +0000 Subject: [PATCH 5/6] Revert EmbeddingConceptFilter changes, make it backwards-compatible --- docs/usage/interface-guide.md | 41 +++-------- src/omop_emb/__init__.py | 2 - src/omop_emb/cli/cli_embeddings.py | 3 +- src/omop_emb/interface.py | 51 +++++++++++--- src/omop_emb/utils/cdm.py | 8 +-- src/omop_emb/utils/embedding_utils.py | 74 ++++--------------- tests/shared_backend_tests.py | 16 +++-- tests/test_cdm_concept_filter.py | 39 ---------- tests/test_interface_concept_similarity.py | 82 +++++++++++++++++++++- 9 files changed, 159 insertions(+), 157 deletions(-) delete mode 100644 tests/test_cdm_concept_filter.py diff --git a/docs/usage/interface-guide.md b/docs/usage/interface-guide.md index cbe06d3..bb94d4c 100644 --- a/docs/usage/interface-guide.md +++ b/docs/usage/interface-guide.md @@ -224,12 +224,8 @@ when `faiss_cache_dir` is not passed directly. ## EmbeddingConceptFilter -`EmbeddingConceptFilter` is an in-database pre-filter applied during KNN search. -All filtering happens before the nearest-neighbour step — only matching concepts -are candidates. It controls only *which* concepts are eligible; it never controls -*how many* results come back, as that's always controlled by `k` (see -[Query nearest concepts](#query-nearest-concepts)). For plain CDM queries (no KNN -involved), use [`CDMConceptFilter`](#cdmconceptfilter) instead. +`EmbeddingConceptFilter` is an in-database pre-filter used both during KNN search +and for plain CDM queries (e.g. `get_concepts_without_embedding`). ```python from omop_emb.utils.embedding_utils import EmbeddingConceptFilter @@ -240,6 +236,7 @@ concept_filter = EmbeddingConceptFilter( concept_ids=(313217, 4329847), # restrict to specific concept IDs require_standard=True, # standard_concept = 'S' or 'C' require_active=True, # invalid_reason NOT IN ('D', 'U') + limit=20, # cap on results returned ) ``` @@ -247,31 +244,13 @@ All fields are optional and combinable. `require_standard` and `require_active` are stored as columns in the embedding table and are resolved entirely inside the primary backend — no CDM round-trip at query time. ---- - -## CDMConceptFilter - -`CDMConceptFilter` is the equivalent filter for plain queries against the OMOP -CDM `concept` table. It should not be used for embeddings or KNN search. Used by -`get_concepts_without_embedding`, `count_concepts_without_embedding`, and -`get_concepts_without_embedding_batched` on `EmbeddingWriterInterface`. Unlike -`EmbeddingConceptFilter`, it carries its own `limit`, since there's no separate -`k`-style parameter for these CDM-only queries. - -```python -from omop_emb.utils.embedding_utils import CDMConceptFilter - -concept_filter = CDMConceptFilter( - domains=("Condition", "Observation"), - require_standard=True, - limit=1000, # cap on CDM rows returned; unrelated to any KNN k -) - -n_missing = embedding_writer.count_concepts_without_embedding( - omop_cdm_engine=cdm_engine, - concept_filter=concept_filter, -) -``` +!!! warning "`limit` is deprecated for KNN search" + For [`get_nearest_concepts`](#query-nearest-concepts) and similar KNN methods, + pass `k` instead of setting `limit` on the filter. `limit` is still honored + as a fallback when `k` isn't given, and validated for consistency when both + are given, but will stop affecting KNN result count in 2.0. `limit` remains + fully supported (and is not deprecated) for CDM-only queries like + `get_concepts_without_embedding`. --- diff --git a/src/omop_emb/__init__.py b/src/omop_emb/__init__.py index 7ff71ce..1ef2deb 100644 --- a/src/omop_emb/__init__.py +++ b/src/omop_emb/__init__.py @@ -28,7 +28,6 @@ RegistryManager, ) from omop_emb.utils.embedding_utils import ( - CDMConceptFilter, EmbeddingConceptFilter, NearestConceptMatch, ) @@ -60,7 +59,6 @@ "EmbeddingModelRecord", "RegistryManager", "EmbeddingConceptFilter", - "CDMConceptFilter", "NearestConceptMatch", "ConceptEmbeddingRecord", "EmbeddingBackend", diff --git a/src/omop_emb/cli/cli_embeddings.py b/src/omop_emb/cli/cli_embeddings.py index 5bb666e..019cdff 100644 --- a/src/omop_emb/cli/cli_embeddings.py +++ b/src/omop_emb/cli/cli_embeddings.py @@ -21,7 +21,6 @@ from omop_emb.embeddings import EmbeddingClient from omop_emb.interface import EmbeddingReaderInterface, EmbeddingWriterInterface from omop_emb.utils.embedding_utils import ( - CDMConceptFilter, EmbeddingConceptFilter, NearestConceptMatch, ) @@ -200,7 +199,7 @@ def add_embeddings( embedding_writer.register_model() # Filter concepts - concept_filter = CDMConceptFilter( + concept_filter = EmbeddingConceptFilter( require_standard=standard_only, domains=tuple(domains) if domains else None, vocabularies=tuple(vocabularies) if vocabularies else None, diff --git a/src/omop_emb/interface.py b/src/omop_emb/interface.py index fc80c36..dfca96e 100644 --- a/src/omop_emb/interface.py +++ b/src/omop_emb/interface.py @@ -17,6 +17,7 @@ from __future__ import annotations import logging +import warnings from dataclasses import replace as dc_replace from typing import ( TYPE_CHECKING, @@ -51,7 +52,6 @@ from omop_emb.backends.index_config import IndexConfig from omop_emb.config import BackendType, MetricType, ProviderType from omop_emb.utils.embedding_utils import ( - CDMConceptFilter, EmbeddingConceptFilter, NearestConceptMatch, ) @@ -244,9 +244,34 @@ def get_embedding_count_by_vocabulary(self) -> Mapping[str, int]: # Search # ------------------------------------------------------------------ - def _resolve_effective_k(self, k: Optional[int]) -> int: - """Resolve the number of nearest neighbours to request.""" - return k or self._k + def _resolve_effective_k( + self, k: Optional[int], concept_filter: Optional[EmbeddingConceptFilter] + ) -> int: + """Resolve the number of nearest neighbours to request. + + ``concept_filter.limit`` is deprecated as a KNN result-count control + (see :class:`EmbeddingConceptFilter`'s docstring) — this is a + transitional shim, not the long-term contract. It's still honored as + a fallback when *k* isn't given, and validated for consistency when + both are given, until removal in 2.0. + """ + filter_limit = concept_filter.limit if concept_filter is not None else None + if filter_limit is not None: + warnings.warn( + "EmbeddingConceptFilter.limit is deprecated for KNN search and will be " + "removed in 2.0. Pass k explicitly instead. It is still honored as a " + "fallback when k is not given, and validated for consistency when both " + "are given, until then.", + DeprecationWarning, + stacklevel=2, + ) + if k is not None and k != filter_limit: + raise ValueError( + f"k={k} and concept_filter.limit={filter_limit} were both given and " + f"disagree. concept_filter.limit for KNN search is deprecated. Pass " + f"k alone." + ) + return k or filter_limit or self._k def get_nearest_concepts( self, @@ -275,7 +300,7 @@ def get_nearest_concepts( exist at all). Enrichment fields of each NearestConceptMatch are ``None`` if no CDM engine was provided to the interface. """ - effective_k = self._resolve_effective_k(k) + effective_k = self._resolve_effective_k(k, concept_filter) if self._faiss_cache is not None: if faiss_index_config is None: @@ -366,11 +391,15 @@ def get_similar_concepts( raise ValueError(f"No stored embedding for concept_ids: {missing}") vectors = np.asarray([stored[cid] for cid in ids], dtype=np.float64) - effective_k = self._resolve_effective_k(k) + effective_k = self._resolve_effective_k(k, concept_filter) + # Prevent ValueError due to mismatch being raised due to effective_k + 1 below + inner_filter = ( + dc_replace(concept_filter, limit=None) if concept_filter is not None else None + ) raw = self.get_nearest_concepts( vectors, - concept_filter=concept_filter, + concept_filter=inner_filter, k=effective_k + 1, # +1 because we will filter out the query concept itself from results faiss_index_config=faiss_index_config, ) @@ -475,7 +504,7 @@ def get_concepts_without_embedding( self, omop_cdm_engine: Engine, *, - concept_filter: Optional[CDMConceptFilter] = None, + concept_filter: Optional[EmbeddingConceptFilter] = None, ) -> Mapping[int, Row]: """Return CDM rows for concepts lacking embeddings, keyed by concept_id. @@ -499,7 +528,7 @@ def count_concepts_without_embedding( self, omop_cdm_engine: Engine, *, - concept_filter: Optional[CDMConceptFilter] = None, + concept_filter: Optional[EmbeddingConceptFilter] = None, ) -> int: """Return how many CDM concepts match *concept_filter* but lack an embedding.""" embedded_ids = self._backend.get_all_stored_concept_ids( @@ -513,7 +542,7 @@ def get_concepts_without_embedding_batched( omop_cdm_engine: Engine, *, batch_size: int, - concept_filter: Optional[CDMConceptFilter] = None, + concept_filter: Optional[EmbeddingConceptFilter] = None, limit: Optional[int] = None, ) -> Iterable[Mapping[int, Row]]: """Yield ``{concept_id: Row}`` batches for concepts lacking embeddings. @@ -562,7 +591,7 @@ def _enrich( return raw unique_ids = {r.concept_id for results in raw for r in results} - concept_filter = CDMConceptFilter(concept_ids=tuple(unique_ids)) + concept_filter = EmbeddingConceptFilter(concept_ids=tuple(unique_ids)) rows = fetch_cdm_concepts_for_filter( concept_filter=concept_filter, cdm_engine=self._cdm_engine ) diff --git a/src/omop_emb/utils/cdm.py b/src/omop_emb/utils/cdm.py index 9bf1e5b..f9c4168 100644 --- a/src/omop_emb/utils/cdm.py +++ b/src/omop_emb/utils/cdm.py @@ -11,7 +11,7 @@ from sqlalchemy.orm import Session, sessionmaker from omop_alchemy.cdm.model.vocabulary import Concept -from omop_emb.utils.embedding_utils import CDMConceptFilter +from omop_emb.utils.embedding_utils import EmbeddingConceptFilter logger = logging.getLogger(__name__) @@ -44,7 +44,7 @@ def check_concept_cdm(cdm_engine: Engine) -> None: def fetch_cdm_concepts_for_filter( - concept_filter: Optional[CDMConceptFilter], + concept_filter: Optional[EmbeddingConceptFilter], cdm_engine: Engine, ) -> dict[int, Row]: """Return CDM rows matching *concept_filter*, keyed by concept_id. @@ -68,7 +68,7 @@ def fetch_cdm_concepts_for_filter( def iter_cdm_concepts_for_filter( - concept_filter: Optional[CDMConceptFilter], + concept_filter: Optional[EmbeddingConceptFilter], cdm_engine: Engine, chunk_size: int = 5_000, ) -> Iterator[Row]: @@ -95,7 +95,7 @@ def iter_cdm_concepts_for_filter( def count_missing_concepts( - concept_filter: Optional[CDMConceptFilter], + concept_filter: Optional[EmbeddingConceptFilter], cdm_engine: Engine, embedded_ids: set[int], chunk_size: int = 10_000, diff --git a/src/omop_emb/utils/embedding_utils.py b/src/omop_emb/utils/embedding_utils.py index 16f8683..31d44a3 100644 --- a/src/omop_emb/utils/embedding_utils.py +++ b/src/omop_emb/utils/embedding_utils.py @@ -18,70 +18,15 @@ @dataclass(frozen=True) class EmbeddingConceptFilter: - """Search constraints applied during KNN retrieval. + """Search constraints applied during KNN retrieval and plain CDM queries. - All fields are optional. Unset fields impose no constraint. This filter - controls only *which* concepts are eligible candidates; it never controls - *how many* results come back. Pass ``k`` to - :meth:`EmbeddingReaderInterface.get_nearest_concepts` (or similar) for that. - For CDM-only queries (e.g. ``get_concepts_without_embedding``), use - :class:`CDMConceptFilter` instead. + All fields are optional. Unset fields impose no constraint. Notes ----- Mirrors OMOP grounding needs without importing ``omop_graph`` or its search-constraint types into ``omop_emb``. - Attributes - ---------- - 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 standard concepts (``standard_concept`` in - ``('S', 'C')``) are returned. Default ``False``. - require_active : bool - When ``True``, only active concepts (``invalid_reason`` not in - ``('D', 'U')``) are returned. Default ``False``. - """ - - concept_ids: Optional[tuple[int, ...]] = None - domains: Optional[tuple[str, ...]] = None - vocabularies: Optional[tuple[str, ...]] = None - require_standard: bool = False - require_active: bool = False - - def is_empty(self) -> bool: - """Return ``True`` if no constraints are set.""" - return ( - self.concept_ids is None - and self.domains is None - and self.vocabularies is None - and not self.require_standard - and not self.require_active - ) - - -@dataclass(frozen=True) -class CDMConceptFilter: - """Search constraints applied to plain CDM ``concept``-table queries. - - All fields are optional. Unset fields impose no constraint. Distinct from - :class:`EmbeddingConceptFilter`: this filter is for CDM-only queries (e.g. - ``get_concepts_without_embedding``, ``count_concepts_without_embedding``), - not KNN search, and ``limit`` caps the number of CDM rows returned. - - Notes - ----- - Mirrors omop_graph.graph.constraints.SearchConstraintConcept as we cannot - import omop-graph into omop_emb. This issues is being noted here: - https://github.com/AustralianCancerDataNetwork/OMOP_Alchemy/issues/11 - Once that is solved, this can be removed again and imported from omop_alchemy - - Attributes ---------- concept_ids : tuple[int, ...], optional @@ -97,8 +42,17 @@ class CDMConceptFilter: When ``True``, only active concepts (``invalid_reason`` not in ``('D', 'U')``) are returned. Default ``False``. limit : int, optional - Maximum number of CDM rows to return. If not set, all matching rows - are returned. + For CDM-only queries (e.g. ``get_concepts_without_embedding``), caps + the number of CDM rows returned. + + .. deprecated:: + Using ``limit`` to control KNN result count is deprecated and + will be removed in 2.0 — pass ``k`` to + :meth:`EmbeddingReaderInterface.get_nearest_concepts` (or similar) + instead. Still honored as a fallback when ``k`` is not given, and + validated for consistency when both are given, until then. This + deprecation does not apply to CDM-only queries, where ``limit`` + remains the intended way to cap row count. """ concept_ids: Optional[tuple[int, ...]] = None @@ -111,7 +65,7 @@ class CDMConceptFilter: def __post_init__(self) -> None: if self.limit is not None and self.limit <= 0: raise ValueError( - f"CDMConceptFilter.limit must be a positive integer, got {self.limit}." + f"EmbeddingConceptFilter.limit must be a positive integer, got {self.limit}." ) def apply(self, query: Select, table: type) -> Select: diff --git a/tests/shared_backend_tests.py b/tests/shared_backend_tests.py index a8f0056..75feaaa 100644 --- a/tests/shared_backend_tests.py +++ b/tests/shared_backend_tests.py @@ -240,7 +240,7 @@ def test_knn_domain_filter(self, backend: EmbeddingBackend): model_name=MODEL_NAME, metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, - concept_filter=EmbeddingConceptFilter(domains=("Drug",)), + concept_filter=EmbeddingConceptFilter(domains=("Drug",), limit=10), ) returned_ids = {r.concept_id for r in results[0]} assert ASPIRIN_ID in returned_ids @@ -253,7 +253,7 @@ def test_knn_vocabulary_filter(self, backend: EmbeddingBackend): model_name=MODEL_NAME, metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, - concept_filter=EmbeddingConceptFilter(vocabularies=("SNOMED",)), + concept_filter=EmbeddingConceptFilter(vocabularies=("SNOMED",), limit=10), ) returned_ids = {r.concept_id for r in results[0]} assert HYPERTENSION_ID in returned_ids @@ -267,7 +267,7 @@ def test_knn_concept_id_filter(self, backend: EmbeddingBackend): metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, concept_filter=EmbeddingConceptFilter( - concept_ids=(HYPERTENSION_ID, ASPIRIN_ID) + concept_ids=(HYPERTENSION_ID, ASPIRIN_ID), limit=10 ), ) returned_ids = {r.concept_id for r in results[0]} @@ -279,21 +279,23 @@ def test_knn_require_standard_filter(self, backend: EmbeddingBackend): model_name=MODEL_NAME, metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, - concept_filter=EmbeddingConceptFilter(require_standard=True), + concept_filter=EmbeddingConceptFilter(require_standard=True, limit=10), ) returned_ids = {r.concept_id for r in results[0]} assert NON_STANDARD_ID not in returned_ids assert HYPERTENSION_ID in returned_ids - def test_knn_k_controls_result_count_with_filter(self, backend: EmbeddingBackend): - """k is the sole result-count control; EmbeddingConceptFilter has no such field.""" + def test_knn_backend_ignores_concept_filter_limit(self, backend: EmbeddingBackend): + """Backend SQL builders never read concept_filter.limit — only k controls + result count at this layer (the deprecation/consistency check lives one + layer up, in EmbeddingReaderInterface, not here).""" self._upsert_all(backend) results = backend.get_nearest_concepts( model_name=MODEL_NAME, metric_type=MetricType.L2, query_embeddings=QUERY_EMBEDDING, k=1, - concept_filter=EmbeddingConceptFilter(vocabularies=("SNOMED",)), + concept_filter=EmbeddingConceptFilter(vocabularies=("SNOMED",), limit=10), ) assert len(results[0]) == 1 assert results[0][0].concept_id == DIABETES_ID # nearest SNOMED concept by L2 diff --git a/tests/test_cdm_concept_filter.py b/tests/test_cdm_concept_filter.py deleted file mode 100644 index a7dfb96..0000000 --- a/tests/test_cdm_concept_filter.py +++ /dev/null @@ -1,39 +0,0 @@ -"""Tests for CDMConceptFilter.apply() — the CDM-only WHERE/LIMIT builder.""" - -import pytest -from sqlalchemy import select - -from omop_alchemy.cdm.model.vocabulary import Concept -from omop_emb.utils.embedding_utils import CDMConceptFilter - - -@pytest.mark.unit -class TestCDMConceptFilterApply: - def test_empty_filter_adds_no_clauses(self): - query = select(Concept.concept_id) - result = CDMConceptFilter().apply(query, Concept) - - assert str(result) == str(query) - - def test_concept_ids_adds_in_clause(self): - query = select(Concept.concept_id) - result = CDMConceptFilter(concept_ids=(1, 2, 3)).apply(query, Concept) - - compiled = str(result) - assert "WHERE" in compiled - assert "concept_id IN" in compiled - - def test_limit_is_applied(self): - query = select(Concept.concept_id) - result = CDMConceptFilter(limit=5).apply(query, Concept) - - assert "LIMIT" in str(result) - - def test_negative_limit_raises(self): - with pytest.raises(ValueError, match="positive integer"): - CDMConceptFilter(limit=0) - - def test_is_empty(self): - assert CDMConceptFilter().is_empty() - assert not CDMConceptFilter(limit=5).is_empty() - assert not CDMConceptFilter(domains=("Drug",)).is_empty() diff --git a/tests/test_interface_concept_similarity.py b/tests/test_interface_concept_similarity.py index a16b93d..3692efb 100644 --- a/tests/test_interface_concept_similarity.py +++ b/tests/test_interface_concept_similarity.py @@ -1,5 +1,6 @@ """Tests for EmbeddingReaderInterface.get_joint_embedding / get_similar_concepts.""" +import warnings from unittest.mock import Mock import numpy as np @@ -7,7 +8,7 @@ from omop_emb.config import MetricType from omop_emb.interface import EmbeddingReaderInterface -from omop_emb.utils.embedding_utils import NearestConceptMatch +from omop_emb.utils.embedding_utils import EmbeddingConceptFilter, NearestConceptMatch def _make_backend() -> Mock: @@ -160,3 +161,82 @@ def test_missing_embedding_raises(self): with pytest.raises(ValueError, match="No stored embedding"): interface.get_similar_concepts((1, 2), k=1) + + +@pytest.mark.unit +class TestResolveEffectiveKDeprecation: + """concept_filter.limit is deprecated for KNN result-count control, but + still honored as a fallback (with a DeprecationWarning) until 2.0, and + validated for consistency against an explicit k.""" + + def _matches(self, backend: Mock, n: int) -> None: + backend.get_nearest_concepts.return_value = ( + tuple(NearestConceptMatch(concept_id=i, similarity=1.0) for i in range(n)), + ) + + def test_k_alone_no_warning(self): + backend = _make_backend() + self._matches(backend, 3) + interface = _make_interface(backend, MetricType.COSINE) + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + interface.get_nearest_concepts(np.zeros((1, 2)), k=3) + + assert backend.get_nearest_concepts.call_args.kwargs["k"] == 3 + + def test_filter_limit_alone_warns_and_is_honored(self): + backend = _make_backend() + self._matches(backend, 5) + interface = _make_interface(backend, MetricType.COSINE) + + with pytest.warns(DeprecationWarning, match="deprecated for KNN search"): + interface.get_nearest_concepts( + np.zeros((1, 2)), concept_filter=EmbeddingConceptFilter(limit=5) + ) + + assert backend.get_nearest_concepts.call_args.kwargs["k"] == 5 + + def test_matching_k_and_filter_limit_warns_but_does_not_raise(self): + backend = _make_backend() + self._matches(backend, 4) + interface = _make_interface(backend, MetricType.COSINE) + + with pytest.warns(DeprecationWarning): + interface.get_nearest_concepts( + np.zeros((1, 2)), k=4, concept_filter=EmbeddingConceptFilter(limit=4) + ) + + assert backend.get_nearest_concepts.call_args.kwargs["k"] == 4 + + def test_conflicting_k_and_filter_limit_raises(self): + backend = _make_backend() + interface = _make_interface(backend, MetricType.COSINE) + + with pytest.raises(ValueError, match="disagree"): + interface.get_nearest_concepts( + np.zeros((1, 2)), k=5, concept_filter=EmbeddingConceptFilter(limit=3) + ) + + def test_get_similar_concepts_inner_call_does_not_spuriously_conflict(self): + """effective_k + 1 (the over-fetch for self-match exclusion) must not be + flagged as conflicting with the same concept_filter.limit it was derived + from.""" + backend = _make_backend() + backend.get_embeddings_by_concept_ids.return_value = {1: [1.0, 0.0]} + backend.get_nearest_concepts.return_value = ( + ( + NearestConceptMatch(concept_id=1, similarity=1.0), + NearestConceptMatch(concept_id=2, similarity=0.9), + NearestConceptMatch(concept_id=3, similarity=0.8), + ), + ) + interface = _make_interface(backend, MetricType.COSINE) + + with pytest.warns(DeprecationWarning): + result = interface.get_similar_concepts( + 1, concept_filter=EmbeddingConceptFilter(limit=2) + ) + + assert [m.concept_id for m in result[0]] == [2, 3] + assert backend.get_nearest_concepts.call_args.kwargs["k"] == 3 From 769db28c0b82cf5e4e3abfc7b891434c6b30685b Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 23 Jul 2026 03:58:19 +0000 Subject: [PATCH 6/6] Restore get_indexed_concept_ids --- src/omop_emb/interface.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/omop_emb/interface.py b/src/omop_emb/interface.py index dfca96e..e1798f1 100644 --- a/src/omop_emb/interface.py +++ b/src/omop_emb/interface.py @@ -443,6 +443,28 @@ def get_embeddings_by_concept_ids( concept_ids=concept_ids, ) + def get_indexed_concept_ids( + self, + concept_filter: Optional[EmbeddingConceptFilter] = None, + ) -> set[int]: + """Return every stored concept_id matching *concept_filter*. + + .. deprecated:: + This method is deprecated and will be removed in 2.0. + Prefer backend- or workflow-specific alternatives instead. + """ + warnings.warn( + "EmbeddingReaderInterface.get_indexed_concept_ids() is deprecated and " + "will be removed in 2.0.", + DeprecationWarning, + stacklevel=2, + ) + return self._backend.get_concept_ids_matching_filter( + model_name=self.canonical_model_name, + metric_type=self._metric_type, + concept_filter=concept_filter or EmbeddingConceptFilter(), + ) + def get_joint_embedding( self, concept_ids: Tuple[int, ...],