class OntologyPatchRetriever(Tool):
"""Combines vector retrieval into one composite ontology graph."""
vector_store: VectorStoreManager = Field(exclude=True)
sparql_tool: Any | None = Field(default=None, exclude=True)
# Typed ``Any`` for the same reason as ``sparql_tool``: OntologyManager holds a
# back-reference to this class, so a concrete annotation would be a cycle.
ontology_manager: Any | None = Field(default=None, exclude=True)
patch: PatchRetrievalConfig = Field(
default_factory=PatchRetrievalConfig,
exclude=True,
)
_last_retrieval_metrics: dict[str, Any] = PrivateAttr(default_factory=dict)
_surface_index: CatalogSurfaceIndex | None = PrivateAttr(default=None)
# Whole-module graphs for the small-module closure, keyed by ontology IRI.
# The catalog is stable for the life of a run, and every content unit hits
# the same handful of modules — refetching per unit multiplies catalog
# reads by the unit count for no new information. ``None`` caches a miss.
_small_module_cache: dict[str, Ontology | None] = PrivateAttr(default_factory=dict)
# Tenancy the cache was filled under. The retriever outlives a tenancy
# switch, and serving one tenant's modules to another would be a leak.
_small_module_cache_scope: str = PrivateAttr(default="")
@property
def last_retrieval_metrics(self) -> dict[str, Any]:
return self._last_retrieval_metrics
def _match_query_unit_signals(self, trigger_source: str) -> dict[str, str]:
"""Match number-adjacent unit tokens against catalog surface forms.
Additive, outside the semantic atom budget (precedent: the lexical
trigger lane). Returns ``{entity_iri: ontology_iri}``; empty when the
lane is disabled or nothing matches.
"""
if not self.vector_store.store_config.query_unit_signals_enabled:
return {}
manager = self.ontology_manager
if manager is None or not trigger_source:
return {}
tokens = number_adjacent_tokens(trigger_source)
if not tokens:
return {}
if self._surface_index is None:
# Symbol/notation predicates come from configuration rather than
# being compiled into query_signals; built lazily because the
# store config is not available at PrivateAttr default time.
self._surface_index = CatalogSurfaceIndex(
symbol_predicates=[
URIRef(iri)
for iri in (
self.vector_store.store_config.induced_subgraph_symbol_predicates
)
]
)
matched = self._surface_index.match(tokens, manager.ontologies)
if matched:
logger.info(
"Query unit signals matched %d entity(ies) from tokens %s",
len(matched),
sorted(tokens),
)
return matched
@staticmethod
def _schema_axiom_graph(
merged_context: tuple[RDFGraph, dict[str, str]] | None,
catalog: list[Ontology] | None,
) -> RDFGraph | None:
"""Pick the graph to read ``rdfs:domain``/``rdfs:range`` axioms from.
Whichever of the two expansion paths materialized the ontologies wins;
neither being available means the induced-subgraph call is fetching on
its own and there is nothing local to close over.
"""
if merged_context is not None:
return merged_context[0]
if catalog:
combined = RDFGraph()
for ontology in catalog:
combined += ontology.graph
return combined
return None
async def _apply_small_module_closure(
self, graph: RDFGraph, hit_ontology_iris: list[str]
) -> None:
"""Merge whole small modules into the snapshot (header-stripped).
A vocabulary small enough to fit entirely (e.g. a qualified-quantity
module of ~20 terms) is included wholesale once any of its atoms is
admitted: partial inclusion of a tiny module is what pushes the
renderer to improvise near-miss property names.
"""
closure_max = self.patch.small_module_closure_max_triples
if closure_max <= 0:
return
modules = await self._asmall_module_candidates(hit_ontology_iris)
closed: list[str] = []
for onto_iri, ontology in modules:
if len(ontology.graph) > closure_max:
continue
module_graph = Ontology.strip_ontology_header_triples(ontology.graph.copy())
_drop_module_contribution(graph, module_graph)
graph += module_graph
for prefix, namespace_uri in ontology.graph.namespaces():
graph.bind(prefix, namespace_uri)
closed.append(onto_iri)
if closed:
self._last_retrieval_metrics["module_closure_iris"] = closed
async def _asmall_module_candidates(
self, hit_ontology_iris: list[str]
) -> list[tuple[str, Ontology]]:
"""Resolve hit ontologies to full graphs, manager first, store second.
The in-memory manager is empty in every deployment that keeps its
catalog in a triple store and fetches per query — which is the normal
server configuration, and where this closure silently did nothing.
"""
store_config = getattr(self.vector_store, "store_config", None)
scope = str(getattr(store_config, "ontology_table", "") or "")
if scope != self._small_module_cache_scope:
self._small_module_cache.clear()
self._small_module_cache_scope = scope
wanted = sorted(set(hit_ontology_iris))
resolved: list[tuple[str, Ontology]] = []
missing: list[str] = []
manager = self.ontology_manager
for onto_iri in wanted:
if onto_iri in self._small_module_cache:
cached = self._small_module_cache[onto_iri]
if cached is not None:
resolved.append((onto_iri, cached))
continue
ontology = (
manager.get_freshest_terminal_ontology_by_iri(onto_iri)
if manager is not None
else None
)
if ontology is None or ontology.is_null():
missing.append(onto_iri)
else:
self._small_module_cache[onto_iri] = ontology
resolved.append((onto_iri, ontology))
store = self.sparql_tool.triple_store_manager if self.sparql_tool else None
if missing and store is not None:
try:
fetched = await store.afetch_ontologies_by_iri(missing)
except Exception as exc:
# Do NOT cache on this path. A None entry is a permanent
# negative (see the miss-caching note above), so memoizing a
# transient store error would silently strip the small-module
# closure from every later unit in the process.
logger.warning(
"Small-module closure catalog fetch failed (not cached, "
"will retry on the next unit): %s",
exc,
)
return sorted(resolved, key=lambda item: item[0])
by_iri = {o.iri: o for o in fetched if o.iri and not o.is_null()}
for onto_iri in missing:
found = by_iri.get(onto_iri)
self._small_module_cache[onto_iri] = found
if found is not None:
resolved.append((onto_iri, found))
return sorted(resolved, key=lambda item: item[0])
async def _acandidate_context(
self,
*,
entity_uris: list[str],
ontology_iris: list[str],
ontology_version_filters: dict[str, set[str]] | None,
ontology_hash_filters: dict[str, set[str]] | None,
depth: int,
) -> tuple[RDFGraph, dict[str, str]]:
"""Build the working graph from a CONSTRUCT instead of merging catalogs.
Version and hash filters are applied to the *headers*, so the CONSTRUCT is
restricted to exactly the named graphs the merge path would have selected.
Prefix bindings cannot come from a CONSTRUCT, so they are rebuilt from the
catalog's author-prefix table; standard vocabulary prefixes are bound
downstream by :func:`_bind_common_vocab_prefixes` as on the merge path.
Returns:
tuple: ``(candidate_graph, prefix_map)``.
"""
manager = self.ontology_manager
assert manager is not None and self.sparql_tool is not None
store = self.sparql_tool.triple_store_manager
headers = select_relevant_ontologies(
dedupe_terminal_ontologies(await manager.aget_catalog_headers()),
ontology_iris,
ontology_version_filters,
ontology_hash_filters,
)
if not headers:
return RDFGraph(), {}
graph_irefs = _sparql_irefs([header.graph_uri for header in headers])
seed_irefs = _sparql_irefs(entity_uris)
if not graph_irefs or not seed_irefs:
return RDFGraph(), {}
candidate = RDFGraph()
for chunk in _chunked(seed_irefs, _MAX_VALUES_TERMS):
partial = await store.aconstruct(
build_candidate_subgraph_query(chunk, graph_irefs, depth=depth)
)
candidate += partial
prefix_map: dict[str, str] = {}
for header in headers:
namespace = str(header.namespace)
prefix = manager.author_prefix_for_namespace(namespace)
if prefix:
prefix_map[prefix] = namespace
prefix_map = filter_overbroad_namespace_map(prefix_map)
for prefix, namespace in prefix_map.items():
candidate.bind(prefix, Namespace(namespace))
# Mirror the merge path: author @prefix names persisted as sh:declare
# triples (pulled by the candidate CONSTRUCT's header branch) win over
# stem-derived recovery, exactly as ontology_from_named_graph binds them
# for merged catalog graphs.
declared = candidate.bind_declared_prefixes()
known_before_declared = set(prefix_map.values())
for namespace, prefix in declared.items():
if namespace not in known_before_declared:
prefix_map[prefix] = namespace
# Graphs served from a triple store carry no author @prefix bindings, so
# stem-derived prefixes fill any remaining gap (see
# ontology_from_named_graph). Recover the same implicit stems here so
# both context paths advertise identical namespaces.
candidate.bind_implicit_namespaces()
known_namespaces = set(prefix_map.values())
for prefix, namespace_uri in candidate.namespaces():
ns = str(namespace_uri)
if (
not prefix
or ns in known_namespaces
or ns in RDFLIB_DEFAULT_NAMESPACE_URIS
):
continue
prefix_map[prefix] = ns
return candidate, prefix_map
async def _aresolve_merged_context(
self,
*,
entity_uris: list[str],
ontology_iris: list[str],
catalog: list[Ontology] | None,
ontology_version_filters: dict[str, set[str]] | None,
ontology_hash_filters: dict[str, set[str]] | None,
depth: int,
candidate_pushdown: bool,
) -> tuple[RDFGraph, dict[str, str]] | None:
"""Resolve the merged ontology context through the catalog, or ``None``.
Returning ``None`` leaves the induced-subgraph call on its own fetch path,
which is what happens when no catalog is registered or a read fails.
``catalog`` being set means the reference-expansion fallback already
materialized everything, so there is nothing left to save here.
Args:
ontology_iris: Ontology IRIs surviving reference expansion.
catalog: Ontologies already materialized by the fallback path, if any.
ontology_version_filters: Allowed versions per ontology IRI.
ontology_hash_filters: Allowed hashes per ontology IRI.
Returns:
tuple | None: ``(merged_graph, prefix_map)``, or ``None`` to fall back.
"""
manager = self.ontology_manager
if manager is None or catalog is not None:
return None
store = self.sparql_tool.triple_store_manager if self.sparql_tool else None
use_pushdown = (
candidate_pushdown
and store is not None
and store.supports_sparql_construct()
)
try:
if use_pushdown:
merged = await self._acandidate_context(
entity_uris=entity_uris,
ontology_iris=ontology_iris,
ontology_version_filters=ontology_version_filters,
ontology_hash_filters=ontology_hash_filters,
depth=depth,
)
mode = "sparql_candidate"
else:
selected = select_relevant_ontologies(
await manager.aget_ontologies_by_iri(ontology_iris),
ontology_iris,
ontology_version_filters,
ontology_hash_filters,
)
merged = await manager.aget_merged_graph(selected)
mode = "merged_catalog"
except Exception as exc:
logger.warning(
"Catalog context via OntologyManager failed (%s); "
"falling back to a direct triple-store read",
exc,
)
return None
self._last_retrieval_metrics.update(manager.catalog_cache_stats())
self._last_retrieval_metrics["catalog_context_mode"] = mode
self._last_retrieval_metrics["catalog_context_triples"] = len(merged[0])
return merged
def _effective_top_k(self, top_k: int | None) -> int:
if top_k is not None:
return top_k
return self.vector_store.store_config.top_k
def _resolve_subgraph_budget(
self,
subgraph_depth: int | None,
max_total_triples: int | None,
estimated_triples_per_query: int | None,
) -> tuple[int, int, int]:
"""Fill unset induced-subgraph budget arguments from configuration."""
sc = self.vector_store.store_config
return (
sc.induced_subgraph_depth if subgraph_depth is None else subgraph_depth,
(
sc.induced_subgraph_max_total_triples
if max_total_triples is None
else max_total_triples
),
(
sc.induced_subgraph_estimated_triples_per_query
if estimated_triples_per_query is None
else estimated_triples_per_query
),
)
def retrieve(
self,
query: str,
top_k: int | None = None,
expand_sparql: bool = True,
subgraph_depth: int | None = None,
max_total_triples: int | None = None,
estimated_triples_per_query: int | None = None,
) -> tuple[RDFGraph, list[str]]:
"""Retrieve top-k hits for one query and optional induced subgraph; returns source ontology IRIs."""
try:
asyncio.get_running_loop()
except RuntimeError:
return asyncio.run(
self.aretrieve(
query=query,
top_k=top_k,
expand_sparql=expand_sparql,
subgraph_depth=subgraph_depth,
max_total_triples=max_total_triples,
estimated_triples_per_query=estimated_triples_per_query,
)
)
raise RuntimeError(
"retrieve() cannot be called from async code; use await aretrieve()"
)
def retrieve_ensemble(
self,
queries: list[str],
top_k: int | None = None,
expand_sparql: bool = True,
subgraph_depth: int | None = None,
max_total_triples: int | None = None,
estimated_triples_per_query: int | None = None,
trigger_text: str | None = None,
) -> tuple[RDFGraph, list[str]]:
"""Sync: one induced graph and source IRIs for the union of vector hits over ``queries``."""
try:
asyncio.get_running_loop()
except RuntimeError:
return asyncio.run(
self.aretrieve_ensemble(
queries=queries,
top_k=top_k,
expand_sparql=expand_sparql,
subgraph_depth=subgraph_depth,
max_total_triples=max_total_triples,
estimated_triples_per_query=estimated_triples_per_query,
trigger_text=trigger_text,
)
)
raise RuntimeError(
"retrieve_ensemble() is not allowed inside async code; use aretrieve_ensemble()"
)
async def aretrieve(
self,
query: str,
top_k: int | None = None,
expand_sparql: bool = True,
subgraph_depth: int | None = None,
max_total_triples: int | None = None,
estimated_triples_per_query: int | None = None,
trigger_text: str | None = None,
) -> tuple[RDFGraph, list[str]]:
"""Async single-query variant of :meth:`aretrieve_ensemble`."""
return await self.aretrieve_ensemble(
queries=[query],
top_k=top_k,
expand_sparql=expand_sparql,
subgraph_depth=subgraph_depth,
max_total_triples=max_total_triples,
estimated_triples_per_query=estimated_triples_per_query,
trigger_text=trigger_text,
)
async def aretrieve_ensemble(
self,
queries: list[str],
top_k: int | None = None,
expand_sparql: bool = True,
subgraph_depth: int | None = None,
max_total_triples: int | None = None,
estimated_triples_per_query: int | None = None,
trigger_text: str | None = None,
) -> tuple[RDFGraph, list[str]]:
"""Vector search over all ``queries`` once, score-filter, dedupe, single subgraph expansion.
``subgraph_depth`` / ``max_total_triples`` / ``estimated_triples_per_query``
default to the configured values (``ONTOLOGY_PATCH_INDUCED_SUBGRAPH_*``).
They previously carried literal defaults of 1 / 300 / 24, which
contradicted the config defaults of 2 / 1200 / 24: the pipeline passed
config explicitly and was unaffected, but any other caller of this
public API silently got a 4x smaller snapshot than the deployment was
configured for.
"""
self._last_retrieval_metrics = {}
subgraph_depth, max_total_triples, estimated_triples_per_query = (
self._resolve_subgraph_budget(
subgraph_depth, max_total_triples, estimated_triples_per_query
)
)
trigger_source = (trigger_text or "").strip()
if not queries and not trigger_source:
return RDFGraph(), []
eff_top_k = self._effective_top_k(top_k)
hits_by_query: list[OntologySearchHitsByChannel] = []
if queries:
hits_by_query = await self.vector_store.asearch_patch_hits_many(
queries=queries,
top_k=eff_top_k,
)
sc = self.vector_store.store_config
pc = self.patch
eff_max_atoms = pc.effective_max_atoms(len(queries))
merged = _filter_and_merge_patch_hits(
hits_by_query,
store_config=sc,
patch_config=pc,
per_query_core_score_ratio=pc.per_query_core_score_ratio,
per_query_neighborhood_score_ratio=pc.per_query_neighborhood_score_ratio,
per_query_bm25_score_ratio=pc.per_query_bm25_score_ratio,
min_core_query_best_score=pc.min_core_query_best_score,
min_neighborhood_query_best_score=pc.min_neighborhood_query_best_score,
min_bm25_query_best_score=pc.min_bm25_query_best_score,
min_merged_max_score=pc.min_merged_max_score,
max_atoms_total=0,
)
atoms_after_dedupe = len(merged)
merged = [atom for atom in merged if not _is_ontology_declaration_atom(atom)]
if merged and pc.merged_score_ratio > 0.0:
merged_top = float(merged[0].score or 0.0)
merged_floor = merged_top * pc.merged_score_ratio
merged = [
atom for atom in merged if float(atom.score or 0.0) >= merged_floor
]
ranked_before_cut = list(merged)
if merged and pc.mmr_lambda < 1.0:
merged = _normalize_relevance_scores(merged)
vectors = await self.vector_store.afetch_vectors(
[atom.atom_id for atom in merged]
)
core_w, neigh_w = normalized_core_neighborhood_weights(sc)
merged = _mmr_rerank(
merged,
vectors,
mmr_lambda=pc.mmr_lambda,
max_atoms=eff_max_atoms,
core_weight=core_w,
neighborhood_weight=neigh_w,
)
elif pc.cross_query_merge_mode in (
CrossQueryMergeMode.MAX_SCORE,
CrossQueryMergeMode.SUM_SCORE,
):
merged = _select_atoms_round_robin_by_ontology(
merged,
per_ontology_seed_quota=pc.per_ontology_seed_quota,
max_atoms=eff_max_atoms,
per_ontology_atom_floor=pc.per_ontology_atom_floor,
per_role_atom_floor=pc.per_role_atom_floor,
)
elif eff_max_atoms > 0:
merged = merged[:eff_max_atoms]
trigger_source = trigger_source or " ".join(queries)
trigger_atoms = await asyncio.to_thread(
self.vector_store.match_lexical_triggers, trigger_source
)
merged, trigger_promoted, trigger_appended = _merge_lexical_trigger_atoms(
merged, trigger_atoms, fusion=sc.lexical_trigger_fusion
)
# After the trigger merge: an exact-case trigger hit is positive
# evidence and exempts the atom; what remains penalizable is the
# case-folded BM25/dense residue.
merged, symbol_case_penalized = _demote_case_mismatched_symbol_atoms(
merged,
trigger_source,
policy=sc.symbol_case_mismatch_policy,
demote_factor=sc.symbol_case_mismatch_demote_factor,
)
if not merged:
self._last_retrieval_metrics = {
"query_count": len(queries),
"top_k": eff_top_k,
"effective_max_atoms": eff_max_atoms,
"atoms_after_dedupe": atoms_after_dedupe,
"atoms_final": 0,
"seed_iris": [],
"lexical_trigger_hits": len(trigger_atoms),
"lexical_trigger_atom_ids": [a.atom_id for a in trigger_atoms],
"lexical_trigger_promoted": trigger_promoted,
"lexical_trigger_appended": trigger_appended,
"symbol_case_penalized": symbol_case_penalized,
}
if pc.dump_ontology_ranks:
self._last_retrieval_metrics["ontology_rank_diagnostics"] = (
build_ontology_rank_diagnostics(
hits_by_query, ranked_before_cut, []
)
)
return RDFGraph(), []
source_iris = _source_iris_from_atoms(merged)
seeds_by_ontology: dict[str, int] = defaultdict(int)
for atom in merged:
if atom.ontology_iri:
seeds_by_ontology[atom.ontology_iri] += 1
self._last_retrieval_metrics = {
"query_count": len(queries),
"top_k": eff_top_k,
"effective_max_atoms": eff_max_atoms,
"merge_mode": pc.cross_query_merge_mode.value,
"atoms_after_dedupe": atoms_after_dedupe,
"atoms_final": len(merged),
"seed_iris": [atom.iri for atom in merged if atom.iri],
"source_ontology_iris": source_iris,
"seeds_by_ontology": dict(seeds_by_ontology),
"lexical_trigger_hits": len(trigger_atoms),
"lexical_trigger_atom_ids": [a.atom_id for a in trigger_atoms],
"lexical_trigger_iris": [a.iri for a in trigger_atoms if a.iri],
"lexical_trigger_promoted": trigger_promoted,
"lexical_trigger_appended": trigger_appended,
"symbol_case_penalized": symbol_case_penalized,
}
if pc.dump_ontology_ranks:
self._last_retrieval_metrics["ontology_rank_diagnostics"] = (
build_ontology_rank_diagnostics(
hits_by_query, ranked_before_cut, merged
)
)
if not expand_sparql or self.sparql_tool is None:
return RDFGraph(), source_iris
entity_uris, entity_relevance, entity_roles = _ranked_entity_weights(merged)
signal_entities = self._match_query_unit_signals(trigger_source)
for signal_iri, signal_onto_iri in sorted(signal_entities.items()):
if signal_iri in entity_relevance:
continue
entity_uris.append(signal_iri)
entity_relevance[signal_iri] = sc.lexical_trigger_score
entity_roles[signal_iri] = "resource"
if signal_entities:
self._last_retrieval_metrics["query_signal_iris"] = sorted(
signal_entities.keys()
)
hit_ontology_iris = sorted(
{atom.ontology_iri for atom in merged if atom.ontology_iri}
| set(signal_entities.values())
)
ontology_version_filters: dict[str, set[str]] = {}
ontology_hash_filters: dict[str, set[str]] = {}
for atom in merged:
if atom.ontology_iri and atom.ontology_version:
ontology_version_filters.setdefault(atom.ontology_iri, set()).add(
str(atom.ontology_version)
)
if atom.ontology_iri and atom.ontology_hash:
ontology_hash_filters.setdefault(atom.ontology_iri, set()).add(
atom.ontology_hash
)
ontology_iris = hit_ontology_iris
catalog: list[Ontology] | None = None
triple_store_manager = self.sparql_tool.triple_store_manager
if triple_store_manager is not None:
ontology_iris, catalog, expansion_metrics = await _aexpand_ontology_iris(
triple_store_manager, entity_uris, hit_ontology_iris
)
expanded = sorted(set(ontology_iris) - set(hit_ontology_iris))
if expanded:
self._last_retrieval_metrics["expanded_ontology_iris"] = expanded
self._last_retrieval_metrics.update(expansion_metrics)
merged_context = await self._aresolve_merged_context(
entity_uris=entity_uris,
ontology_iris=ontology_iris,
catalog=catalog,
ontology_version_filters=ontology_version_filters or None,
ontology_hash_filters=ontology_hash_filters or None,
depth=subgraph_depth,
candidate_pushdown=sc.induced_subgraph_candidate_pushdown,
)
schema_graph = self._schema_axiom_graph(merged_context, catalog)
if schema_graph is not None:
closure = _schema_closure_entities(
schema_graph,
entity_uris,
max_entities=pc.schema_closure_max_entities,
ancestor_depth=pc.schema_closure_ancestor_depth,
seed_relevance=entity_relevance,
)
if closure:
closure_score = _closure_floor_score(entity_relevance)
for closure_iri, closure_role in closure.items():
entity_uris.append(closure_iri)
entity_relevance[closure_iri] = closure_score
entity_roles[closure_iri] = closure_role
self._last_retrieval_metrics["schema_closure_iris"] = sorted(closure)
hub_seed_count = sc.induced_subgraph_hub_seed_count
ancestor_depth = sc.induced_subgraph_ancestor_closure_depth
entity_groups: dict[str, str] = {
atom.iri: atom.ontology_iri
for atom in merged
if atom.iri and atom.ontology_iri
}
for signal_iri, signal_onto_iri in signal_entities.items():
entity_groups.setdefault(signal_iri, signal_onto_iri)
symbol_predicates = tuple(
URIRef(iri) for iri in sc.induced_subgraph_symbol_predicates
)
graph = await self.sparql_tool.aget_induced_subgraph(
ontologies=catalog,
merged=merged_context,
entity_uris=entity_uris,
entity_relevance=entity_relevance,
entity_roles=entity_roles,
ontology_iris=ontology_iris,
depth=subgraph_depth,
max_total_triples=max_total_triples,
estimated_triples_per_query=estimated_triples_per_query,
ontology_version_filters=ontology_version_filters or None,
ontology_hash_filters=ontology_hash_filters or None,
hub_seed_count=hub_seed_count,
ancestor_closure_depth=ancestor_depth,
type_promotion_score_factor=(
sc.induced_subgraph_type_promotion_score_factor
),
seed_order=sc.induced_subgraph_seed_order.value,
entity_groups=entity_groups,
extra_description_predicates=symbol_predicates,
)
await self._apply_small_module_closure(graph, hit_ontology_iris)
self._last_retrieval_metrics["snapshot_triple_count"] = len(graph)
self._last_retrieval_metrics["ontology_iris_for_expansion"] = ontology_iris
self._last_retrieval_metrics.update(self.sparql_tool.last_finalize_metrics)
_bind_common_vocab_prefixes(graph)
return graph, source_iris