From f277198e0dfa08ef6621268327eaf0b0bb6b7ea1 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 10 Jul 2026 05:53:01 +0000 Subject: [PATCH 1/2] Provide simple additional context to the grounding pipeline --- src/omop_graph/reasoning/grounding.py | 17 +++- tests/test_grounding.py | 119 +++++++++++++++++++++++++- 2 files changed, 134 insertions(+), 2 deletions(-) diff --git a/src/omop_graph/reasoning/grounding.py b/src/omop_graph/reasoning/grounding.py index 689b9fb..a674632 100644 --- a/src/omop_graph/reasoning/grounding.py +++ b/src/omop_graph/reasoning/grounding.py @@ -87,6 +87,13 @@ def __post_init__(self) -> None: ) +def _query_text_with_context(query: str, context: Optional[str]) -> str: + """Combine query with optional disambiguating context for embedding.""" + if not context: + return query + return f"{query}\n\n{context}" + + def ground_term( resolver_pipeline: ResolverPipeline, kg: KnowledgeGraph, @@ -94,6 +101,7 @@ def ground_term( 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 +121,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 +153,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, ) diff --git a/tests/test_grounding.py b/tests/test_grounding.py index 9400d6f..46da72d 100644 --- a/tests/test_grounding.py +++ b/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() From 45001737aeba4922a4f174e37a4f360b3c050106 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 16 Jul 2026 04:51:00 +0000 Subject: [PATCH 2/2] Leave note about the resolving of context length --- src/omop_graph/reasoning/grounding.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/omop_graph/reasoning/grounding.py b/src/omop_graph/reasoning/grounding.py index a674632..df5a9b6 100644 --- a/src/omop_graph/reasoning/grounding.py +++ b/src/omop_graph/reasoning/grounding.py @@ -88,7 +88,12 @@ def __post_init__(self) -> None: def _query_text_with_context(query: str, context: Optional[str]) -> str: - """Combine query with optional disambiguating context for embedding.""" + """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}"