feat(celery Wave 6 #33 chunk 1): LightRAG-style query layer Protocol stubs on LineageGraphStore (#1737)

Per §K.11.11 sub-spec — first chunk of Wave 6 PR-A "graphindex
elimination" cross-cutting refactor. Adds the read API the retrieval
pipeline + graph-curation flows will consume in chunk 3 (caller
migration), so chunk 2 (per-backend implementations) can be reviewed
against a stable Protocol surface without spec drift.

`aperag/indexing/graph.py`:
* New `LineageGraphStore` Protocol methods (declarations + docstrings
  only, ~85 LOC):
  - `query_entities_by_keyword(query, top_k)` — lexical recall
  - `query_entities_by_vector(embedding, top_k)` — vector recall (caller
    pre-embeds; backend stays embedder-agnostic per
    simple-stable-directive separation of concerns)
  - `expand_neighbors_n_hops(entity_names, hops=1)` — graph traversal
    bounded by `hops` (no per-collection traversal-budget knob;
    operator/dev pick `hops=1` for plain 1-hop or pass higher
    explicitly).
* Returns are the canonical `EntityWithLineage` /
  `RelationWithLineage` shapes already used by `get_*` — uniform read
  surface. Callers compose context text themselves rather than
  receiving a backend-formatted string (cleaner than legacy
  `query_context() -> ctx.text`).
* `InMemoryLineageGraphStore` ships `NotImplementedError("Wave 6 #33
  chunk 2 — pending")` stubs so chunk 2 can switch to real behaviour
  without rewriting test scaffolding, and so any caller that
  accidentally invokes the API before chunk 2 lands gets a clear
  diagnostic.

`tests/unit_test/indexing/test_lineage_query_protocol.py`: 5 new
unit tests pinning the Protocol surface (3 method names + signature
shape via `inspect.signature` + the `NotImplementedError` chunk-1
contract on the in-memory store).

Out of scope (deferred to chunk 2 / chunk 3 per §K.11.11):
* Postgres / Neo4j / Nebula backend implementations (chunk 2).
* Caller migration: retrieval/pipeline.py + knowledge_graph/service.py
  + graph_curation/* (chunk 3).
* Legacy `aperag/domains/knowledge_graph/graphindex/` package delete
  + legacy tests delete + grep-zero verify (chunk 3).

Production-readiness 三类 (chunk 1 specific):
- must-be-real: Protocol surface declared with full docstrings; tests
  pin signatures and chunk-1 stub contract.
- may-be-gated: in-memory + production backends raise
  `NotImplementedError` until chunk 2 — gate fires loudly so a caller
  cannot silently fall through to wrong behaviour.
- fully-resolves: §K.11.11 chunk 1 acceptance scope (Protocol stubs
  ~200 LOC).

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

Author: earayu

aperag/indexing/graph.py

@@ -507,6 +507,89 @@ async def get_entity(self, entity_name: str) -> EntityWithLineage | None:
     async def get_relation(self, source: str, target: str, type: str) -> RelationWithLineage | None:
         """Read-path helper for relations."""
 
+    # ------------------------------------------------------------------
+    # Wave 6 #33 chunk 1 — LightRAG-style retrieval query layer (Protocol
+    # stubs; cross-backend implementations land in chunk 2 per
+    # §K.11.11). The retrieval pipeline (§G.5) and graph-curation flows
+    # consume these to compose a graph-recall context for a user query
+    # without going through the legacy ``GraphIndexService.query_context``.
+    #
+    # Design notes:
+    # * Returns are the canonical :class:`EntityWithLineage` /
+    #   :class:`RelationWithLineage` shapes already used by ``get_*`` —
+    #   keeps the read surface uniform; callers compose context text
+    #   themselves rather than relying on a backend-formatted string.
+    # * ``query_by_vector`` takes a pre-computed embedding so backends
+    #   stay agnostic of the embedding model; the retrieval pipeline
+    #   already owns the embedder.
+    # * ``expand_neighbors_n_hops`` returns BOTH neighbour entities and
+    #   the relation edges that connect them (1-hop or multi-hop). Per
+    #   simple-stable directive the multi-hop frontier is bounded by
+    #   ``hops`` (caller picks; default 1) so backends don't need a
+    #   query planner.
+    # ------------------------------------------------------------------
+
+    async def query_entities_by_keyword(
+        self,
+        *,
+        query: str,
+        top_k: int,
+    ) -> list[EntityWithLineage]:
+        """Return up to ``top_k`` entities whose ``name`` (or backend
+        equivalent text-searchable field) matches ``query``.
+
+        Match semantics are backend-defined (case-insensitive substring
+        on Postgres / Cypher CONTAINS on Neo4j / etc.); the contract
+        is "best-effort lexical recall, not strict equality". An empty
+        ``query`` returns ``[]`` (no spurious recall).
+
+        Wave 6 #33 chunk 1 — Protocol stub; backend implementations
+        land chunk 2.
+        """
+
+    async def query_entities_by_vector(
+        self,
+        *,
+        embedding: list[float],
+        top_k: int,
+    ) -> list[EntityWithLineage]:
+        """Return up to ``top_k`` entities ranked by similarity of
+        ``embedding`` against an entity-vector index the backend
+        maintains. Backends without a native vector index may compose
+        this against the collection's existing Qdrant collection
+        (entity-name-keyed points) — chunk 2 picks the per-backend
+        approach.
+
+        Empty ``embedding`` is invalid and raises ``ValueError`` (the
+        retrieval pipeline always passes a real vector — the embedder
+        has already validated the query upstream).
+
+        Wave 6 #33 chunk 1 — Protocol stub.
+        """
+
+    async def expand_neighbors_n_hops(
+        self,
+        *,
+        entity_names: list[str],
+        hops: int = 1,
+    ) -> tuple[list[EntityWithLineage], list[RelationWithLineage]]:
+        """Return entities + relations reachable from any of
+        ``entity_names`` within ``hops`` graph traversal steps.
+
+        ``hops=1`` returns immediate neighbours plus the connecting
+        relations. ``hops=N`` returns the N-hop frontier; the result
+        UNION includes the seed entities themselves so callers can
+        feed it directly to a context formatter without a separate
+        join step.
+
+        Implementations MUST bound the result by ``hops`` to keep the
+        traversal cost predictable; per simple-stable directive
+        operators don't manage tunable per-collection traversal
+        budgets — this is the only knob.
+
+        Wave 6 #33 chunk 1 — Protocol stub.
+        """
+
 
 # ---------------------------------------------------------------------
 # In-memory reference implementation — usable by tests and as the
@@ -705,6 +788,40 @@ async def get_relation(self, source: str, target: str, type: str) -> RelationWit
                 description_parts=tuple(_sorted_description_parts(row.description_parts.values())),
             )
 
+    # ------------------------------------------------------------------
+    # Wave 6 #33 chunk 1 — LightRAG-style query layer stubs.
+    # Real implementations land chunk 2 (per §K.11.11 3-chunk
+    # decomposition). The in-memory store also defers to chunk 2 so
+    # tests that assert the Protocol surface can pin the
+    # NotImplementedError contract today and switch to the real
+    # behaviour when chunk 2 ships — without rewriting the test
+    # scaffolding.
+    # ------------------------------------------------------------------
+
+    async def query_entities_by_keyword(
+        self,
+        *,
+        query: str,
+        top_k: int,
+    ) -> list[EntityWithLineage]:
+        raise NotImplementedError("Wave 6 #33 chunk 2 — query_entities_by_keyword pending implementation")
+
+    async def query_entities_by_vector(
+        self,
+        *,
+        embedding: list[float],
+        top_k: int,
+    ) -> list[EntityWithLineage]:
+        raise NotImplementedError("Wave 6 #33 chunk 2 — query_entities_by_vector pending implementation")
+
+    async def expand_neighbors_n_hops(
+        self,
+        *,
+        entity_names: list[str],
+        hops: int = 1,
+    ) -> tuple[list[EntityWithLineage], list[RelationWithLineage]]:
+        raise NotImplementedError("Wave 6 #33 chunk 2 — expand_neighbors_n_hops pending implementation")
+
 
 def _sorted_lineage(members: Iterable[LineageMember]) -> list[LineageMember]:
     return sorted(members, key=lambda m: (m.document_id, m.parse_version))

tests/unit_test/indexing/test_lineage_query_protocol.py

@@ -0,0 +1,105 @@
+# Copyright 2025 ApeCloud, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Wave 6 #33 chunk 1 — Protocol surface for the LightRAG-style query
+layer on :class:`aperag.indexing.graph.LineageGraphStore`.
+
+Pin the Protocol method signatures so chunk 2 (cross-backend impl) can
+be reviewed against a stable surface, and so callers (chunk 3 caller
+migration) can rely on the methods existing even before the per-backend
+implementations land.
+
+Per §K.11.11 sub-spec the read API is:
+* ``query_entities_by_keyword(query, top_k)`` — lexical recall
+* ``query_entities_by_vector(embedding, top_k)`` — vector recall
+* ``expand_neighbors_n_hops(entity_names, hops)`` — graph traversal
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+
+import pytest
+
+from aperag.indexing.graph import (
+    InMemoryLineageGraphStore,
+    LineageGraphStore,
+)
+
+
+def test_lineage_query_protocol_exposes_three_read_methods():
+    """The new query layer surface must declare exactly the three
+    LightRAG-style read methods agreed in §K.11.11. Adding a method
+    here implies a chunk 2 backend implementation contract."""
+
+    members = set(LineageGraphStore.__dict__)
+    assert "query_entities_by_keyword" in members
+    assert "query_entities_by_vector" in members
+    assert "expand_neighbors_n_hops" in members
+
+
+def test_query_by_keyword_signature_takes_query_and_top_k():
+    sig = inspect.signature(LineageGraphStore.query_entities_by_keyword)
+    params = sig.parameters
+    assert "query" in params
+    assert "top_k" in params
+    # kw-only by spec — the retrieval pipeline always names the args.
+    assert params["query"].kind is inspect.Parameter.KEYWORD_ONLY
+    assert params["top_k"].kind is inspect.Parameter.KEYWORD_ONLY
+
+
+def test_query_by_vector_signature_takes_embedding_and_top_k():
+    sig = inspect.signature(LineageGraphStore.query_entities_by_vector)
+    params = sig.parameters
+    assert "embedding" in params
+    assert "top_k" in params
+    assert params["embedding"].kind is inspect.Parameter.KEYWORD_ONLY
+    assert params["top_k"].kind is inspect.Parameter.KEYWORD_ONLY
+
+
+def test_expand_neighbors_signature_takes_entity_names_and_hops():
+    sig = inspect.signature(LineageGraphStore.expand_neighbors_n_hops)
+    params = sig.parameters
+    assert "entity_names" in params
+    assert "hops" in params
+    assert params["entity_names"].kind is inspect.Parameter.KEYWORD_ONLY
+    assert params["hops"].kind is inspect.Parameter.KEYWORD_ONLY
+    # Default hops=1 is the simple-stable-directive choice — caller
+    # must opt into multi-hop explicitly.
+    assert params["hops"].default == 1
+
+
+def test_in_memory_store_query_methods_raise_not_implemented_in_chunk_1():
+    """Chunk 1 ships Protocol stubs only — the in-memory test store
+    must surface a clear NotImplementedError so callers (chunk 3
+    migration) can rely on the gate firing if they accidentally
+    invoke the API before chunk 2 lands."""
+
+    store = InMemoryLineageGraphStore()
+
+    async def _run() -> None:
+        with pytest.raises(NotImplementedError) as exc:
+            await store.query_entities_by_keyword(query="anything", top_k=5)
+        assert "chunk 2" in str(exc.value)
+
+        with pytest.raises(NotImplementedError) as exc:
+            await store.query_entities_by_vector(embedding=[0.1, 0.2, 0.3], top_k=5)
+        assert "chunk 2" in str(exc.value)
+
+        with pytest.raises(NotImplementedError) as exc:
+            await store.expand_neighbors_n_hops(entity_names=["X"], hops=1)
+        assert "chunk 2" in str(exc.value)
+
+    asyncio.run(_run())