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)
patch: PatchRetrievalConfig = Field(
default_factory=PatchRetrievalConfig,
exclude=True,
)
_last_retrieval_metrics: dict[str, Any] = PrivateAttr(default_factory=dict)
@property
def last_retrieval_metrics(self) -> dict[str, Any]:
return self._last_retrieval_metrics
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 retrieve(
self,
query: str,
top_k: int | None = None,
expand_sparql: bool = True,
subgraph_depth: int = 1,
max_total_triples: int = 300,
estimated_triples_per_query: int = 24,
) -> 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 = 1,
max_total_triples: int = 300,
estimated_triples_per_query: int = 24,
) -> 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,
)
)
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 = 1,
max_total_triples: int = 300,
estimated_triples_per_query: int = 24,
) -> 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,
)
async def aretrieve_ensemble(
self,
queries: list[str],
top_k: int | None = None,
expand_sparql: bool = True,
subgraph_depth: int = 1,
max_total_triples: int = 300,
estimated_triples_per_query: int = 24,
) -> tuple[RDFGraph, list[str]]:
"""Vector search over all ``queries`` once, score-filter, dedupe, single subgraph expansion."""
self._last_retrieval_metrics = {}
if not queries:
return RDFGraph(), []
eff_top_k = self._effective_top_k(top_k)
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
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,
)
atoms_after_merge = 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
]
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=pc.max_atoms,
core_weight=core_w,
neighborhood_weight=neigh_w,
)
elif pc.max_atoms > 0:
merged = merged[: pc.max_atoms]
if not merged:
self._last_retrieval_metrics = {
"query_count": len(queries),
"top_k": eff_top_k,
"atoms_after_merge": atoms_after_merge,
"atoms_final": 0,
}
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,
"merge_mode": pc.cross_query_merge_mode.value,
"atoms_after_merge": atoms_after_merge,
"atoms_final": len(merged),
"source_ontology_iris": source_iris,
"seeds_by_ontology": dict(seeds_by_ontology),
}
if not expand_sparql or self.sparql_tool is None:
return RDFGraph(), source_iris
entity_uris, entity_relevance, entity_roles = _ranked_entity_weights(merged)
hit_ontology_iris = sorted(
{atom.ontology_iri for atom in merged if atom.ontology_iri}
)
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
if self.sparql_tool.triple_store_manager is not None:
catalog = await self.sparql_tool.triple_store_manager.afetch_ontologies()
ontology_iris = _expand_ontology_iris_by_reference(
entity_uris,
hit_ontology_iris,
catalog,
)
expanded = sorted(set(ontology_iris) - set(hit_ontology_iris))
if expanded:
self._last_retrieval_metrics["expanded_ontology_iris"] = expanded
hub_seed_count = sc.induced_subgraph_hub_seed_count
ancestor_depth = sc.induced_subgraph_ancestor_closure_depth
graph = await self.sparql_tool.aget_induced_subgraph(
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,
)
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