feat(enh-005): directional GraphRAG expansion (in/out/any)

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

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
 

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)
 

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)

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,
         )

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": [],

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],