diff --git a/CHANGELOG.md b/CHANGELOG.md index e5b066a..3b4a3b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,16 @@ slice between two versions to see which of your filed issues a bump fixes. ### Added +- **enh-005 — directional GraphRAG expansion.** `retrieval.graph_expansion` + (and the `GraphExpansion` value) gain `direction: in | out | any`, so graph + expansion can follow **asymmetric** edges the right way — `in` for callers / + who-cites-X / dependents, `out` for callees / what-X-cites / dependencies. + `any` (the default) is the original undirected `traverse` behaviour, so + existing agents are **byte-for-byte unchanged**. Directional expansion rides + the already-locked `GraphStore.get_edges(direction=...)` primitive via a BFS + in the `Retriever` — **no ABC change, no major bump, no new package**. + Improves feat-023. First piece of the 0.4 graph train (see the enh-005 spec). + - **enh-006 — upgrade drift report (`agentforge upgrade --notes`).** After bumping the framework pin, a consumer can now see which fixes (and the issues they close) the bump shipped, so dead workarounds get cleaned up diff --git a/docs/enhancements/enh-005-directional-graph-expansion.md b/docs/enhancements/enh-005-directional-graph-expansion.md index dae1ac3..4ba146e 100644 --- a/docs/enhancements/enh-005-directional-graph-expansion.md +++ b/docs/enhancements/enh-005-directional-graph-expansion.md @@ -15,7 +15,7 @@ |---|---| | **ID** | enh-005 | | **Title** | `direction` on `GraphExpansion` (in / out / any) | -| **Status** | accepted (targeted at the 0.4 train) | +| **Status** | in-progress (implemented; targeting the 0.4 train) | | **Owner** | kjoshi | | **Created** | 2026-06-17 | | **Target version** | 0.4 | @@ -201,11 +201,21 @@ scheduled. - `GraphStore` contract — `get_edges(direction=...)` (the locked primitive used). ## 11. Implementation status (Python) -**Status: accepted, not yet implemented.** Suggested chunking: -1. Spec + catalogue/roadmap pointer. -2. `direction` on `GraphExpansion` + `GraphExpansionConfig`; `_expand_via_graph` - directional branch via `get_edges`; unit tests incl. backward-compat. -3. Status flip + CHANGELOG + runbook note on feat-023. +**Status: implemented (not yet released).** Shipped on the enh-005 branch: +- `direction: Literal["out","in","any"] = "any"` on `GraphExpansion` + (`agentforge_core.values.retrieval`) and `GraphExpansionConfig` + (`agentforge_core.config.schema`); threaded by `build_retriever_from_config`. +- `Retriever._expand_via_graph` now branches: `any` keeps the native + `traverse()` path (unchanged); `in`/`out` run a BFS over the locked + `get_edges(direction=...)` primitive (`_reach_via_get_edges`), collecting + `edge.dst` (`out`) or `edge.src` (`in`) and fetching each neighbour via + `get_node`. Merge/decay/dedup pipeline reused verbatim. +- Tests in `test_retriever_graphrag.py`: out=successors, in=predecessors, + in surfaces what `traverse` cannot, multi-hop caller chain, edge-type + filtering, no-predecessor tolerance, value + config defaults / rejection, + and the existing feat-023 suite (backward-compat for `any`). Fully offline. + +Remaining: flip Status → shipped + roadmap row when the 0.4 train tags. ## 12. Runbook diff --git a/packages/agentforge-core/src/agentforge_core/config/schema.py b/packages/agentforge-core/src/agentforge_core/config/schema.py index 0a91fec..e3bbc29 100644 --- a/packages/agentforge-core/src/agentforge_core/config/schema.py +++ b/packages/agentforge-core/src/agentforge_core/config/schema.py @@ -164,6 +164,10 @@ class GraphExpansionConfig(BaseModel): """Edge-type filter. YAML lists deserialize to `list[str]`; converted to a tuple by `build_retriever_from_config` before constructing the `GraphExpansion` value.""" + direction: Literal["out", "in", "any"] = "any" + """Edge direction to follow during expansion (enh-005): `out` + (callees / what-X-cites), `in` (callers / who-cites-X), or `any` + (default — the original undirected `traverse` behaviour).""" text_property: str = "text" decay: float = Field(default=0.5, gt=0.0, le=1.0) diff --git a/packages/agentforge-core/src/agentforge_core/values/retrieval.py b/packages/agentforge-core/src/agentforge_core/values/retrieval.py index d5de820..195ab52 100644 --- a/packages/agentforge-core/src/agentforge_core/values/retrieval.py +++ b/packages/agentforge-core/src/agentforge_core/values/retrieval.py @@ -13,6 +13,8 @@ from __future__ import annotations +from typing import Literal + from pydantic import BaseModel, ConfigDict, Field from agentforge_core.contracts.graph_store import GraphStore @@ -32,6 +34,11 @@ class GraphExpansion(BaseModel): fan-out — tune cautiously. edge_types: If set, restricts traversal to these edge types. ``None`` means all edge types. + direction: Which way to follow edges during expansion + (enh-005). ``"out"`` follows ``src → dst`` (e.g. + callees, what-X-cites); ``"in"`` follows ``dst → src`` + (e.g. callers, who-cites-X); ``"any"`` (default) is the + original undirected behaviour via ``store.traverse``. text_property: Graph-node property used to populate the synthesised ``VectorMatch.text``. Defaults to ``"text"``. @@ -49,5 +56,6 @@ class GraphExpansion(BaseModel): store: GraphStore max_hops: int = Field(default=2, ge=1) edge_types: tuple[str, ...] | None = None + direction: Literal["out", "in", "any"] = "any" text_property: str = Field(default="text", min_length=1) decay: float = Field(default=0.5, gt=0.0, le=1.0) diff --git a/packages/agentforge/src/agentforge/cli/_build.py b/packages/agentforge/src/agentforge/cli/_build.py index 16039a0..f934f76 100644 --- a/packages/agentforge/src/agentforge/cli/_build.py +++ b/packages/agentforge/src/agentforge/cli/_build.py @@ -251,6 +251,7 @@ def build_retriever_from_config(config: AgentForgeConfig) -> Retriever | None: store=graph_store_instance, max_hops=ge_cfg.max_hops, edge_types=tuple(ge_cfg.edge_types) if ge_cfg.edge_types is not None else None, + direction=ge_cfg.direction, text_property=ge_cfg.text_property, decay=ge_cfg.decay, ) diff --git a/packages/agentforge/src/agentforge/cli/release_notes.json b/packages/agentforge/src/agentforge/cli/release_notes.json index 255f1ed..0f5de8e 100644 --- a/packages/agentforge/src/agentforge/cli/release_notes.json +++ b/packages/agentforge/src/agentforge/cli/release_notes.json @@ -6,6 +6,14 @@ "date": null, "entries": { "Added": [ + { + "label": "enh-005 — directional GraphRAG expansion.", + "closes": [], + "refs": [ + "enh-005", + "feat-023" + ] + }, { "label": "enh-006 — upgrade drift report (`agentforge upgrade --notes`).", "closes": [], diff --git a/packages/agentforge/src/agentforge/retrieval.py b/packages/agentforge/src/agentforge/retrieval.py index 4603858..e92a022 100644 --- a/packages/agentforge/src/agentforge/retrieval.py +++ b/packages/agentforge/src/agentforge/retrieval.py @@ -29,7 +29,7 @@ from agentforge_core.contracts.embedding import EmbeddingClient from agentforge_core.contracts.reranker import Reranker from agentforge_core.contracts.vector_store import VectorStore -from agentforge_core.values.graph import Path as GraphPath +from agentforge_core.values.graph import GraphNode from agentforge_core.values.retrieval import GraphExpansion from agentforge_core.values.vector import VectorItem, VectorMatch from ulid import ULID @@ -313,53 +313,44 @@ async def _expand_via_graph( if not seeds: return [] - async def _traverse(seed: VectorMatch) -> tuple[VectorMatch, list[GraphPath]]: + async def _reach(seed: VectorMatch) -> tuple[VectorMatch, list[tuple[int, GraphNode]]]: try: - paths = await expansion.store.traverse( - start_id=seed.id, - edge_types=expansion.edge_types, - max_depth=expansion.max_hops, - limit=max(len(seeds), 1) * expansion.max_hops * 4, - ) + if expansion.direction == "any": + reaches = await self._reach_via_traverse(seed, expansion, n_seeds=len(seeds)) + else: + reaches = await self._reach_via_get_edges(seed, expansion) except Exception: - log.debug("graph traverse failed for seed %s", seed.id, exc_info=True) + log.debug("graph expansion failed for seed %s", seed.id, exc_info=True) return seed, [] - return seed, paths + return seed, reaches - results = await asyncio.gather(*(_traverse(s) for s in seeds)) + results = await asyncio.gather(*(_reach(s) for s in seeds)) seed_ids = {s.id for s in seeds} - # Keyed by node id; track best (highest) decayed score per id. + # Keyed by node id; track best (highest) decayed score per id — + # since score = decay**depth (decay <= 1), the lowest-depth reach + # wins, so this also picks the shortest path to each node. expanded_by_id: dict[str, tuple[float, int, VectorMatch]] = {} - for seed, paths in results: - if not paths: - log.debug("no graph paths found for seed %s", seed.id) + for seed, reaches in results: + if not reaches: + log.debug("no graph expansion found for seed %s", seed.id) continue - for path in paths: - # path.nodes[0] is the seed; nodes[i] is at depth i. - for depth, node in enumerate(path.nodes): - if depth == 0: - continue - if node.id in seed_ids: - continue - score = float(seed.score) * (float(expansion.decay) ** depth) - prior = expanded_by_id.get(node.id) - if prior is not None and prior[0] >= score: - continue - text = str(node.properties.get(expansion.text_property, "")) - merged_meta: dict[str, Any] = dict(node.properties) - merged_meta["agentforge.expanded_from"] = seed.id - merged_meta["agentforge.hop"] = depth - expanded_by_id[node.id] = ( - score, - depth, - VectorMatch( - id=node.id, - text=text, - metadata=merged_meta, - score=score, - ), - ) + for depth, node in reaches: + if node.id in seed_ids: + continue + score = float(seed.score) * (float(expansion.decay) ** depth) + prior = expanded_by_id.get(node.id) + if prior is not None and prior[0] >= score: + continue + text = str(node.properties.get(expansion.text_property, "")) + merged_meta: dict[str, Any] = dict(node.properties) + merged_meta["agentforge.expanded_from"] = seed.id + merged_meta["agentforge.hop"] = depth + expanded_by_id[node.id] = ( + score, + depth, + VectorMatch(id=node.id, text=text, metadata=merged_meta, score=score), + ) expansion_matches = sorted( (m for _, _, m in expanded_by_id.values()), @@ -368,6 +359,62 @@ async def _traverse(seed: VectorMatch) -> tuple[VectorMatch, list[GraphPath]]: ) return list(seeds) + expansion_matches + @staticmethod + async def _reach_via_traverse( + seed: VectorMatch, + expansion: GraphExpansion, + *, + n_seeds: int, + ) -> list[tuple[int, GraphNode]]: + """Undirected expansion (``direction="any"``) via the store's + native ``traverse`` — the original feat-023 path, unchanged.""" + paths = await expansion.store.traverse( + start_id=seed.id, + edge_types=expansion.edge_types, + max_depth=expansion.max_hops, + limit=max(n_seeds, 1) * expansion.max_hops * 4, + ) + # path.nodes[0] is the seed; nodes[i] is at depth i. Emit every + # (depth, node); the caller's best-score dedup keeps the shortest. + return [ + (depth, node) for path in paths for depth, node in enumerate(path.nodes) if depth > 0 + ] + + @staticmethod + async def _reach_via_get_edges( + seed: VectorMatch, + expansion: GraphExpansion, + ) -> list[tuple[int, GraphNode]]: + """Directional expansion (``direction in {"out", "in"}``) via a + BFS over the locked ``get_edges(direction=...)`` primitive — no + ABC change (enh-005). ``out`` collects ``edge.dst``; ``in`` + collects ``edge.src``.""" + direction = expansion.direction + edge_types = expansion.edge_types or (None,) + visited: set[str] = {seed.id} + frontier: set[str] = {seed.id} + reaches: list[tuple[int, GraphNode]] = [] + for depth in range(1, expansion.max_hops + 1): + neighbour_ids: set[str] = set() + for node_id in frontier: + for edge_type in edge_types: + edges = await expansion.store.get_edges( + node_id, edge_type=edge_type, direction=direction + ) + for edge in edges: + nbr = edge.dst if direction == "out" else edge.src + if nbr not in visited: + neighbour_ids.add(nbr) + if not neighbour_ids: + break + for nbr in neighbour_ids: + visited.add(nbr) + node = await expansion.store.get_node(nbr) + if node is not None: + reaches.append((depth, node)) + frontier = neighbour_ids + return reaches + def _rrf_fuse( self, vec: list[VectorMatch], diff --git a/packages/agentforge/tests/unit/test_retriever_graphrag.py b/packages/agentforge/tests/unit/test_retriever_graphrag.py index 96730de..87dc004 100644 --- a/packages/agentforge/tests/unit/test_retriever_graphrag.py +++ b/packages/agentforge/tests/unit/test_retriever_graphrag.py @@ -12,6 +12,7 @@ import pytest from agentforge import InMemoryGraphStore, InMemoryVectorStore, Retriever +from agentforge_core.config.schema import GraphExpansionConfig, ModuleEntry from agentforge_core.contracts.embedding import EmbeddingClient from agentforge_core.contracts.reranker import Reranker from agentforge_core.values.graph import GraphEdge, GraphNode @@ -245,6 +246,160 @@ async def test_missing_graph_node_is_silently_skipped() -> None: assert [m.id for m in results] == ["orphan"] +# ---------- Directional expansion (enh-005) ---------- + + +async def _directional_stores() -> tuple[InMemoryVectorStore, InMemoryGraphStore]: + """`a` cites `b`; `c` cites `a`. So out(a) = {b} (what a cites); + in(a) = {c} (who cites a). The reference `traverse` follows + out-edges only, so only `direction="in"` surfaces `c`.""" + return await _seeded_stores( + docs={"a": "alpha", "b": "beta", "c": "gamma"}, + edges=[("a", "b", "CITES"), ("c", "a", "CITES")], + ) + + +def test_graph_expansion_direction_defaults_to_any() -> None: + gs = InMemoryGraphStore() + assert GraphExpansion(store=gs).direction == "any" + assert GraphExpansion(store=gs, direction="in").direction == "in" + + +def test_graph_expansion_config_direction_field() -> None: + entry = ModuleEntry(driver="kuzu", config={"path": ".ckg"}) + assert GraphExpansionConfig(store=entry).direction == "any" # default + assert GraphExpansionConfig(store=entry, direction="in").direction == "in" + with pytest.raises(ValueError, match="direction"): + GraphExpansionConfig(store=entry, direction="backwards") # type: ignore[arg-type] + + +def test_graph_expansion_rejects_bad_direction() -> None: + with pytest.raises(ValueError, match="direction"): + GraphExpansion(store=InMemoryGraphStore(), direction="sideways") # type: ignore[arg-type] + + +@pytest.mark.asyncio +async def test_direction_out_returns_successors() -> None: + vec_store, graph_store = await _directional_stores() + retriever = Retriever( + store=vec_store, + embedder=_FakeEmbedder(dim=4), + top_k=1, + over_fetch_factor=1, + graph_expansion=GraphExpansion(store=graph_store, max_hops=1, direction="out", decay=0.5), + ) + ids = [m.id for m in await retriever.retrieve("alpha")] + assert ids == ["a", "b"] # what `a` cites; `c` (a caller) is excluded + + +@pytest.mark.asyncio +async def test_direction_in_returns_predecessors() -> None: + vec_store, graph_store = await _directional_stores() + retriever = Retriever( + store=vec_store, + embedder=_FakeEmbedder(dim=4), + top_k=1, + over_fetch_factor=1, + graph_expansion=GraphExpansion(store=graph_store, max_hops=1, direction="in", decay=0.5), + ) + ids = [m.id for m in await retriever.retrieve("alpha")] + assert ids == ["a", "c"] # who cites `a`; `b` (a callee) is excluded + + +@pytest.mark.asyncio +async def test_direction_in_surfaces_what_traverse_cannot() -> None: + """`direction="in"` reaches predecessors the undirected `traverse` + path (`any`) misses, and excludes the out-neighbours it would add.""" + vec_store, graph_store = await _directional_stores() + base = { + "store": vec_store, + "embedder": _FakeEmbedder(dim=4), + "top_k": 1, + "over_fetch_factor": 1, + } + any_ids = { + m.id + for m in await Retriever( + **base, + graph_expansion=GraphExpansion(store=graph_store, max_hops=1, direction="any"), + ).retrieve("alpha") + } + in_ids = { + m.id + for m in await Retriever( + **base, + graph_expansion=GraphExpansion(store=graph_store, max_hops=1, direction="in"), + ).retrieve("alpha") + } + assert "c" in in_ids # `direction=in` reaches the caller `c` + assert "c" not in any_ids # the undirected traverse path misses it + assert "b" in any_ids # `any` reaches the callee `b` + assert "b" not in in_ids # `direction=in` excludes it + + +@pytest.mark.asyncio +async def test_direction_in_multi_hop_walks_caller_chain() -> None: + """`b` cites `c`, `c` cites `a`. From `a`, direction=in, 2 hops → + who-cites-a (`c`) then who-cites-c (`b`).""" + vec_store, graph_store = await _seeded_stores( + docs={"a": "alpha", "b": "beta", "c": "gamma"}, + edges=[("c", "a", "CITES"), ("b", "c", "CITES")], + ) + retriever = Retriever( + store=vec_store, + embedder=_FakeEmbedder(dim=4), + top_k=1, + over_fetch_factor=1, + graph_expansion=GraphExpansion(store=graph_store, max_hops=2, direction="in", decay=0.5), + ) + results = await retriever.retrieve("alpha") + by_id = {m.id: m for m in results} + assert {"a", "c", "b"} <= set(by_id) + assert by_id["c"].metadata["agentforge.hop"] == 1 + assert by_id["b"].metadata["agentforge.hop"] == 2 + assert by_id["c"].score == pytest.approx(by_id["a"].score * 0.5) + assert by_id["b"].score == pytest.approx(by_id["a"].score * 0.25) + assert by_id["c"].metadata["agentforge.expanded_from"] == "a" + + +@pytest.mark.asyncio +async def test_direction_in_respects_edge_types() -> None: + """`b` CALLS `a`, `c` IMPORTS `a`. direction=in + edge_types=(CALLS,) + surfaces only the caller `b`, not the importer `c`.""" + vec_store, graph_store = await _seeded_stores( + docs={"a": "alpha", "b": "beta", "c": "gamma"}, + edges=[("b", "a", "CALLS"), ("c", "a", "IMPORTS")], + ) + retriever = Retriever( + store=vec_store, + embedder=_FakeEmbedder(dim=4), + top_k=1, + over_fetch_factor=1, + graph_expansion=GraphExpansion( + store=graph_store, max_hops=1, direction="in", edge_types=("CALLS",), decay=0.5 + ), + ) + ids = {m.id for m in await retriever.retrieve("alpha")} + assert "b" in ids + assert "c" not in ids + + +@pytest.mark.asyncio +async def test_direction_in_tolerates_seed_with_no_predecessors() -> None: + vec_store, graph_store = await _directional_stores() + # In `a`→`b`, `c`→`a`, node `c` has no in-edges (nobody cites c) → + # no expansion, no error. + retriever = Retriever( + store=vec_store, + embedder=_FakeEmbedder(dim=4), + top_k=1, + over_fetch_factor=1, + graph_expansion=GraphExpansion(store=graph_store, max_hops=2, direction="in", decay=0.5), + ) + results = await retriever.retrieve("gamma") + assert [m.id for m in results] == ["c"] + + # ---------- Reranker layering ----------