From 6232e690fc9bc946af4b27f746e62cb5cccc2384 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Fri, 9 Jan 2026 16:43:31 +0100 Subject: [PATCH] fix: improve graph retrieval on large memory banks (#141) --- .gitignore | 1 + .../hindsight_api/engine/memory_engine.py | 18 ++- .../engine/search/mpfp_retrieval.py | 98 ++++++++------- .../hindsight_api/engine/search/types.py | 17 +-- hindsight-api/tests/test_mpfp_retrieval.py | 114 +++++++++--------- 5 files changed, 136 insertions(+), 112 deletions(-) diff --git a/.gitignore b/.gitignore index adc3105a..71aaea52 100644 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,7 @@ nltk_data/ # Monitoring stack (Prometheus/Grafana binaries and data) .monitoring/ +.pgbouncer # Large benchmark datasets (will be downloaded automatically) **/longmemeval_s_cleaned.json diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 94de2c16..27f11079 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -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}" ) + # 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 if tracer: # Convert RetrievalResult to old tuple format for tracer @@ -2337,9 +2350,10 @@ class MemoryEngine(MemoryEngineInterface): await self._authenticate_tenant(request_context) pool = await self._get_pool() 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(): + # 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: if fact_type: # Delete only memories of a specific fact type diff --git a/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py b/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py index cef90f2b..2f20810c 100644 --- a/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py +++ b/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py @@ -48,12 +48,16 @@ class EdgeCache: Grows per-hop as edges are loaded for frontier nodes. 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 graphs: dict[str, dict[str, list[EdgeTarget]]] = field(default_factory=dict) - # Track which (edge_type, node_id) have been loaded - _loaded: set[tuple[str, str]] = field(default_factory=set) + # Track which nodes have been fully loaded (all edge types) + _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]: """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] - def is_loaded(self, edge_type: str, node_id: str) -> bool: - """Check if edges for this node+type have been loaded.""" - return (edge_type, node_id) in self._loaded + def is_fully_loaded(self, node_id: str) -> bool: + """Check if all edges for this node have been loaded.""" + return node_id in self._fully_loaded - def get_uncached(self, edge_type: str, node_ids: list[str]) -> list[str]: - """Get node IDs that haven't been loaded yet for this edge type.""" - return [n for n in node_ids if not self.is_loaded(edge_type, n)] + def get_uncached(self, node_ids: list[str]) -> list[str]: + """Get node IDs that haven't been fully loaded yet.""" + 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: - edge_type: Type of edges - edges: Dict mapping from_node_id -> list of EdgeTarget - all_queried: All node IDs that were queried (marks them as loaded even if no edges) + edges_by_type: Dict mapping edge_type -> from_node_id -> list of EdgeTarget + all_queried: All node IDs that were queried (marks them as fully loaded) """ - if edge_type not in self.graphs: - self.graphs[edge_type] = {} + for edge_type, edges in edges_by_type.items(): + if edge_type not in self.graphs: + self.graphs[edge_type] = {} + for node_id, neighbors in edges.items(): + self.graphs[edge_type][node_id] = neighbors - for node_id, neighbors in edges.items(): - self.graphs[edge_type][node_id] = neighbors - - # Mark all queried nodes as loaded (even if they have no edges) - for node_id in all_queried: - self._loaded.add((edge_type, node_id)) + # Mark all queried nodes as fully loaded (even if they have no edges) + self._fully_loaded.update(all_queried) @dataclass @@ -148,23 +150,19 @@ class SeedNode: # ----------------------------------------------------------------------------- -async def load_edges_for_frontier( +async def load_all_edges_for_frontier( pool, - bank_id: str, - edge_type: 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: 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 Returns: - Dict mapping from_node_id -> list of EdgeTarget + Dict mapping edge_type -> from_node_id -> list of EdgeTarget """ if not node_ids: return {} @@ -172,25 +170,26 @@ async def load_edges_for_frontier( async with acquire_with_retry(pool) as conn: rows = await conn.fetch( 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 WHERE ml.from_unit_id = ANY($1::uuid[]) - AND ml.link_type = $2 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, - 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: + edge_type = row["link_type"] from_id = str(row["from_unit_id"]) to_id = str(row["to_unit_id"]) 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( pool, - bank_id: str, seeds: list[SeedNode], pattern: list[str], config: MPFPConfig, @@ -209,11 +207,10 @@ async def mpfp_traverse_async( """ 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: pool: Database connection pool - bank_id: Memory bank ID seeds: Entry point nodes with initial scores pattern: Sequence of edge types to follow config: Algorithm parameters @@ -242,13 +239,18 @@ async def mpfp_traverse_async( if not active_nodes: break - # Find nodes that need edge loading - uncached = cache.get_uncached(edge_type, active_nodes) + # Find nodes that need edge loading (all edge types at once) + 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: - edges = await load_edges_for_frontier(pool, bank_id, edge_type, uncached) - cache.add_edges(edge_type, edges, uncached) + import time + + 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 next_frontier: dict[str, float] = {} @@ -406,7 +408,9 @@ class MPFPGraphRetriever(GraphRetriever): # If no semantic seeds provided, fall back to finding our own 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) + timings.seeds_time = time.time() - seeds_start # Collect all pattern jobs pattern_jobs = [] @@ -432,13 +436,15 @@ class MPFPGraphRetriever(GraphRetriever): # Run all patterns in parallel (each does lazy edge loading) step_start = time.time() 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) 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.db_queries = cache.db_queries + timings.edge_load_time = cache.edge_load_time # Fuse results step_start = time.time() diff --git a/hindsight-api/hindsight_api/engine/search/types.py b/hindsight-api/hindsight_api/engine/search/types.py index 29f0dec1..eb7f48db 100644 --- a/hindsight-api/hindsight_api/engine/search/types.py +++ b/hindsight-api/hindsight_api/engine/search/types.py @@ -15,14 +15,15 @@ class MPFPTimings: """Timing breakdown for a single MPFP retrieval call.""" fact_type: str - adjacency_query: float = 0.0 - adjacency_process: float = 0.0 - edge_count: int = 0 - traverse: float = 0.0 - pattern_count: int = 0 - fusion: float = 0.0 - fetch: float = 0.0 - result_count: int = 0 + edge_count: int = 0 # Total edges loaded + db_queries: int = 0 # Number of DB queries for edge loading + edge_load_time: float = 0.0 # Time spent loading edges from DB + traverse: float = 0.0 # Total traversal time (includes edge loading) + pattern_count: int = 0 # Number of patterns executed + fusion: float = 0.0 # Time for RRF fusion + fetch: float = 0.0 # Time to fetch memory unit details + seeds_time: float = 0.0 # Time to find semantic seeds (if fallback used) + result_count: int = 0 # Number of results returned @dataclass diff --git a/hindsight-api/tests/test_mpfp_retrieval.py b/hindsight-api/tests/test_mpfp_retrieval.py index 18464628..5325303a 100644 --- a/hindsight-api/tests/test_mpfp_retrieval.py +++ b/hindsight-api/tests/test_mpfp_retrieval.py @@ -20,7 +20,7 @@ from hindsight_api.engine.search.mpfp_retrieval import ( MPFPGraphRetriever, PatternResult, SeedNode, - load_edges_for_frontier, + load_all_edges_for_frontier, mpfp_traverse_async, rrf_fusion, ) @@ -36,32 +36,32 @@ class TestEdgeCache: neighbors = cache.get_neighbors("semantic", "node-1") assert neighbors == [] - def test_is_loaded_false_for_uncached(self): - """is_loaded should return False for nodes not yet loaded.""" + def test_is_fully_loaded_false_for_uncached(self): + """is_fully_loaded should return False for nodes not yet loaded.""" 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): - """Adding edges should mark nodes as loaded.""" + def test_add_all_edges_marks_as_fully_loaded(self): + """Adding edges should mark nodes as fully loaded.""" cache = EdgeCache() - edges = { - "node-1": [EdgeTarget("node-2", 0.8), EdgeTarget("node-3", 0.6)], + edges_by_type = { + "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_loaded("semantic", "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-1") is True + assert cache.is_fully_loaded("node-4") is True # Marked even with no edges + assert cache.is_fully_loaded("node-2") is False # Target, not source 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() - edges = { - "node-1": [EdgeTarget("node-2", 0.8), EdgeTarget("node-3", 0.6)], + edges_by_type = { + "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") assert len(neighbors) == 2 @@ -69,28 +69,30 @@ class TestEdgeCache: assert neighbors[0].weight == 0.8 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() - # Load some nodes - cache.add_edges("semantic", {"node-1": []}, ["node-1", "node-2"]) + # Load some nodes (all edge types) + cache.add_all_edges({"semantic": {"node-1": []}}, ["node-1", "node-2"]) # 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"} def test_get_normalized_neighbors_normalizes_weights(self): """get_normalized_neighbors should normalize weights to sum to 1.""" cache = EdgeCache() - edges = { - "node-1": [ - EdgeTarget("node-2", 0.8), - EdgeTarget("node-3", 0.4), - EdgeTarget("node-4", 0.2), - ], + edges_by_type = { + "semantic": { + "node-1": [ + EdgeTarget("node-2", 0.8), + EdgeTarget("node-3", 0.4), + 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 neighbors = cache.get_normalized_neighbors("semantic", "node-1", top_k=2) @@ -111,8 +113,11 @@ class TestEdgeCache: """Different edge types should be stored separately.""" cache = EdgeCache() - cache.add_edges("semantic", {"node-1": [EdgeTarget("node-2", 0.8)]}, ["node-1"]) - cache.add_edges("temporal", {"node-1": [EdgeTarget("node-3", 0.5)]}, ["node-1"]) + edges_by_type = { + "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") temporal_neighbors = cache.get_neighbors("temporal", "node-1") @@ -198,7 +203,6 @@ class TestMPFPTraverseAsync: result = await mpfp_traverse_async( pool=None, # Not used when no seeds - bank_id="test", seeds=[], pattern=["semantic"], config=config, @@ -213,19 +217,18 @@ class TestMPFPTraverseAsync: cache = EdgeCache() config = MPFPConfig(alpha=0.15, threshold=1e-6) - # Pre-populate cache with empty edges for seed - cache.add_edges("semantic", {}, ["seed-1"]) + # Pre-populate cache with empty edges for seed (marks as fully loaded) + cache.add_all_edges({}, ["seed-1"]) seeds = [SeedNode("seed-1", 1.0)] 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, return_value={}, ): result = await mpfp_traverse_async( pool=MagicMock(), - bank_id="test", seeds=seeds, pattern=["semantic"], config=config, @@ -244,24 +247,25 @@ class TestMPFPTraverseAsync: seeds = [SeedNode("seed-1", 1.0)] - # Mock edge loading - async def mock_load_edges(pool, bank_id, edge_type, node_ids): + # Mock edge loading (returns all edge types at once) + async def mock_load_all_edges(pool, node_ids): if "seed-1" in node_ids: return { - "seed-1": [ - EdgeTarget("neighbor-1", 0.8), - EdgeTarget("neighbor-2", 0.4), - ] + "semantic": { + "seed-1": [ + EdgeTarget("neighbor-1", 0.8), + EdgeTarget("neighbor-2", 0.4), + ] + } } return {} with patch( - "hindsight_api.engine.search.mpfp_retrieval.load_edges_for_frontier", - side_effect=mock_load_edges, + "hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier", + side_effect=mock_load_all_edges, ): result = await mpfp_traverse_async( pool=MagicMock(), - bank_id="test", seeds=seeds, pattern=["semantic"], config=config, @@ -287,22 +291,21 @@ class TestMPFPTraverseAsync: seeds = [SeedNode("seed-1", 1.0)] - # Mock edge loading for two hops - async def mock_load_edges(pool, bank_id, edge_type, node_ids): - edges = {} + # Mock edge loading for two hops (returns all edge types at once) + async def mock_load_all_edges(pool, node_ids): + edges: dict[str, dict[str, list[EdgeTarget]]] = {"semantic": {}} 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: - edges["hop1-node"] = [EdgeTarget("hop2-node", 1.0)] + edges["semantic"]["hop1-node"] = [EdgeTarget("hop2-node", 1.0)] return edges with patch( - "hindsight_api.engine.search.mpfp_retrieval.load_edges_for_frontier", - side_effect=mock_load_edges, + "hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier", + side_effect=mock_load_all_edges, ): result = await mpfp_traverse_async( pool=MagicMock(), - bank_id="test", seeds=seeds, pattern=["semantic", "semantic"], # Two hops config=config, @@ -320,27 +323,26 @@ class TestMPFPTraverseAsync: cache = EdgeCache() config = MPFPConfig(alpha=0.15, threshold=1e-6) - # Pre-load cache - cache.add_edges("semantic", {"seed-1": [EdgeTarget("neighbor-1", 1.0)]}, ["seed-1"]) + # Pre-load cache (marks seed-1 as fully loaded) + cache.add_all_edges({"semantic": {"seed-1": [EdgeTarget("neighbor-1", 1.0)]}}, ["seed-1"]) seeds = [SeedNode("seed-1", 1.0)] load_mock = AsyncMock(return_value={}) 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, ): await mpfp_traverse_async( pool=MagicMock(), - bank_id="test", seeds=seeds, pattern=["semantic"], config=config, 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()