holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

fix: improve graph retrieval on large memory banks (#141)

Browse source
This commit is contained in:
Nicolò Boschi 2026-01-09 16:43:31 +01:00 committed by GitHub
parent 4135a6cee5
commit 6232e690fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 136 additions and 112 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -29,6 +29,7 @@ nltk_data/
# Monitoring stack (Prometheus/Grafana binaries and data) # Monitoring stack (Prometheus/Grafana binaries and data)
.monitoring/ .monitoring/
.pgbouncer
# Large benchmark datasets (will be downloaded automatically) # Large benchmark datasets (will be downloaded automatically)
**/longmemeval_s_cleaned.json **/longmemeval_s_cleaned.json

18
hindsight-api/hindsight_api/engine/memory_engine.py
View file

@ -1709,6 +1709,19 @@ class MemoryEngine(MemoryEngineInterface):
f" [2] Parallel retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {parallel_duration:.3f}s{setup_info}{temporal_info}" f" [2] Parallel retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {parallel_duration:.3f}s{setup_info}{temporal_info}"
) )
# Log MPFP timing breakdown if available
if all_mpfp_timings:
mpfp_total = all_mpfp_timings[0] # Take first fact type's timing as representative
mpfp_parts = [
f"db_queries={mpfp_total.db_queries}",
f"edge_load={mpfp_total.edge_load_time:.3f}s",
f"edges={mpfp_total.edge_count}",
f"patterns={mpfp_total.pattern_count}",
]
if mpfp_total.seeds_time > 0.01:
mpfp_parts.append(f"seeds={mpfp_total.seeds_time:.3f}s")
log_buffer.append(f" [MPFP] {', '.join(mpfp_parts)}")
# Record retrieval results for tracer - per fact type # Record retrieval results for tracer - per fact type
if tracer: if tracer:
# Convert RetrievalResult to old tuple format for tracer # Convert RetrievalResult to old tuple format for tracer
@ -2337,9 +2350,10 @@ class MemoryEngine(MemoryEngineInterface):
await self._authenticate_tenant(request_context) await self._authenticate_tenant(request_context)
pool = await self._get_pool() pool = await self._get_pool()
async with acquire_with_retry(pool) as conn: async with acquire_with_retry(pool) as conn:
# Ensure connection is not in read-only mode (can happen with connection poolers)
await conn.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE")
async with conn.transaction(): async with conn.transaction():
# Ensure transaction is not in read-only mode (can happen with connection poolers)
# Using SET LOCAL so it only affects this transaction, not the session
await conn.execute("SET LOCAL transaction_read_only TO off")
try: try:
if fact_type: if fact_type:
# Delete only memories of a specific fact type # Delete only memories of a specific fact type

90
hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py
View file

@ -48,12 +48,16 @@ class EdgeCache:
Grows per-hop as edges are loaded for frontier nodes. Grows per-hop as edges are loaded for frontier nodes.
Shared across patterns to avoid redundant loads. Shared across patterns to avoid redundant loads.
Loads ALL edge types at once to minimize DB queries.
""" """
# edge_type -> from_node_id -> list of EdgeTarget # edge_type -> from_node_id -> list of EdgeTarget
graphs: dict[str, dict[str, list[EdgeTarget]]] = field(default_factory=dict) graphs: dict[str, dict[str, list[EdgeTarget]]] = field(default_factory=dict)
# Track which (edge_type, node_id) have been loaded # Track which nodes have been fully loaded (all edge types)
_loaded: set[tuple[str, str]] = field(default_factory=set) _fully_loaded: set[str] = field(default_factory=set)
# Timing stats
db_queries: int = 0
edge_load_time: float = 0.0
def get_neighbors(self, edge_type: str, node_id: str) -> list[EdgeTarget]: def get_neighbors(self, edge_type: str, node_id: str) -> list[EdgeTarget]:
"""Get neighbors for a node via a specific edge type.""" """Get neighbors for a node via a specific edge type."""
@ -71,32 +75,30 @@ class EdgeCache:
return [EdgeTarget(node_id=n.node_id, weight=n.weight / total) for n in neighbors] return [EdgeTarget(node_id=n.node_id, weight=n.weight / total) for n in neighbors]
def is_loaded(self, edge_type: str, node_id: str) -> bool: def is_fully_loaded(self, node_id: str) -> bool:
"""Check if edges for this node+type have been loaded.""" """Check if all edges for this node have been loaded."""
return (edge_type, node_id) in self._loaded return node_id in self._fully_loaded
def get_uncached(self, edge_type: str, node_ids: list[str]) -> list[str]: def get_uncached(self, node_ids: list[str]) -> list[str]:
"""Get node IDs that haven't been loaded yet for this edge type.""" """Get node IDs that haven't been fully loaded yet."""
return [n for n in node_ids if not self.is_loaded(edge_type, n)] return [n for n in node_ids if not self.is_fully_loaded(n)]
def add_edges(self, edge_type: str, edges: dict[str, list[EdgeTarget]], all_queried: list[str]): def add_all_edges(self, edges_by_type: dict[str, dict[str, list[EdgeTarget]]], all_queried: list[str]):
""" """
Add loaded edges to the cache. Add loaded edges to the cache (all edge types at once).
Args: Args:
edge_type: Type of edges edges_by_type: Dict mapping edge_type -> from_node_id -> list of EdgeTarget
edges: Dict mapping from_node_id -> list of EdgeTarget all_queried: All node IDs that were queried (marks them as fully loaded)
all_queried: All node IDs that were queried (marks them as loaded even if no edges)
""" """
for edge_type, edges in edges_by_type.items():
if edge_type not in self.graphs: if edge_type not in self.graphs:
self.graphs[edge_type] = {} self.graphs[edge_type] = {}
for node_id, neighbors in edges.items(): for node_id, neighbors in edges.items():
self.graphs[edge_type][node_id] = neighbors self.graphs[edge_type][node_id] = neighbors
# Mark all queried nodes as loaded (even if they have no edges) # Mark all queried nodes as fully loaded (even if they have no edges)
for node_id in all_queried: self._fully_loaded.update(all_queried)
self._loaded.add((edge_type, node_id))
@dataclass @dataclass
@ -148,23 +150,19 @@ class SeedNode:
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------
async def load_edges_for_frontier( async def load_all_edges_for_frontier(
pool, pool,
bank_id: str,
edge_type: str,
node_ids: list[str], node_ids: list[str],
) -> dict[str, list[EdgeTarget]]: ) -> dict[str, dict[str, list[EdgeTarget]]]:
""" """
Load edges for specific frontier nodes only. Load ALL edge types for frontier nodes in one query.
Args: Args:
pool: Database connection pool pool: Database connection pool
bank_id: Memory bank ID
edge_type: Type of edges to load
node_ids: Frontier node IDs to load edges for node_ids: Frontier node IDs to load edges for
Returns: Returns:
Dict mapping from_node_id -> list of EdgeTarget Dict mapping edge_type -> from_node_id -> list of EdgeTarget
""" """
if not node_ids: if not node_ids:
return {} return {}
@ -172,25 +170,26 @@ async def load_edges_for_frontier(
async with acquire_with_retry(pool) as conn: async with acquire_with_retry(pool) as conn:
rows = await conn.fetch( rows = await conn.fetch(
f""" f"""
SELECT ml.from_unit_id, ml.to_unit_id, ml.weight SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight
FROM {fq_table("memory_links")} ml FROM {fq_table("memory_links")} ml
WHERE ml.from_unit_id = ANY($1::uuid[]) WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = $2
AND ml.weight >= 0.1 AND ml.weight >= 0.1
ORDER BY ml.from_unit_id, ml.weight DESC ORDER BY ml.from_unit_id, ml.link_type, ml.weight DESC
""", """,
node_ids, node_ids,
edge_type,
) )
result: dict[str, list[EdgeTarget]] = defaultdict(list) # Group by edge_type -> from_node -> neighbors
result: dict[str, dict[str, list[EdgeTarget]]] = defaultdict(lambda: defaultdict(list))
for row in rows: for row in rows:
edge_type = row["link_type"]
from_id = str(row["from_unit_id"]) from_id = str(row["from_unit_id"])
to_id = str(row["to_unit_id"]) to_id = str(row["to_unit_id"])
weight = row["weight"] weight = row["weight"]
result[from_id].append(EdgeTarget(node_id=to_id, weight=weight)) result[edge_type][from_id].append(EdgeTarget(node_id=to_id, weight=weight))
return dict(result) # Convert nested defaultdicts to regular dicts
return {edge_type: dict(edges) for edge_type, edges in result.items()}
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------
@ -200,7 +199,6 @@ async def load_edges_for_frontier(
async def mpfp_traverse_async( async def mpfp_traverse_async(
pool, pool,
bank_id: str,
seeds: list[SeedNode], seeds: list[SeedNode],
pattern: list[str], pattern: list[str],
config: MPFPConfig, config: MPFPConfig,
@ -209,11 +207,10 @@ async def mpfp_traverse_async(
""" """
Async Forward Push traversal with lazy edge loading. Async Forward Push traversal with lazy edge loading.
Loads edges on-demand per hop, only for frontier nodes. Loads ALL edge types per hop to minimize DB queries.
Args: Args:
pool: Database connection pool pool: Database connection pool
bank_id: Memory bank ID
seeds: Entry point nodes with initial scores seeds: Entry point nodes with initial scores
pattern: Sequence of edge types to follow pattern: Sequence of edge types to follow
config: Algorithm parameters config: Algorithm parameters
@ -242,13 +239,18 @@ async def mpfp_traverse_async(
if not active_nodes: if not active_nodes:
break break
# Find nodes that need edge loading # Find nodes that need edge loading (all edge types at once)
uncached = cache.get_uncached(edge_type, active_nodes) uncached = cache.get_uncached(active_nodes)
# Batch load edges for uncached nodes # Batch load ALL edges for uncached nodes (one query for all edge types)
if uncached: if uncached:
edges = await load_edges_for_frontier(pool, bank_id, edge_type, uncached) import time
cache.add_edges(edge_type, edges, uncached)
load_start = time.time()
edges_by_type = await load_all_edges_for_frontier(pool, uncached)
cache.edge_load_time += time.time() - load_start
cache.db_queries += 1
cache.add_all_edges(edges_by_type, uncached)
# Propagate mass # Propagate mass
next_frontier: dict[str, float] = {} next_frontier: dict[str, float] = {}
@ -406,7 +408,9 @@ class MPFPGraphRetriever(GraphRetriever):
# If no semantic seeds provided, fall back to finding our own # If no semantic seeds provided, fall back to finding our own
if not semantic_seed_nodes: if not semantic_seed_nodes:
seeds_start = time.time()
semantic_seed_nodes = await self._find_semantic_seeds(pool, query_embedding_str, bank_id, fact_type) semantic_seed_nodes = await self._find_semantic_seeds(pool, query_embedding_str, bank_id, fact_type)
timings.seeds_time = time.time() - seeds_start
# Collect all pattern jobs # Collect all pattern jobs
pattern_jobs = [] pattern_jobs = []
@ -432,13 +436,15 @@ class MPFPGraphRetriever(GraphRetriever):
# Run all patterns in parallel (each does lazy edge loading) # Run all patterns in parallel (each does lazy edge loading)
step_start = time.time() step_start = time.time()
pattern_tasks = [ pattern_tasks = [
mpfp_traverse_async(pool, bank_id, seeds, pattern, self.config, cache) for seeds, pattern in pattern_jobs mpfp_traverse_async(pool, seeds, pattern, self.config, cache) for seeds, pattern in pattern_jobs
] ]
pattern_results = await asyncio.gather(*pattern_tasks) pattern_results = await asyncio.gather(*pattern_tasks)
timings.traverse = time.time() - step_start timings.traverse = time.time() - step_start
# Count edges loaded # Record edge loading stats from cache
timings.edge_count = sum(len(neighbors) for g in cache.graphs.values() for neighbors in g.values()) timings.edge_count = sum(len(neighbors) for g in cache.graphs.values() for neighbors in g.values())
timings.db_queries = cache.db_queries
timings.edge_load_time = cache.edge_load_time
# Fuse results # Fuse results
step_start = time.time() step_start = time.time()

17
hindsight-api/hindsight_api/engine/search/types.py
View file

@ -15,14 +15,15 @@ class MPFPTimings:
"""Timing breakdown for a single MPFP retrieval call.""" """Timing breakdown for a single MPFP retrieval call."""
fact_type: str fact_type: str
adjacency_query: float = 0.0 edge_count: int = 0 # Total edges loaded
adjacency_process: float = 0.0 db_queries: int = 0 # Number of DB queries for edge loading
edge_count: int = 0 edge_load_time: float = 0.0 # Time spent loading edges from DB
traverse: float = 0.0 traverse: float = 0.0 # Total traversal time (includes edge loading)
pattern_count: int = 0 pattern_count: int = 0 # Number of patterns executed
fusion: float = 0.0 fusion: float = 0.0 # Time for RRF fusion
fetch: float = 0.0 fetch: float = 0.0 # Time to fetch memory unit details
result_count: int = 0 seeds_time: float = 0.0 # Time to find semantic seeds (if fallback used)
result_count: int = 0 # Number of results returned
@dataclass @dataclass

96
hindsight-api/tests/test_mpfp_retrieval.py
View file

@ -20,7 +20,7 @@ from hindsight_api.engine.search.mpfp_retrieval import (
MPFPGraphRetriever, MPFPGraphRetriever,
PatternResult, PatternResult,
SeedNode, SeedNode,
load_edges_for_frontier, load_all_edges_for_frontier,
mpfp_traverse_async, mpfp_traverse_async,
rrf_fusion, rrf_fusion,
) )
@ -36,32 +36,32 @@ class TestEdgeCache:
neighbors = cache.get_neighbors("semantic", "node-1") neighbors = cache.get_neighbors("semantic", "node-1")
assert neighbors == [] assert neighbors == []
def test_is_loaded_false_for_uncached(self): def test_is_fully_loaded_false_for_uncached(self):
"""is_loaded should return False for nodes not yet loaded.""" """is_fully_loaded should return False for nodes not yet loaded."""
cache = EdgeCache() cache = EdgeCache()
assert cache.is_loaded("semantic", "node-1") is False assert cache.is_fully_loaded("node-1") is False
def test_add_edges_marks_as_loaded(self): def test_add_all_edges_marks_as_fully_loaded(self):
"""Adding edges should mark nodes as loaded.""" """Adding edges should mark nodes as fully loaded."""
cache = EdgeCache() cache = EdgeCache()
edges = { edges_by_type = {
"node-1": [EdgeTarget("node-2", 0.8), EdgeTarget("node-3", 0.6)], "semantic": {"node-1": [EdgeTarget("node-2", 0.8), EdgeTarget("node-3", 0.6)]},
} }
cache.add_edges("semantic", edges, ["node-1", "node-4"]) # node-4 has no edges cache.add_all_edges(edges_by_type, ["node-1", "node-4"]) # node-4 has no edges
assert cache.is_loaded("semantic", "node-1") is True assert cache.is_fully_loaded("node-1") is True
assert cache.is_loaded("semantic", "node-4") is True # Marked even with no edges assert cache.is_fully_loaded("node-4") is True # Marked even with no edges
assert cache.is_loaded("semantic", "node-2") is False # Target, not source assert cache.is_fully_loaded("node-2") is False # Target, not source
def test_get_neighbors_returns_added_edges(self): def test_get_neighbors_returns_added_edges(self):
"""get_neighbors should return edges after add_edges.""" """get_neighbors should return edges after add_all_edges."""
cache = EdgeCache() cache = EdgeCache()
edges = { edges_by_type = {
"node-1": [EdgeTarget("node-2", 0.8), EdgeTarget("node-3", 0.6)], "semantic": {"node-1": [EdgeTarget("node-2", 0.8), EdgeTarget("node-3", 0.6)]},
} }
cache.add_edges("semantic", edges, ["node-1"]) cache.add_all_edges(edges_by_type, ["node-1"])
neighbors = cache.get_neighbors("semantic", "node-1") neighbors = cache.get_neighbors("semantic", "node-1")
assert len(neighbors) == 2 assert len(neighbors) == 2
@ -69,28 +69,30 @@ class TestEdgeCache:
assert neighbors[0].weight == 0.8 assert neighbors[0].weight == 0.8
def test_get_uncached_filters_loaded_nodes(self): def test_get_uncached_filters_loaded_nodes(self):
"""get_uncached should only return nodes not yet loaded.""" """get_uncached should only return nodes not yet fully loaded."""
cache = EdgeCache() cache = EdgeCache()
# Load some nodes # Load some nodes (all edge types)
cache.add_edges("semantic", {"node-1": []}, ["node-1", "node-2"]) cache.add_all_edges({"semantic": {"node-1": []}}, ["node-1", "node-2"])
# Check uncached # Check uncached
uncached = cache.get_uncached("semantic", ["node-1", "node-2", "node-3", "node-4"]) uncached = cache.get_uncached(["node-1", "node-2", "node-3", "node-4"])
assert set(uncached) == {"node-3", "node-4"} assert set(uncached) == {"node-3", "node-4"}
def test_get_normalized_neighbors_normalizes_weights(self): def test_get_normalized_neighbors_normalizes_weights(self):
"""get_normalized_neighbors should normalize weights to sum to 1.""" """get_normalized_neighbors should normalize weights to sum to 1."""
cache = EdgeCache() cache = EdgeCache()
edges = { edges_by_type = {
"semantic": {
"node-1": [ "node-1": [
EdgeTarget("node-2", 0.8), EdgeTarget("node-2", 0.8),
EdgeTarget("node-3", 0.4), EdgeTarget("node-3", 0.4),
EdgeTarget("node-4", 0.2), EdgeTarget("node-4", 0.2),
], ],
},
} }
cache.add_edges("semantic", edges, ["node-1"]) cache.add_all_edges(edges_by_type, ["node-1"])
# Get top 2, normalized # Get top 2, normalized
neighbors = cache.get_normalized_neighbors("semantic", "node-1", top_k=2) neighbors = cache.get_normalized_neighbors("semantic", "node-1", top_k=2)
@ -111,8 +113,11 @@ class TestEdgeCache:
"""Different edge types should be stored separately.""" """Different edge types should be stored separately."""
cache = EdgeCache() cache = EdgeCache()
cache.add_edges("semantic", {"node-1": [EdgeTarget("node-2", 0.8)]}, ["node-1"]) edges_by_type = {
cache.add_edges("temporal", {"node-1": [EdgeTarget("node-3", 0.5)]}, ["node-1"]) "semantic": {"node-1": [EdgeTarget("node-2", 0.8)]},
"temporal": {"node-1": [EdgeTarget("node-3", 0.5)]},
}
cache.add_all_edges(edges_by_type, ["node-1"])
semantic_neighbors = cache.get_neighbors("semantic", "node-1") semantic_neighbors = cache.get_neighbors("semantic", "node-1")
temporal_neighbors = cache.get_neighbors("temporal", "node-1") temporal_neighbors = cache.get_neighbors("temporal", "node-1")
@ -198,7 +203,6 @@ class TestMPFPTraverseAsync:
result = await mpfp_traverse_async( result = await mpfp_traverse_async(
pool=None, # Not used when no seeds pool=None, # Not used when no seeds
bank_id="test",
seeds=[], seeds=[],
pattern=["semantic"], pattern=["semantic"],
config=config, config=config,
@ -213,19 +217,18 @@ class TestMPFPTraverseAsync:
cache = EdgeCache() cache = EdgeCache()
config = MPFPConfig(alpha=0.15, threshold=1e-6) config = MPFPConfig(alpha=0.15, threshold=1e-6)
# Pre-populate cache with empty edges for seed # Pre-populate cache with empty edges for seed (marks as fully loaded)
cache.add_edges("semantic", {}, ["seed-1"]) cache.add_all_edges({}, ["seed-1"])
seeds = [SeedNode("seed-1", 1.0)] seeds = [SeedNode("seed-1", 1.0)]
with patch( with patch(
"hindsight_api.engine.search.mpfp_retrieval.load_edges_for_frontier", "hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier",
new_callable=AsyncMock, new_callable=AsyncMock,
return_value={}, return_value={},
): ):
result = await mpfp_traverse_async( result = await mpfp_traverse_async(
pool=MagicMock(), pool=MagicMock(),
bank_id="test",
seeds=seeds, seeds=seeds,
pattern=["semantic"], pattern=["semantic"],
config=config, config=config,
@ -244,24 +247,25 @@ class TestMPFPTraverseAsync:
seeds = [SeedNode("seed-1", 1.0)] seeds = [SeedNode("seed-1", 1.0)]
# Mock edge loading # Mock edge loading (returns all edge types at once)
async def mock_load_edges(pool, bank_id, edge_type, node_ids): async def mock_load_all_edges(pool, node_ids):
if "seed-1" in node_ids: if "seed-1" in node_ids:
return { return {
"semantic": {
"seed-1": [ "seed-1": [
EdgeTarget("neighbor-1", 0.8), EdgeTarget("neighbor-1", 0.8),
EdgeTarget("neighbor-2", 0.4), EdgeTarget("neighbor-2", 0.4),
] ]
} }
}
return {} return {}
with patch( with patch(
"hindsight_api.engine.search.mpfp_retrieval.load_edges_for_frontier", "hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier",
side_effect=mock_load_edges, side_effect=mock_load_all_edges,
): ):
result = await mpfp_traverse_async( result = await mpfp_traverse_async(
pool=MagicMock(), pool=MagicMock(),
bank_id="test",
seeds=seeds, seeds=seeds,
pattern=["semantic"], pattern=["semantic"],
config=config, config=config,
@ -287,22 +291,21 @@ class TestMPFPTraverseAsync:
seeds = [SeedNode("seed-1", 1.0)] seeds = [SeedNode("seed-1", 1.0)]
# Mock edge loading for two hops # Mock edge loading for two hops (returns all edge types at once)
async def mock_load_edges(pool, bank_id, edge_type, node_ids): async def mock_load_all_edges(pool, node_ids):
edges = {} edges: dict[str, dict[str, list[EdgeTarget]]] = {"semantic": {}}
if "seed-1" in node_ids: if "seed-1" in node_ids:
edges["seed-1"] = [EdgeTarget("hop1-node", 1.0)] edges["semantic"]["seed-1"] = [EdgeTarget("hop1-node", 1.0)]
if "hop1-node" in node_ids: if "hop1-node" in node_ids:
edges["hop1-node"] = [EdgeTarget("hop2-node", 1.0)] edges["semantic"]["hop1-node"] = [EdgeTarget("hop2-node", 1.0)]
return edges return edges
with patch( with patch(
"hindsight_api.engine.search.mpfp_retrieval.load_edges_for_frontier", "hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier",
side_effect=mock_load_edges, side_effect=mock_load_all_edges,
): ):
result = await mpfp_traverse_async( result = await mpfp_traverse_async(
pool=MagicMock(), pool=MagicMock(),
bank_id="test",
seeds=seeds, seeds=seeds,
pattern=["semantic", "semantic"], # Two hops pattern=["semantic", "semantic"], # Two hops
config=config, config=config,
@ -320,27 +323,26 @@ class TestMPFPTraverseAsync:
cache = EdgeCache() cache = EdgeCache()
config = MPFPConfig(alpha=0.15, threshold=1e-6) config = MPFPConfig(alpha=0.15, threshold=1e-6)
# Pre-load cache # Pre-load cache (marks seed-1 as fully loaded)
cache.add_edges("semantic", {"seed-1": [EdgeTarget("neighbor-1", 1.0)]}, ["seed-1"]) cache.add_all_edges({"semantic": {"seed-1": [EdgeTarget("neighbor-1", 1.0)]}}, ["seed-1"])
seeds = [SeedNode("seed-1", 1.0)] seeds = [SeedNode("seed-1", 1.0)]
load_mock = AsyncMock(return_value={}) load_mock = AsyncMock(return_value={})
with patch( with patch(
"hindsight_api.engine.search.mpfp_retrieval.load_edges_for_frontier", "hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier",
load_mock, load_mock,
): ):
await mpfp_traverse_async( await mpfp_traverse_async(
pool=MagicMock(), pool=MagicMock(),
bank_id="test",
seeds=seeds, seeds=seeds,
pattern=["semantic"], pattern=["semantic"],
config=config, config=config,
cache=cache, cache=cache,
) )
# Should not call load_edges_for_frontier since seed-1 is already cached # Should not call load_all_edges_for_frontier since seed-1 is already cached
load_mock.assert_not_called() load_mock.assert_not_called()