diff --git a/aperag/indexing/graph_search_service.py b/aperag/indexing/graph_search_service.py index e153202f4..c34ea60a2 100644 --- a/aperag/indexing/graph_search_service.py +++ b/aperag/indexing/graph_search_service.py @@ -57,12 +57,21 @@ candidates are more expensive than false-positive recall.""" -# Vector point indexer tag — pinned by §K.12 invariant #5. Task #3 +# Vector point indexer tags — pinned by §K.12 invariant #5. Task #3 # writes graph-entity vectors with ``payload["indexer"] == -# "graph_entity"``; this service filters on the same tag so chunk / -# summary vectors sharing the physical collection never bleed into -# graph recall. +# "graph_entity"`` and graph-relation vectors with +# ``payload["indexer"] == "graph_relation"``; this service filters on +# the matching tag so chunk / summary vectors sharing the physical +# collection never bleed into graph recall. GRAPH_ENTITY_INDEXER: str = "graph_entity" +GRAPH_RELATION_INDEXER: str = "graph_relation" + + +# ``entity_name`` payload field for ``graph_relation`` points carries the +# composite ``f"{source}->{target}"`` per the task #3 writer +# (``aperag/indexing/graph.py:1631``). The arrow has no surrounding +# whitespace; we split exactly once to recover the endpoints. +RELATION_PAYLOAD_ARROW: str = "->" class GraphSearchService: @@ -73,8 +82,11 @@ class GraphSearchService: * :meth:`search_entities` — embed the query, ANN-search the entity vector index, fetch the matching ``EntityWithLineage`` rows. - * :meth:`search_relations` — derive relations attached to the - vector-recalled entities (1-hop expansion). + * :meth:`search_relations` — embed the query, ANN-search the + relation vector index, fetch the matching ``RelationWithLineage`` + rows. Wave 8 W8-1 (task #12) replaced the Wave 7 conservative + 1-hop expansion with direct vector recall now that task #3 is + writing relation vectors. * :meth:`get_subgraph` — wrap ``expand_neighbors_n_hops`` as a stable read primitive for MCP tools (task #7). * :meth:`compose_context` — render the ``EntityWithLineage`` / @@ -184,7 +196,7 @@ async def _batch_get_entities(self, names: Sequence[str]) -> list[EntityWithLine return [e for e in results if e is not None] # ------------------------------------------------------------------ - # search_relations — derived from search_entities + # search_relations — vector recall path (Wave 8 W8-1 task #12) # ------------------------------------------------------------------ async def search_relations( @@ -192,23 +204,99 @@ async def search_relations( query: str, top_k: int | None = None, ) -> list[RelationWithLineage]: - """Recall relations attached to the entities that vector-search - retrieved for ``query``. - - Vector store does not carry separate relation vectors (Wave 7 - scope: only entity vectors are written by task #3; per-relation - vectors are out-of-scope per architect msg=acbd0003 §K.12.5). - We therefore derive relations as the 1-hop expansion of the - entity-search result. + """Vector-recall the top-K relations matching ``query``. + + Wave 8 W8-1 (architect ratify task #12): replaces the Wave 7 + conservative 1-hop expansion with a direct ANN search filtered + on ``Eq("indexer", "graph_relation")`` — task #3 (PR #1757) has + been writing relation vectors all along; this is finally the + consumer side. Mirrors :meth:`search_entities` shape for + symmetry. + + Hits resolve to full :class:`RelationWithLineage` rows via + per-hit ``store.get_relation`` reverse lookup so + :meth:`compose_context` keeps its byte-parity with the legacy + rendering (per task #5 invariant). Hits whose payload doesn't + parse cleanly (missing arrow / empty endpoints / no + ``entity_type``) are dropped silently — better to skip a hit + than to surface a relation we can't reconstruct. """ - entities = await self.search_entities(query, top_k=top_k) - if not entities: + text = (query or "").strip() + if not text: + return [] + k = top_k if top_k is not None else self._top_k + if k <= 0: + return [] + + try: + embedding = await asyncio.to_thread(self._embedder.embed_query, text) + except Exception: + logger.warning( + "graph_search_service: embed failed for relation query=%r", + text, + exc_info=True, + ) return [] - _neighbour_entities, relations = await self._store.expand_neighbors_n_hops( - entity_names=[e.name for e in entities], - hops=1, + + request = QueryRequest( + embedding=embedding, + top_k=k, + flt=Eq("indexer", GRAPH_RELATION_INDEXER), + score_threshold=self._score_threshold or None, + ) + try: + hits = await asyncio.to_thread(self._vector_connector.search, request) + except Exception: + logger.warning( + "graph_search_service: vector search failed for relation query=%r", + text, + exc_info=True, + ) + return [] + + keys = self._relation_keys_from_hits(hits) + if not keys: + return [] + return await self._batch_get_relations(keys) + + @staticmethod + def _relation_keys_from_hits( + hits: Iterable[SearchHit], + ) -> list[tuple[str, str, str]]: + # Preserve hit order (vector store ranked by score) and de-dup + # so the same edge returned twice (e.g. alias-redirect side-effect) + # only fetches once. + seen: set[tuple[str, str, str]] = set() + ordered: list[tuple[str, str, str]] = [] + for hit in hits: + payload = hit.payload or {} + composite = payload.get("entity_name") or payload.get("name") + relation_type = payload.get("entity_type") or payload.get("relation_type") + if not composite or not relation_type: + continue + composite_str = str(composite) + if RELATION_PAYLOAD_ARROW not in composite_str: + continue + source, _, target = composite_str.partition(RELATION_PAYLOAD_ARROW) + if not source or not target: + continue + key = (source, target, str(relation_type)) + if key in seen: + continue + seen.add(key) + ordered.append(key) + return ordered + + async def _batch_get_relations(self, keys: Sequence[tuple[str, str, str]]) -> list[RelationWithLineage]: + # Same ``asyncio.gather`` shape as :meth:`_batch_get_entities`: + # N is bounded by ``top_k`` (≤ 10 in practice) so per-key + # round-trip cost is acceptable, and the Protocol surface stays + # narrow (no batch-get method). + results = await asyncio.gather( + *(self._store.get_relation(s, t, ty) for s, t, ty in keys), + return_exceptions=False, ) - return relations + return [r for r in results if r is not None] # ------------------------------------------------------------------ # get_subgraph — MCP-facing read primitive (task #7) @@ -338,5 +426,7 @@ def build_graph_search_service_for(collection: Any) -> GraphSearchService: "DEFAULT_TOP_K", "DEFAULT_SCORE_THRESHOLD", "GRAPH_ENTITY_INDEXER", + "GRAPH_RELATION_INDEXER", + "RELATION_PAYLOAD_ARROW", "build_graph_search_service_for", ] diff --git a/tests/unit_test/indexing/test_graph_search_service.py b/tests/unit_test/indexing/test_graph_search_service.py index eccf087c3..5224b1963 100644 --- a/tests/unit_test/indexing/test_graph_search_service.py +++ b/tests/unit_test/indexing/test_graph_search_service.py @@ -13,18 +13,23 @@ # limitations under the License. """Unit tests for ``aperag.indexing.graph_search_service`` — Wave 7 -task #5. +task #5 + Wave 8 W8-1 task #12. -Pins the Wave 7 vector-recall contract: +Pins the vector-recall contract: * ``search_entities`` embeds the query, ANN-searches the ``graph_entity`` indexer slice (filter + threshold pinned), then fetches matching ``EntityWithLineage`` rows via per-name ``get_entity`` (asyncio.gather). De-dups payload names so an aliased entity returned twice doesn't double-fetch. -* ``search_relations`` derives relations as the 1-hop expansion of the - vector-recalled entities — vector store carries no per-relation - vectors in Wave 7. +* ``search_relations`` (Wave 8 W8-1 upgrade) embeds the query, + ANN-searches the ``graph_relation`` indexer slice, parses each hit's + ``entity_name="src->tgt"`` + ``entity_type=relation_type`` payload + per the task #3 writer, then reverse-looks-up + ``RelationWithLineage`` via ``store.get_relation``. Hits whose + payload doesn't parse cleanly are skipped silently. Replaces the + Wave 7 conservative 1-hop expansion path now that task #3 is + writing relation vectors. * ``get_subgraph`` is a thin pass-through to ``expand_neighbors_n_hops`` for MCP / retrieval callers. * ``compose_context`` renders byte-for-byte the same LightRAG-style @@ -124,15 +129,22 @@ def __init__( self, entities: dict[str, EntityWithLineage] | None = None, expansions: dict[tuple[str, ...], tuple[list[EntityWithLineage], list[RelationWithLineage]]] | None = None, + relations: dict[tuple[str, str, str], RelationWithLineage] | None = None, ) -> None: self._entities = entities or {} self._expansions = expansions or {} + self._relations = relations or {} self.get_entity_calls: list[str] = [] + self.get_relation_calls: list[tuple[str, str, str]] = [] async def get_entity(self, entity_name: str) -> EntityWithLineage | None: self.get_entity_calls.append(entity_name) return self._entities.get(entity_name) + async def get_relation(self, source: str, target: str, type: str) -> RelationWithLineage | None: + self.get_relation_calls.append((source, target, type)) + return self._relations.get((source, target, type)) + async def expand_neighbors_n_hops( self, *, @@ -185,13 +197,14 @@ def _make_service( *, entities: dict[str, EntityWithLineage] | None = None, expansions: dict[tuple[str, ...], tuple[list[EntityWithLineage], list[RelationWithLineage]]] | None = None, + relations: dict[tuple[str, str, str], RelationWithLineage] | None = None, hits: list[SearchHit] | None = None, embedder: Any | None = None, connector: Any | None = None, top_k: int = 10, score_threshold: float = 0.0, ) -> tuple[GraphSearchService, _FakeStore, _FakeVectorConnector | Any, _FakeEmbedder | Any]: - store = _FakeStore(entities=entities, expansions=expansions) + store = _FakeStore(entities=entities, expansions=expansions, relations=relations) connector = connector if connector is not None else _FakeVectorConnector(hits=hits) embedder = embedder if embedder is not None else _FakeEmbedder() service = GraphSearchService( @@ -322,31 +335,147 @@ async def test_search_entities_swallows_vector_store_failure(): # --------------------------------------------------------------------- -# search_relations +# search_relations (Wave 8 W8-1 vector recall path) # --------------------------------------------------------------------- @pytest.mark.asyncio -async def test_search_relations_empty_when_no_entities_match(): - service, _, _, _ = _make_service(entities={}, hits=[]) - assert await service.search_relations("query") == [] +async def test_search_relations_empty_query_returns_empty(): + service, _, connector, embedder = _make_service(entities={}) + assert await service.search_relations("") == [] + assert await service.search_relations(" ") == [] + assert connector.searches == [] + assert embedder.calls == [] @pytest.mark.asyncio -async def test_search_relations_returns_one_hop_expansion_of_entity_results(): - a = _entity("Alpha") - b = _entity("Beta") - rel = _relation("Alpha", "Beta") - service, _, _, _ = _make_service( - entities={"Alpha": a, "Beta": b}, +async def test_search_relations_zero_topk_returns_empty(): + service, _, connector, embedder = _make_service(entities={}) + assert await service.search_relations("query", top_k=0) == [] + assert connector.searches == [] + assert embedder.calls == [] + + +@pytest.mark.asyncio +async def test_search_relations_uses_graph_relation_filter(): + """Wave 8 W8-1: filter pinned to ``Eq("indexer", "graph_relation")`` + so the ANN never bleeds into entity / chunk vectors sharing the + physical collection.""" + service, _, connector, _ = _make_service(entities={}, hits=[], score_threshold=0.42, top_k=5) + await service.search_relations("query") + assert len(connector.searches) == 1 + request = connector.searches[0].request + from aperag.indexing.graph_search_service import GRAPH_RELATION_INDEXER + + assert request.flt == Eq("indexer", GRAPH_RELATION_INDEXER) + assert request.score_threshold == 0.42 + assert request.top_k == 5 + + +@pytest.mark.asyncio +async def test_search_relations_parses_payload_and_resolves_via_get_relation(): + """Hit payload carries ``entity_name="src->tgt"`` + + ``entity_type=relation_type`` (per task #3 writer + ``aperag/indexing/graph.py:1631``); we split, reverse-lookup via + ``store.get_relation`` and return the full ``RelationWithLineage`` + so :meth:`compose_context` keeps its byte-parity rendering.""" + rel = _relation("Alpha", "Beta", relation_type="founded") + service, store, _, _ = _make_service( + relations={("Alpha", "Beta", "founded"): rel}, hits=[ - SearchHit(id="1", score=0.9, payload={"entity_name": "Alpha"}), - SearchHit(id="2", score=0.8, payload={"entity_name": "Beta"}), + SearchHit( + id="1", + score=0.9, + payload={"entity_name": "Alpha->Beta", "entity_type": "founded"}, + ), ], - expansions={("Alpha", "Beta"): ([a, b], [rel])}, ) relations = await service.search_relations("query") assert relations == [rel] + assert store.get_relation_calls == [("Alpha", "Beta", "founded")] + + +@pytest.mark.asyncio +async def test_search_relations_preserves_hit_order_and_dedupes(): + rel_ab = _relation("Alpha", "Beta", relation_type="founded") + rel_bc = _relation("Beta", "Gamma", relation_type="acquired") + service, store, _, _ = _make_service( + relations={ + ("Alpha", "Beta", "founded"): rel_ab, + ("Beta", "Gamma", "acquired"): rel_bc, + }, + hits=[ + SearchHit(id="1", score=0.9, payload={"entity_name": "Beta->Gamma", "entity_type": "acquired"}), + SearchHit(id="2", score=0.8, payload={"entity_name": "Alpha->Beta", "entity_type": "founded"}), + # Duplicate of the first hit (e.g. alias-redirect side-effect) + # — must dedupe so we don't double-fetch the same edge. + SearchHit(id="3", score=0.7, payload={"entity_name": "Beta->Gamma", "entity_type": "acquired"}), + ], + ) + relations = await service.search_relations("query") + assert [(r.source, r.target) for r in relations] == [ + ("Beta", "Gamma"), + ("Alpha", "Beta"), + ] + assert store.get_relation_calls == [("Beta", "Gamma", "acquired"), ("Alpha", "Beta", "founded")] + + +@pytest.mark.asyncio +async def test_search_relations_skips_payload_missing_arrow_or_type(): + """A hit whose payload doesn't parse cleanly is dropped silently — + better to skip than to surface an edge we can't reconstruct.""" + rel_ab = _relation("Alpha", "Beta", relation_type="founded") + service, store, _, _ = _make_service( + relations={("Alpha", "Beta", "founded"): rel_ab}, + hits=[ + SearchHit(id="ghost1", score=1.0, payload={}), # no payload at all + SearchHit( + id="ghost2", score=0.99, payload={"entity_name": "AlphaBeta", "entity_type": "founded"} + ), # missing arrow + SearchHit(id="ghost3", score=0.98, payload={"entity_name": "Alpha->Beta"}), # missing entity_type + SearchHit( + id="ghost4", score=0.97, payload={"entity_name": "->Beta", "entity_type": "founded"} + ), # empty source + SearchHit( + id="ghost5", score=0.96, payload={"entity_name": "Alpha->", "entity_type": "founded"} + ), # empty target + SearchHit(id="real", score=0.9, payload={"entity_name": "Alpha->Beta", "entity_type": "founded"}), + ], + ) + relations = await service.search_relations("query") + assert [(r.source, r.target) for r in relations] == [("Alpha", "Beta")] + # Only the real hit hit the store. + assert store.get_relation_calls == [("Alpha", "Beta", "founded")] + + +@pytest.mark.asyncio +async def test_search_relations_drops_edge_gced_from_store(): + """Vector hit for an edge that was deleted between sync and search + → store.get_relation returns None → drop from result, no exception.""" + service, _, _, _ = _make_service( + relations={}, # store empty: every get_relation returns None + hits=[ + SearchHit(id="1", score=0.9, payload={"entity_name": "Alpha->Beta", "entity_type": "founded"}), + ], + ) + assert await service.search_relations("query") == [] + + +@pytest.mark.asyncio +async def test_search_relations_swallows_embedder_failure(): + service, store, connector, _ = _make_service(entities={}, embedder=_FailingEmbedder()) + assert await service.search_relations("query") == [] + assert connector.searches == [] + assert store.get_relation_calls == [] + + +@pytest.mark.asyncio +async def test_search_relations_swallows_vector_store_failure(): + service, store, connector, embedder = _make_service(entities={}, connector=_FailingVectorConnector()) + assert await service.search_relations("query") == [] + assert embedder.calls == ["query"] + assert len(connector.searches) == 1 + assert store.get_relation_calls == [] # ---------------------------------------------------------------------