Add joint/centroid embedding and concept-to-concept similarity queries

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
@@ -195,9 +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.
+`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
@@ -216,6 +244,14 @@ 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.
 
+!!! 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`.
+
 ---
 
 ## EmbeddingClient and providers

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

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]

src/omop_emb/cli/cli_embeddings.py

@@ -20,7 +20,10 @@
 )
 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 (
+    EmbeddingConceptFilter,
+    NearestConceptMatch,
+)
 
 logger = logging.getLogger(__name__)
 app = typer.Typer(
@@ -650,7 +653,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(

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,
@@ -50,7 +51,10 @@
 )
 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 (
+    EmbeddingConceptFilter,
+    NearestConceptMatch,
+)
 
 if TYPE_CHECKING:
     from omop_emb.storage.faiss import FAISSCache
@@ -240,6 +244,35 @@ def get_embedding_count_by_vocabulary(self) -> Mapping[str, int]:
     # Search
     # ------------------------------------------------------------------
 
+    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,
         query_embedding: np.ndarray,
@@ -262,9 +295,12 @@ 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
+        effective_k = self._resolve_effective_k(k, concept_filter)
 
         if self._faiss_cache is not None:
             if faiss_index_config is None:
@@ -303,6 +339,75 @@ 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 = 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=inner_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 +427,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,
@@ -344,22 +449,74 @@ def get_indexed_concept_ids(
     ) -> 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, ...],
+        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,
+            *weights* sums to zero, 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)})."
+            )
+        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]
+        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)

src/omop_emb/utils/embedding_utils.py

@@ -18,10 +18,9 @@
 
 @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. ``limit``
-    maps directly to the ``k`` nearest neighbours returned.
+    All fields are optional. Unset fields impose no constraint.
 
     Notes
     -----
@@ -43,8 +42,17 @@ 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.
+        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

tests/shared_backend_tests.py

@@ -285,6 +285,21 @@ def test_knn_require_standard_filter(self, backend: EmbeddingBackend):
         assert NON_STANDARD_ID not in returned_ids
         assert HYPERTENSION_ID in returned_ids
 
+    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",), limit=10),
+        )
+        assert len(results[0]) == 1
+        assert results[0][0].concept_id == DIABETES_ID  # nearest SNOMED concept by L2
+
     # ------------------------------------------------------------------
     # KNN — similarity math
     # ------------------------------------------------------------------