Bound graph retrieval by query relevance

Author: Coding-Dev-Tools

engraphis/core/recall.py

@@ -364,23 +364,39 @@ def _code_arm(
         if not symbols:
             return {}
 
+        aliases: dict[str, str] = {}
+        for symbol_id, symbol in symbols.items():
+            for key in ("id", "name", "fqname"):
+                value = str(symbol.get(key) or "")
+                if value:
+                    aliases[value] = symbol_id
         # Expand one stored code edge to capture callers/callees, bounded by a
-        # multiple of candidate_k.  Both symbol ids and parser-emitted names are
-        # accepted because language backends use both representations.
+        # multiple of candidate_k. Query only edges incident to matched aliases
+        # before applying that cap, so later files cannot be hidden by a global prefix.
+        edge_kwargs = {
+            "limit": max(100, min(2000, candidate_k * 20)),
+            "layers": flt.graph_layers,
+        }
+        # ``endpoints`` is a v2 Store optimization. Preserve compatibility with
+        # external code stores that have not added the optional filter yet.
+        try:
+            edge_parameters = inspect.signature(self.store.list_code_edges).parameters.values()
+            supports_endpoints = any(
+                parameter.name == "endpoints"
+                or parameter.kind == inspect.Parameter.VAR_KEYWORD
+                for parameter in edge_parameters
+            )
+        except (TypeError, ValueError):
+            supports_endpoints = False
+        if supports_endpoints:
+            edge_kwargs["endpoints"] = list(aliases)
         code_edges = _call_temporal_store(
             self.store.list_code_edges,
             flt,
             flt.repo_id,
-            limit=max(100, min(2000, candidate_k * 20)),
-            layers=flt.graph_layers,
             requested_historical=historical,
+            **edge_kwargs,
         )
-        aliases: dict[str, str] = {}
-        for symbol_id, symbol in symbols.items():
-            for key in ("id", "name", "fqname"):
-                value = str(symbol.get(key) or "")
-                if value:
-                    aliases[value] = symbol_id
         related_names: dict[str, float] = {}
         for edge in code_edges:
             src, dst = str(edge.get("src") or ""), str(edge.get("dst") or "")
@@ -493,8 +509,31 @@ def connect(a: str, b: str, w: float) -> None:
             adj.setdefault(a, []).append((b, w))
             adj.setdefault(b, []).append((a, w))
 
-        edges = self.store.edges_in_scope(flt, at=now, limit=4000)
-        for e in edges:
+        # Build a bounded edge set outward from the query entities.  A global
+        # ULID-ordered cap would let old unrelated edges crowd out a new relation
+        # required by this query before PPR sees it.
+        edge_cap = 4000
+        edges_by_id = {}
+        frontier = set(seeds)
+        expanded: set[str] = set()
+        while frontier and len(edges_by_id) < edge_cap:
+            batch = sorted(frontier - expanded)[:400]
+            if not batch:
+                break
+            frontier.difference_update(batch)
+            expanded.update(batch)
+            next_frontier: set[str] = set()
+            for edge in self.store.neighbors(
+                    batch, at=now, layers=flt.graph_layers, flt=flt,
+                    limit=edge_cap - len(edges_by_id)):
+                if edge.id in edges_by_id:
+                    continue
+                edges_by_id[edge.id] = edge
+                next_frontier.update((edge.src, edge.dst))
+                if len(edges_by_id) >= edge_cap:
+                    break
+            frontier.update(next_frontier - expanded)
+        for e in edges_by_id.values():
             connect(ent(e.src), ent(e.dst), max(float(e.weight or 1.0), 1e-6))
 
         incidence = self.store.list_memory_entities(flt, limit=12_000)
@@ -505,9 +544,23 @@ def connect(a: str, b: str, w: float) -> None:
         # retrieval arms so PPR can traverse that edge without widening scope. Keep
         # the incidence frontier as well when independent caps choose a different
         # subset of the scoped memory universe.
-        memory_ids = sorted({
+        incidence_memory_ids = {
             str(row.get("memory_id") or "")
             for row in incidence if row.get("memory_id")
+        }
+        frontier_links = self.store.links_touching(
+            sorted(incidence_memory_ids),
+            layers=flt.graph_layers,
+            flt=flt,
+            limit=20_000,
+        )
+        # Expand from the entity-incidence frontier before adding the bounded newest
+        # memory window. An older unmentioned endpoint can then participate in PPR
+        # through its visible link instead of being silently dropped by that window.
+        memory_ids = sorted(incidence_memory_ids | {
+            endpoint
+            for link in frontier_links
+            for endpoint in (link["a"], link["b"])
         } | {
             memory.id for memory in self.store.list_memories(flt, limit=12_000)
         })

engraphis/core/store.py

@@ -3168,9 +3168,68 @@ def links_among(self, ids: list[str], *,
                     break
         return rows
 
+    def links_touching(self, ids: list[str], *,
+                       layers: Optional[list[GraphLayer]] = None,
+                       flt: Optional[SearchFilter] = None,
+                       include_invalid: bool = False,
+                       limit: Optional[int] = None) -> list[dict]:
+        """Return visible links with at least one endpoint in ``ids``.
+
+        This bounded frontier expansion is distinct from :meth:`links_among`: graph
+        recall uses it to retain an unmentioned endpoint linked to an entity-attached
+        memory, without first materializing every memory in a large scope.
+        """
+        if not ids:
+            return []
+        if layers is not None and not layers:
+            return []
+        row_cap = None if limit is None else max(0, int(limit))
+        if row_cap == 0:
+            return []
+        ordered_ids = sorted(set(ids))
+        visibility_sql, visibility_params = _temporal_visibility_sql("", flt)
+        rows: list[dict] = []
+        seen: set[tuple] = set()
+        # Each id appears once for each endpoint predicate; reserve parameters for
+        # time/layer filters so this remains under SQLite's portable bind limit.
+        chunk_size = max(1, (IN_CLAUSE_CHUNK - 16) // 2)
+        for start in range(0, len(ordered_ids), chunk_size):
+            if row_cap is not None and len(rows) >= row_cap:
+                break
+            chunk = ordered_ids[start:start + chunk_size]
+            marks = ",".join("?" for _ in chunk)
+            sql = (
+                "SELECT a, b, relation, layer, reason, created_at, valid_from, valid_to, "
+                "valid_to_recorded_at, ingested_at, expired_at FROM mem_links "
+                f"WHERE (a IN ({marks}) OR b IN ({marks}))"
+            )
+            params: list[Any] = [*chunk, *chunk]
+            if not include_invalid:
+                sql += f" AND {visibility_sql}"
+                params.extend(visibility_params)
+            if layers is not None:
+                layer_marks = ",".join("?" for _ in layers)
+                sql += f" AND layer IN ({layer_marks})"
+                params.extend(_enum(layer) for layer in layers)
+            sql += " ORDER BY a, b, relation, valid_from, ingested_at"
+            for row in self.conn.execute(sql, params).fetchall():
+                item = dict(row)
+                key = (
+                    item["a"], item["b"], item["relation"], item["layer"],
+                    item["valid_from"], item["valid_to"], item["ingested_at"],
+                )
+                if key in seen:
+                    continue
+                seen.add(key)
+                rows.append(item)
+                if row_cap is not None and len(rows) >= row_cap:
+                    break
+        return rows
+
     def neighbors(self, node_ids: list[str], *, at: Optional[float] = None,
                   layers: Optional[list[GraphLayer]] = None,
-                  flt: Optional[SearchFilter] = None) -> list[Edge]:
+                  flt: Optional[SearchFilter] = None,
+                  limit: Optional[int] = None) -> list[Edge]:
         if not node_ids:
             return []
         valid_at, known_at = _temporal_anchors(flt, valid_at=at)
@@ -3213,6 +3272,10 @@ def neighbors(self, node_ids: list[str], *, at: Optional[float] = None,
             else:
                 sql += " AND repo_id=?"
             params.append(flt.repo_id)
+        sql += " ORDER BY id"
+        if limit is not None:
+            sql += " LIMIT ?"
+            params.append(max(0, int(limit)))
         rows = self.conn.execute(sql, params).fetchall()
         return [_row_to_edge(r) for r in rows]
 
@@ -3409,6 +3472,7 @@ def list_symbols_page(self, repo_id: str, *,
 
     def list_code_edges(self, repo_id: str, *, limit: Optional[int] = None,
                         layers: Optional[list[GraphLayer]] = None,
+                        endpoints: Optional[list[str]] = None,
                         flt: Optional[SearchFilter] = None) -> list[dict]:
         temporal, params = _temporal_visibility_sql("", flt)
         sql = "SELECT * FROM code_edges WHERE repo_id=? AND " + temporal
@@ -3419,6 +3483,13 @@ def list_code_edges(self, repo_id: str, *, limit: Optional[int] = None,
             marks = ",".join("?" for _ in layers)
             sql += f" AND layer IN ({marks})"
             params.extend(_enum(layer) for layer in layers)
+        if endpoints is not None:
+            if not endpoints:
+                return []
+            marks = ",".join("?" for _ in endpoints)
+            sql += f" AND (src IN ({marks}) OR dst IN ({marks}))"
+            params.extend(endpoints)
+            params.extend(endpoints)
         sql += " ORDER BY file, line, id"
         if limit is not None:
             sql += " LIMIT ?"

tests/test_core_store.py

@@ -491,6 +491,22 @@ def test_code_listing_helpers_honor_limit(store):
     assert store.list_code_files("repo_x", languages={"rust"}, limit=3) == []
 
 
+def test_code_edge_endpoint_filter_applies_before_limit(store):
+    for index in range(4):
+        store.add_code_edge(
+            repo_id="repo_x", src=f"noise_{index}", dst=f"other_{index}",
+            relation="calls", file="a_noise.py", line=index,
+        )
+    store.add_code_edge(
+        repo_id="repo_x", src="target", dst="caller", relation="calls",
+        file="z_target.py", line=1,
+    )
+
+    edges = store.list_code_edges("repo_x", endpoints=["target"], limit=1)
+
+    assert [(edge["src"], edge["dst"]) for edge in edges] == [("target", "caller")]
+
+
 def test_memory_links_infer_and_filter_graph_layers(store):
     wid = store.get_or_create_workspace("w")
     a = store.add_memory(MemoryRecord(id="", content="cause", workspace_id=wid))

tests/test_recall.py

@@ -95,6 +95,37 @@ def test_graph_arm_pulls_related_via_entities():
     assert any("checkout" in c["content"].lower() for c in res.chunks)
 
 
+def test_graph_arm_selects_seed_frontier_edges_before_global_edge_cap(monkeypatch):
+    from engraphis.core.interfaces import Edge, Node
+
+    store, emb, eng = _engine()
+    wid = store.get_or_create_workspace("w")
+    rid = store.get_or_create_repo(wid, "r")
+    redis = store.upsert_entity(Node(
+        id="", name="Redis", ntype="tech", workspace_id=wid, repo_id=rid))
+    checkout = store.upsert_entity(Node(
+        id="", name="checkout", ntype="module", workspace_id=wid, repo_id=rid))
+    store.upsert_edge(Edge(
+        id="", src=redis, dst=checkout, relation="used_by",
+        workspace_id=wid, repo_id=rid))
+    memory_id = _add(store, emb, wid, rid, "Checkout uses the Redis cache.")
+    store.link_memory_entity(
+        memory_id=memory_id, entity_id=checkout, workspace_id=wid, repo_id=rid,
+        source_kind="test", confidence=1.0,
+    )
+    def no_global_edges(*_args, **_kwargs):
+        raise AssertionError("PPR must traverse from query seeds, not global edges")
+
+    monkeypatch.setattr(store, "edges_in_scope", no_global_edges)
+
+    scores = eng._graph_arm_ppr(
+        "How does Redis relate to the checkout service?",
+        SearchFilter(workspace_id=wid, repo_id=rid), now=10**12,
+    )
+
+    assert memory_id in scores
+
+
 def test_graph_arm_backfills_text_memory_when_its_entity_is_added_later():
     from engraphis.core.interfaces import Edge, Node
 
@@ -169,6 +200,41 @@ def test_graph_arm_backfills_workspace_mentions_for_a_later_repo_entity():
     )
 
 
+def test_graph_arm_expands_an_older_unmentioned_link_endpoint_from_incidence(monkeypatch):
+    from engraphis.core.interfaces import Node
+
+    store, emb, eng = _engine()
+    wid = store.get_or_create_workspace("w")
+    rid = store.get_or_create_repo(wid, "r")
+    redis = store.upsert_entity(Node(
+        id="", name="Redis", ntype="technology", workspace_id=wid, repo_id=rid,
+    ))
+    attached = _add(store, emb, wid, rid, "The cache migration is attached evidence.")
+    older_unmentioned = _add(store, emb, wid, rid, "The old rollout required a staged cutover.")
+    store.link_memory_entity(
+        memory_id=attached, entity_id=redis, workspace_id=wid, repo_id=rid,
+        source_kind="test", confidence=1.0,
+    )
+    store.add_link(attached, older_unmentioned, relation="supports")
+
+    # Simulate a full scope whose bounded newest-memory window excludes the older
+    # endpoint. The incidence frontier still contains ``attached``.
+    monkeypatch.setattr(
+        store,
+        "list_memories",
+        lambda *_args, **_kwargs: [
+            MemoryRecord(id=f"mem_new_{i}", content="") for i in range(12_000)
+        ],
+    )
+
+    scores = eng._graph_arm_ppr(
+        "How does Redis relate to the rollout?",
+        SearchFilter(workspace_id=wid, repo_id=rid), now=10**12,
+    )
+
+    assert older_unmentioned in scores
+
+
 def test_entity_backfill_preserves_closed_workspace_memory_history():
     from engraphis.core.interfaces import Node