fix: improve graph retrieval on large memory banks

.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

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

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

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

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