feat(Wave 8 W8-1 #12): activate search_relations vector recall

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,31 +196,107 @@ 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(
         self,
         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",
 ]

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 == []
 
 
 # ---------------------------------------------------------------------