Add optional context parameter to embedding-based grounding

src/omop_graph/reasoning/grounding.py

@@ -87,13 +87,26 @@ def __post_init__(self) -> None:
             )
 
 
+def _query_text_with_context(query: str, context: Optional[str]) -> str:
+    """Combine query with optional disambiguating context for embedding.
+
+    No length guard here by design. omop-emb's EmbeddingClient.embeddings()
+    wraps provider API errors (including oversized-input failures) in a clear
+    EmbeddingClientError instead of letting a raw provider error propagate.
+    """
+    if not context:
+        return query
+    return f"{query}\n\n{context}"
+
+
 def ground_term(
     resolver_pipeline: ResolverPipeline,
     kg: KnowledgeGraph,
     query: str,
     query_embedding: Optional[np.ndarray],
     constraints: GroundingConstraints,
     max_candidates: Optional[int] = None,
+    context: Optional[str] = None,
 ) -> List[StandardConceptWithScore]:
     """
     Ground a text string to a ranked list of standard OMOP concepts.
@@ -113,6 +126,13 @@ def ground_term(
         Contextual constraints (parents, domains, etc.) to apply.
     max_candidates : int, optional
         Limit for the number of candidates returned. If None, returns all candidates.
+    context : str, optional
+        Auxiliary disambiguating context (e.g. surrounding text, expected
+        domain, Q&A context) folded into the text used for the on-demand
+        query embedding. 
+        Has no effect when *query_embedding* is already supplied by the
+        caller, since no embedding is computed in that case. Does not affect
+        candidate resolution or structural filtering (``constraints``).
 
     Returns
     -------
@@ -138,7 +158,7 @@ def ground_term(
             from omop_emb.embeddings import EmbeddingRole
 
             query_embedding = embedding_writer.embed_texts(
-                texts=(query,),
+                texts=(_query_text_with_context(query, context),),
                 embedding_role=EmbeddingRole.QUERY,
             )
 

tests/test_grounding.py

@@ -1,13 +1,21 @@
 from __future__ import annotations
 
+from unittest.mock import Mock
+
+import numpy as np
 import pytest
 
 from omop_graph.extensions.omop_alchemy import PredicateKind
 from omop_graph.graph.constraints import SearchConstraintConcept
 from omop_graph.graph.kg import KnowledgeGraph
-from omop_graph.reasoning.grounding import GroundingConstraints, ground_term
+from omop_graph.reasoning.grounding import (
+    GroundingConstraints,
+    _query_text_with_context,
+    ground_term,
+)
 from omop_graph.reasoning.resolvers.resolver_pipeline import ResolverPipeline
 from omop_graph.reasoning.resolvers.resolvers import (
+    EmbeddingResolver,
     ExactLabelResolver,
     ExactSynonymResolver,
     PartialLabelResolver,
@@ -152,3 +160,112 @@ def test_grounding_standard_candidate_without_parent_ids_is_zero_hop(
     assert ranked[0].concept_id == 196653
     assert ranked[0].identity_hops == 0
     assert ranked[0].separation == 0
+
+
+class TestQueryTextWithContext:
+    """_query_text_with_context: plain concatenation, no context is a no-op."""
+
+    def test_no_context_returns_query_unchanged(self):
+        assert _query_text_with_context("EGFR", None) == "EGFR"
+
+    def test_empty_context_returns_query_unchanged(self):
+        assert _query_text_with_context("EGFR", "") == "EGFR"
+
+    def test_context_is_appended(self):
+        result = _query_text_with_context("EGFR", "genomic mutation panel")
+        assert result == "EGFR\n\ngenomic mutation panel"
+
+
+def test_ground_term_folds_context_into_on_demand_embedding_text(
+    mock_cdm_kg: KnowledgeGraph,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """context is concatenated into the text passed to embed_texts when the
+    query embedding is computed on demand (query_embedding=None)."""
+    fake_writer = Mock()
+    fake_writer.embed_texts.return_value = np.zeros((1, 8), dtype=np.float32)
+    monkeypatch.setattr(
+        "omop_graph.reasoning.grounding.get_embedding_writer_interface",
+        lambda kg: fake_writer,
+    )
+    # Avoid needing a real embedding-backed nearest-neighbour lookup — the
+    # resolution outcome is irrelevant to this test, only the embed_texts call.
+    monkeypatch.setattr(
+        "omop_graph.reasoning.resolvers.resolvers.get_neareast_concepts",
+        lambda **kwargs: None,
+    )
+
+    pipeline = ResolverPipeline(resolvers=(EmbeddingResolver(),))
+
+    ground_term(
+        resolver_pipeline=pipeline,
+        kg=mock_cdm_kg,
+        query="EGFR",
+        query_embedding=None,
+        constraints=_unconstrained_constraints(),
+        context="genomic mutation panel",
+    )
+
+    fake_writer.embed_texts.assert_called_once()
+    assert fake_writer.embed_texts.call_args.kwargs["texts"] == (
+        "EGFR\n\ngenomic mutation panel",
+    )
+
+
+def test_ground_term_omits_context_separator_when_not_given(
+    mock_cdm_kg: KnowledgeGraph,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fake_writer = Mock()
+    fake_writer.embed_texts.return_value = np.zeros((1, 8), dtype=np.float32)
+    monkeypatch.setattr(
+        "omop_graph.reasoning.grounding.get_embedding_writer_interface",
+        lambda kg: fake_writer,
+    )
+    monkeypatch.setattr(
+        "omop_graph.reasoning.resolvers.resolvers.get_neareast_concepts",
+        lambda **kwargs: None,
+    )
+
+    pipeline = ResolverPipeline(resolvers=(EmbeddingResolver(),))
+
+    ground_term(
+        resolver_pipeline=pipeline,
+        kg=mock_cdm_kg,
+        query="EGFR",
+        query_embedding=None,
+        constraints=_unconstrained_constraints(),
+    )
+
+    assert fake_writer.embed_texts.call_args.kwargs["texts"] == ("EGFR",)
+
+
+def test_ground_term_context_has_no_effect_when_query_embedding_supplied(
+    mock_cdm_kg: KnowledgeGraph,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """context only feeds the on-demand embedding path; a caller-supplied
+    query_embedding means no embedding is computed, so embed_texts is never
+    called at all."""
+    fake_writer = Mock()
+    monkeypatch.setattr(
+        "omop_graph.reasoning.grounding.get_embedding_writer_interface",
+        lambda kg: fake_writer,
+    )
+    monkeypatch.setattr(
+        "omop_graph.reasoning.resolvers.resolvers.get_neareast_concepts",
+        lambda **kwargs: None,
+    )
+
+    pipeline = ResolverPipeline(resolvers=(EmbeddingResolver(),))
+
+    ground_term(
+        resolver_pipeline=pipeline,
+        kg=mock_cdm_kg,
+        query="EGFR",
+        query_embedding=np.zeros((1, 8), dtype=np.float32),
+        constraints=_unconstrained_constraints(),
+        context="genomic mutation panel",
+    )
+
+    fake_writer.embed_texts.assert_not_called()