fix: batch queries on recall (#149)

* fix: batch queries on recall

* fix: batch queries on recall

hindsight-api/hindsight_api/config.py

@@ -64,6 +64,7 @@ ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
 ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
 ENV_MPFP_TOP_K_NEIGHBORS = "HINDSIGHT_API_MPFP_TOP_K_NEIGHBORS"
 ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
+ENV_RECALL_CONNECTION_BUDGET = "HINDSIGHT_API_RECALL_CONNECTION_BUDGET"
 ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID"
 ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
 
@@ -128,6 +129,7 @@ DEFAULT_MCP_ENABLED = True
 DEFAULT_GRAPH_RETRIEVER = "link_expansion"  # Options: "link_expansion", "mpfp", "bfs"
 DEFAULT_MPFP_TOP_K_NEIGHBORS = 20  # Fan-out limit per node in MPFP graph traversal
 DEFAULT_RECALL_MAX_CONCURRENT = 32  # Max concurrent recall operations per worker
+DEFAULT_RECALL_CONNECTION_BUDGET = 4  # Max concurrent DB connections per recall operation
 DEFAULT_MCP_LOCAL_BANK_ID = "mcp"
 
 # Observation thresholds
@@ -241,6 +243,7 @@ class HindsightConfig:
     graph_retriever: str
     mpfp_top_k_neighbors: int
     recall_max_concurrent: int
+    recall_connection_budget: int
 
     # Observation thresholds
     observation_min_facts: int
@@ -315,6 +318,9 @@ class HindsightConfig:
             graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
             mpfp_top_k_neighbors=int(os.getenv(ENV_MPFP_TOP_K_NEIGHBORS, str(DEFAULT_MPFP_TOP_K_NEIGHBORS))),
             recall_max_concurrent=int(os.getenv(ENV_RECALL_MAX_CONCURRENT, str(DEFAULT_RECALL_MAX_CONCURRENT))),
+            recall_connection_budget=int(
+                os.getenv(ENV_RECALL_CONNECTION_BUDGET, str(DEFAULT_RECALL_CONNECTION_BUDGET))
+            ),
             # Optimization flags
             skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
             lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",

hindsight-api/hindsight_api/engine/db_budget.py

@@ -0,0 +1,284 @@
+"""
+Database connection budget management.
+
+Limits concurrent database connections per operation to prevent
+a single operation (e.g., recall with parallel queries) from
+exhausting the connection pool.
+"""
+
+import asyncio
+import logging
+import uuid
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, AsyncIterator
+
+if TYPE_CHECKING:
+    import asyncpg
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OperationBudget:
+    """
+    Tracks connection budget for a single operation.
+
+    Each operation gets a semaphore limiting its concurrent connections.
+    """
+
+    operation_id: str
+    max_connections: int
+    semaphore: asyncio.Semaphore = field(init=False)
+    active_count: int = field(default=0, init=False)
+
+    def __post_init__(self):
+        self.semaphore = asyncio.Semaphore(self.max_connections)
+
+
+class ConnectionBudgetManager:
+    """
+    Manages per-operation connection budgets.
+
+    Usage:
+        manager = ConnectionBudgetManager(default_budget=4)
+
+        # Start an operation
+        async with manager.operation(max_connections=2) as op:
+            # Acquire connections within the budget
+            async with op.acquire(pool) as conn:
+                await conn.fetch(...)
+
+            # Multiple connections respect the budget
+            async with op.acquire(pool) as conn1, op.acquire(pool) as conn2:
+                # At most 2 concurrent connections for this operation
+                ...
+    """
+
+    def __init__(self, default_budget: int = 4):
+        """
+        Initialize the budget manager.
+
+        Args:
+            default_budget: Default max connections per operation
+        """
+        self.default_budget = default_budget
+        self._operations: dict[str, OperationBudget] = {}
+        self._lock = asyncio.Lock()
+
+    @asynccontextmanager
+    async def operation(
+        self,
+        max_connections: int | None = None,
+        operation_id: str | None = None,
+    ) -> AsyncIterator["BudgetedOperation"]:
+        """
+        Create a budgeted operation context.
+
+        Args:
+            max_connections: Max concurrent connections for this operation.
+                           Defaults to manager's default_budget.
+            operation_id: Optional custom operation ID. Auto-generated if not provided.
+
+        Yields:
+            BudgetedOperation context for acquiring connections
+        """
+        op_id = operation_id or f"op-{uuid.uuid4().hex[:12]}"
+        budget = max_connections or self.default_budget
+
+        async with self._lock:
+            if op_id in self._operations:
+                raise ValueError(f"Operation {op_id} already exists")
+            self._operations[op_id] = OperationBudget(op_id, budget)
+
+        try:
+            yield BudgetedOperation(self, op_id)
+        finally:
+            async with self._lock:
+                self._operations.pop(op_id, None)
+
+    def _get_budget(self, operation_id: str) -> OperationBudget:
+        """Get budget for an operation (internal use)."""
+        budget = self._operations.get(operation_id)
+        if not budget:
+            raise ValueError(f"Operation {operation_id} not found")
+        return budget
+
+
+class BudgetedOperation:
+    """
+    A single operation with connection budget.
+
+    Provides methods to acquire connections within the budget.
+    """
+
+    def __init__(self, manager: ConnectionBudgetManager, operation_id: str):
+        self._manager = manager
+        self.operation_id = operation_id
+
+    @property
+    def budget(self) -> OperationBudget:
+        """Get the budget for this operation."""
+        return self._manager._get_budget(self.operation_id)
+
+    @asynccontextmanager
+    async def acquire(self, pool: "asyncpg.Pool") -> AsyncIterator["asyncpg.Connection"]:
+        """
+        Acquire a connection within the operation's budget.
+
+        Blocks if the operation has reached its connection limit.
+
+        Args:
+            pool: asyncpg connection pool
+
+        Yields:
+            Database connection
+        """
+        budget = self.budget
+        async with budget.semaphore:
+            budget.active_count += 1
+            conn = await pool.acquire()
+            try:
+                yield conn
+            finally:
+                budget.active_count -= 1
+                await pool.release(conn)
+
+    def wrap_pool(self, pool: "asyncpg.Pool") -> "BudgetedPool":
+        """
+        Wrap a pool with this operation's budget.
+
+        The returned BudgetedPool can be passed to functions expecting a pool,
+        and all acquire() calls will be limited by this operation's budget.
+
+        Args:
+            pool: asyncpg connection pool to wrap
+
+        Returns:
+            BudgetedPool that limits connections to this operation's budget
+        """
+        return BudgetedPool(pool, self)
+
+    async def acquire_many(
+        self,
+        pool: "asyncpg.Pool",
+        count: int,
+    ) -> AsyncIterator[list["asyncpg.Connection"]]:
+        """
+        Acquire multiple connections within the budget.
+
+        Note: This acquires connections sequentially to respect the budget.
+        For parallel acquisition, use multiple acquire() calls with asyncio.gather().
+
+        Args:
+            pool: asyncpg connection pool
+            count: Number of connections to acquire
+
+        Yields:
+            List of database connections
+        """
+        connections = []
+        try:
+            for _ in range(count):
+                conn = await pool.acquire()
+                connections.append(conn)
+            yield connections
+        finally:
+            for conn in connections:
+                await pool.release(conn)
+
+
+# Global default manager instance
+_default_manager: ConnectionBudgetManager | None = None
+
+
+def get_budget_manager(default_budget: int = 4) -> ConnectionBudgetManager:
+    """
+    Get or create the global budget manager.
+
+    Args:
+        default_budget: Default max connections per operation
+
+    Returns:
+        Global ConnectionBudgetManager instance
+    """
+    global _default_manager
+    if _default_manager is None:
+        _default_manager = ConnectionBudgetManager(default_budget=default_budget)
+    return _default_manager
+
+
+@asynccontextmanager
+async def budgeted_operation(
+    max_connections: int | None = None,
+    operation_id: str | None = None,
+    default_budget: int = 4,
+) -> AsyncIterator[BudgetedOperation]:
+    """
+    Convenience function to create a budgeted operation.
+
+    Args:
+        max_connections: Max concurrent connections for this operation
+        operation_id: Optional custom operation ID
+        default_budget: Default budget if manager not yet created
+
+    Yields:
+        BudgetedOperation context
+
+    Example:
+        async with budgeted_operation(max_connections=2) as op:
+            async with op.acquire(pool) as conn:
+                await conn.fetch(...)
+    """
+    manager = get_budget_manager(default_budget)
+    async with manager.operation(max_connections, operation_id) as op:
+        yield op
+
+
+class BudgetedPool:
+    """
+    A pool wrapper that limits concurrent connection acquisitions.
+
+    This can be passed to functions expecting a pool, and acquire()
+    calls will be limited by the budget semaphore.
+
+    Usage:
+        async with budgeted_operation(max_connections=4) as op:
+            budgeted_pool = op.wrap_pool(pool)
+            # Pass budgeted_pool to functions that expect a pool
+            await some_function(budgeted_pool, ...)
+    """
+
+    def __init__(self, pool: "asyncpg.Pool", operation: BudgetedOperation):
+        self._pool = pool
+        self._operation = operation
+
+    async def acquire(self) -> "asyncpg.Connection":
+        """
+        Acquire a connection within the budget.
+
+        Note: Caller must release the connection when done.
+        Prefer using as context manager via acquire_with_retry or op.acquire().
+        """
+        budget = self._operation.budget
+        await budget.semaphore.acquire()
+        budget.active_count += 1
+        try:
+            return await self._pool.acquire()
+        except Exception:
+            budget.active_count -= 1
+            budget.semaphore.release()
+            raise
+
+    async def release(self, conn: "asyncpg.Connection") -> None:
+        """Release a connection back to the pool."""
+        budget = self._operation.budget
+        try:
+            await self._pool.release(conn)
+        finally:
+            budget.active_count -= 1
+            budget.semaphore.release()
+
+    def __getattr__(self, name):
+        """Proxy other attributes to the underlying pool."""
+        return getattr(self._pool, name)

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any
 
 from ..config import get_config
 from ..metrics import get_metrics_collector
+from .db_budget import budgeted_operation
 
 # Context variable for current schema (async-safe, per-task isolation)
 _current_schema: contextvars.ContextVar[str] = contextvars.ContextVar("current_schema", default="public")
@@ -1609,37 +1610,39 @@ class MemoryEngine(MemoryEngineInterface):
                 tracer.record_query_embedding(query_embedding)
                 tracer.add_phase_metric("generate_query_embedding", step_duration)
 
-            # Step 2: N*4-Way Parallel Retrieval (N fact types × 4 retrieval methods)
+            # Step 2: Optimized parallel retrieval using batched queries
+            # - Semantic + BM25 combined in 1 CTE query for ALL fact types
+            # - Graph runs per fact type (complex traversal)
+            # - Temporal runs per fact type (if constraint detected)
             step_start = time.time()
             query_embedding_str = str(query_embedding)
 
-            from .search.retrieval import get_default_graph_retriever, retrieve_parallel
+            from .search.retrieval import (
+                get_default_graph_retriever,
+                retrieve_all_fact_types_parallel,
+            )
 
             # Track each retrieval start time
             retrieval_start = time.time()
 
-            # Temporal extraction now runs IN PARALLEL with other retrievals inside retrieve_parallel
+            # Run optimized retrieval with connection budget
-            # This prevents slow dateparser from blocking semantic/BM25/graph retrieval
+            config = get_config()
-            tc_duration = 0.0  # Will be tracked inside temporal retrieval timing
+            async with budgeted_operation(
-
+                max_connections=config.recall_connection_budget,
-            # Run retrieval for each fact type in parallel
+                operation_id=f"recall-{recall_id}",
-            # Each retrieve_parallel uses ~4 connections, so 3 fact types = ~12 concurrent connections
+            ) as op:
-            retrieval_tasks = [
+                budgeted_pool = op.wrap_pool(pool)
-                retrieve_parallel(
+                parallel_start = time.time()
-                    pool,
+                multi_result = await retrieve_all_fact_types_parallel(
+                    budgeted_pool,
                     query,
                     query_embedding_str,
                     bank_id,
-                    ft,
+                    fact_type,  # Pass all fact types at once
                     thinking_budget,
                     question_date,
                     self.query_analyzer,
-                    temporal_constraint=None,  # Extracted in parallel inside retrieve_parallel
                 )
-                for ft in fact_type
-            ]
-            parallel_start = time.time()
-            all_retrievals = await asyncio.gather(*retrieval_tasks)
                 parallel_duration = time.time() - parallel_start
 
             # Combine all results from all fact types and aggregate timings
@@ -1657,12 +1660,15 @@ class MemoryEngine(MemoryEngineInterface):
             all_mpfp_timings = []
 
             detected_temporal_constraint = None
-            max_conn_wait = 0.0
+            max_conn_wait = multi_result.max_conn_wait
-            for idx, retrieval_result in enumerate(all_retrievals):
+            for ft in fact_type:
+                retrieval_result = multi_result.results_by_fact_type.get(ft)
+                if not retrieval_result:
+                    continue
+
                 # Log fact types in this retrieval batch
-                ft_name = fact_type[idx] if idx < len(fact_type) else "unknown"
                 logger.debug(
-                    f"[RECALL {recall_id}] Fact type '{ft_name}': semantic={len(retrieval_result.semantic)}, bm25={len(retrieval_result.bm25)}, graph={len(retrieval_result.graph)}, temporal={len(retrieval_result.temporal) if retrieval_result.temporal else 0}"
+                    f"[RECALL {recall_id}] Fact type '{ft}': semantic={len(retrieval_result.semantic)}, bm25={len(retrieval_result.bm25)}, graph={len(retrieval_result.graph)}, temporal={len(retrieval_result.temporal) if retrieval_result.temporal else 0}"
                 )
 
                 semantic_results.extend(retrieval_result.semantic)
@@ -1678,8 +1684,6 @@ class MemoryEngine(MemoryEngineInterface):
                     detected_temporal_constraint = retrieval_result.temporal_constraint
                 # Collect MPFP timings
                 all_mpfp_timings.extend(retrieval_result.mpfp_timings)
-                # Track max connection wait
-                max_conn_wait = max(max_conn_wait, retrieval_result.max_conn_wait)
 
             # If no temporal results from any fact type, set to None
             if not temporal_results:
@@ -1711,10 +1715,8 @@ class MemoryEngine(MemoryEngineInterface):
                 temporal_count = len(temporal_results) if temporal_results else 0
                 timing_parts.append(f"temporal={temporal_count}({aggregated_timings['temporal']:.3f}s)")
                 temporal_info = f" | temporal_range={start_dt.strftime('%Y-%m-%d')} to {end_dt.strftime('%Y-%m-%d')}"
-            # Only tc is sequential setup now (adjacency loads in parallel with retrieval)
-            setup_info = f", tc={tc_duration:.3f}s" if tc_duration > 0.01 else ""
             log_buffer.append(
-                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{temporal_info}"
             )
 
             # Log graph retriever timing breakdown if available
@@ -1751,8 +1753,10 @@ class MemoryEngine(MemoryEngineInterface):
                     return [(r.id, r.__dict__) for r in results]
 
                 # Add retrieval results per fact type (to show parallel execution in UI)
-                for idx, rr in enumerate(all_retrievals):
+                for ft_name in fact_type:
-                    ft_name = fact_type[idx] if idx < len(fact_type) else "unknown"
+                    rr = multi_result.results_by_fact_type.get(ft_name)
+                    if not rr:
+                        continue
 
                     # Add semantic retrieval results for this fact type
                     tracer.add_retrieval_results(

hindsight-api/hindsight_api/engine/search/retrieval.py

@@ -39,6 +39,18 @@ class ParallelRetrievalResult:
     max_conn_wait: float = 0.0  # Maximum connection acquisition wait time across all methods
 
 
+@dataclass
+class MultiFactTypeRetrievalResult:
+    """Result from retrieval across all fact types."""
+
+    # Results per fact type
+    results_by_fact_type: dict[str, ParallelRetrievalResult]
+    # Aggregate timings
+    timings: dict[str, float] = field(default_factory=dict)
+    # Max connection wait across all operations
+    max_conn_wait: float = 0.0
+
+
 # Default graph retriever instance (can be overridden)
 _default_graph_retriever: GraphRetriever | None = None
 
@@ -158,6 +170,356 @@ async def retrieve_bm25(conn, query_text: str, bank_id: str, fact_type: str, lim
     return [RetrievalResult.from_db_row(dict(r)) for r in results]
 
 
+async def retrieve_semantic_bm25_combined(
+    conn,
+    query_emb_str: str,
+    query_text: str,
+    bank_id: str,
+    fact_types: list[str],
+    limit: int,
+) -> dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]]:
+    """
+    Combined semantic + BM25 retrieval for multiple fact types in a single query.
+
+    Uses CTEs with window functions to get top-N results per fact type per method,
+    all in one database round-trip.
+
+    Args:
+        conn: Database connection
+        query_emb_str: Query embedding as string
+        query_text: Query text for BM25
+        bank_id: Bank ID
+        fact_types: List of fact types to retrieve
+        limit: Maximum results per method per fact type
+
+    Returns:
+        Dict mapping fact_type -> (semantic_results, bm25_results)
+    """
+    import re
+
+    # Sanitize query text for BM25 (same as retrieve_bm25)
+    sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
+    tokens = [token for token in sanitized_text.split() if token]
+
+    # If no valid tokens for BM25, just run semantic
+    if not tokens:
+        results = await conn.fetch(
+            f"""
+            WITH semantic_ranked AS (
+                SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                       1 - (embedding <=> $1::vector) AS similarity,
+                       NULL::float AS bm25_score,
+                       'semantic' AS source,
+                       ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
+                FROM {fq_table("memory_units")}
+                WHERE bank_id = $2
+                  AND embedding IS NOT NULL
+                  AND fact_type = ANY($3)
+                  AND (1 - (embedding <=> $1::vector)) >= 0.3
+            )
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   similarity, bm25_score, source
+            FROM semantic_ranked
+            WHERE rn <= $4
+            """,
+            query_emb_str,
+            bank_id,
+            fact_types,
+            limit,
+        )
+        # Group by fact_type
+        result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {
+            ft: ([], []) for ft in fact_types
+        }
+        for r in results:
+            row = dict(r)
+            ft = row.get("fact_type")
+            row.pop("source", None)
+            if ft in result_dict:
+                result_dict[ft][0].append(RetrievalResult.from_db_row(row))
+        return result_dict
+
+    query_tsquery = " | ".join(tokens)
+
+    # Combined CTE query for both semantic and BM25 across all fact types
+    # Uses window functions to limit per fact_type per method
+    results = await conn.fetch(
+        f"""
+        WITH semantic_ranked AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   1 - (embedding <=> $1::vector) AS similarity,
+                   NULL::float AS bm25_score,
+                   'semantic' AS source,
+                   ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
+            FROM {fq_table("memory_units")}
+            WHERE bank_id = $2
+              AND embedding IS NOT NULL
+              AND fact_type = ANY($3)
+              AND (1 - (embedding <=> $1::vector)) >= 0.3
+        ),
+        bm25_ranked AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   NULL::float AS similarity,
+                   ts_rank_cd(search_vector, to_tsquery('english', $5)) AS bm25_score,
+                   'bm25' AS source,
+                   ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY ts_rank_cd(search_vector, to_tsquery('english', $5)) DESC) AS rn
+            FROM {fq_table("memory_units")}
+            WHERE bank_id = $2
+              AND fact_type = ANY($3)
+              AND search_vector @@ to_tsquery('english', $5)
+        ),
+        semantic AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   similarity, bm25_score, source
+            FROM semantic_ranked WHERE rn <= $4
+        ),
+        bm25 AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   similarity, bm25_score, source
+            FROM bm25_ranked WHERE rn <= $4
+        )
+        SELECT * FROM semantic
+        UNION ALL
+        SELECT * FROM bm25
+        """,
+        query_emb_str,
+        bank_id,
+        fact_types,
+        limit,
+        query_tsquery,
+    )
+
+    # Group results by fact_type and source
+    result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types}
+    for r in results:
+        row = dict(r)
+        source = row.pop("source", None)
+        ft = row.get("fact_type")
+        if ft in result_dict:
+            if source == "semantic":
+                result_dict[ft][0].append(RetrievalResult.from_db_row(row))
+            else:
+                result_dict[ft][1].append(RetrievalResult.from_db_row(row))
+
+    return result_dict
+
+
+async def retrieve_temporal_combined(
+    conn,
+    query_emb_str: str,
+    bank_id: str,
+    fact_types: list[str],
+    start_date: datetime,
+    end_date: datetime,
+    budget: int,
+    semantic_threshold: float = 0.1,
+) -> dict[str, list[RetrievalResult]]:
+    """
+    Temporal retrieval for multiple fact types in a single query.
+
+    Batches the entry point query using window functions to get top-N per fact type,
+    then runs spreading for each fact type.
+
+    Args:
+        conn: Database connection
+        query_emb_str: Query embedding as string
+        bank_id: Bank ID
+        fact_types: List of fact types to retrieve
+        start_date: Start of time range
+        end_date: End of time range
+        budget: Node budget for spreading per fact type
+        semantic_threshold: Minimum semantic similarity to include
+
+    Returns:
+        Dict mapping fact_type -> list of RetrievalResult
+    """
+    from ..memory_engine import fq_table
+
+    # Ensure dates are timezone-aware
+    if start_date.tzinfo is None:
+        start_date = start_date.replace(tzinfo=UTC)
+    if end_date.tzinfo is None:
+        end_date = end_date.replace(tzinfo=UTC)
+
+    # Batch query: Get entry points for ALL fact types at once with window function
+    entry_points = await conn.fetch(
+        f"""
+        WITH ranked_entries AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   1 - (embedding <=> $1::vector) AS similarity,
+                   ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, embedding <=> $1::vector) AS rn
+            FROM {fq_table("memory_units")}
+            WHERE bank_id = $2
+              AND fact_type = ANY($3)
+              AND embedding IS NOT NULL
+              AND (
+                  (occurred_start IS NOT NULL AND occurred_end IS NOT NULL
+                   AND occurred_start <= $5 AND occurred_end >= $4)
+                  OR
+                  (mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
+                  OR
+                  (occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
+                  OR
+                  (occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
+              )
+              AND (1 - (embedding <=> $1::vector)) >= $6
+        )
+        SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id, similarity
+        FROM ranked_entries
+        WHERE rn <= 10
+        """,
+        query_emb_str,
+        bank_id,
+        fact_types,
+        start_date,
+        end_date,
+        semantic_threshold,
+    )
+
+    if not entry_points:
+        return {ft: [] for ft in fact_types}
+
+    # Group entry points by fact type
+    entries_by_ft: dict[str, list] = {ft: [] for ft in fact_types}
+    for ep in entry_points:
+        ft = ep["fact_type"]
+        if ft in entries_by_ft:
+            entries_by_ft[ft].append(ep)
+
+    # Calculate shared temporal parameters
+    total_days = (end_date - start_date).total_seconds() / 86400
+    mid_date = start_date + (end_date - start_date) / 2
+
+    # Process each fact type (spreading needs to stay per fact type due to link filtering)
+    results_by_ft: dict[str, list[RetrievalResult]] = {}
+
+    for ft in fact_types:
+        ft_entry_points = entries_by_ft.get(ft, [])
+        if not ft_entry_points:
+            results_by_ft[ft] = []
+            continue
+
+        results = []
+        visited = set()
+        node_scores = {}
+
+        # Process entry points
+        for ep in ft_entry_points:
+            unit_id = str(ep["id"])
+            visited.add(unit_id)
+
+            # Calculate temporal proximity
+            best_date = None
+            if ep["occurred_start"] is not None and ep["occurred_end"] is not None:
+                best_date = ep["occurred_start"] + (ep["occurred_end"] - ep["occurred_start"]) / 2
+            elif ep["occurred_start"] is not None:
+                best_date = ep["occurred_start"]
+            elif ep["occurred_end"] is not None:
+                best_date = ep["occurred_end"]
+            elif ep["mentioned_at"] is not None:
+                best_date = ep["mentioned_at"]
+
+            if best_date:
+                days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
+                temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
+            else:
+                temporal_proximity = 0.5
+
+            ep_result = RetrievalResult.from_db_row(dict(ep))
+            ep_result.temporal_score = temporal_proximity
+            ep_result.temporal_proximity = temporal_proximity
+            results.append(ep_result)
+            node_scores[unit_id] = (ep["similarity"], 1.0)
+
+        # Spreading through temporal links (same as single-fact-type version)
+        frontier = list(node_scores.keys())
+        budget_remaining = budget - len(ft_entry_points)
+        batch_size = 20
+
+        while frontier and budget_remaining > 0:
+            batch_ids = frontier[:batch_size]
+            frontier = frontier[batch_size:]
+
+            neighbors = await conn.fetch(
+                f"""
+                SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id, mu.chunk_id,
+                       ml.weight, ml.link_type, ml.from_unit_id,
+                       1 - (mu.embedding <=> $1::vector) AS similarity
+                FROM {fq_table("memory_links")} ml
+                JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
+                WHERE ml.from_unit_id = ANY($2::uuid[])
+                  AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
+                  AND ml.weight >= 0.1
+                  AND mu.fact_type = $3
+                  AND mu.embedding IS NOT NULL
+                  AND (1 - (mu.embedding <=> $1::vector)) >= $4
+                ORDER BY ml.weight DESC
+                LIMIT $5
+                """,
+                query_emb_str,
+                batch_ids,
+                ft,
+                semantic_threshold,
+                batch_size * 10,
+            )
+
+            for n in neighbors:
+                neighbor_id = str(n["id"])
+                if neighbor_id in visited:
+                    continue
+
+                visited.add(neighbor_id)
+                budget_remaining -= 1
+
+                parent_id = str(n["from_unit_id"])
+                _, parent_temporal_score = node_scores.get(parent_id, (0.5, 0.5))
+
+                neighbor_best_date = None
+                if n["occurred_start"] is not None and n["occurred_end"] is not None:
+                    neighbor_best_date = n["occurred_start"] + (n["occurred_end"] - n["occurred_start"]) / 2
+                elif n["occurred_start"] is not None:
+                    neighbor_best_date = n["occurred_start"]
+                elif n["occurred_end"] is not None:
+                    neighbor_best_date = n["occurred_end"]
+                elif n["mentioned_at"] is not None:
+                    neighbor_best_date = n["mentioned_at"]
+
+                if neighbor_best_date:
+                    days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
+                    neighbor_temporal_proximity = (
+                        1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
+                    )
+                else:
+                    neighbor_temporal_proximity = 0.3
+
+                link_type = n["link_type"]
+                if link_type in ("causes", "caused_by"):
+                    causal_boost = 2.0
+                elif link_type in ("enables", "prevents"):
+                    causal_boost = 1.5
+                else:
+                    causal_boost = 1.0
+
+                propagated_temporal = parent_temporal_score * n["weight"] * causal_boost * 0.7
+                combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)
+
+                neighbor_result = RetrievalResult.from_db_row(dict(n))
+                neighbor_result.temporal_score = combined_temporal
+                neighbor_result.temporal_proximity = neighbor_temporal_proximity
+                results.append(neighbor_result)
+
+                if budget_remaining > 0 and combined_temporal > 0.2:
+                    node_scores[neighbor_id] = (n["similarity"], combined_temporal)
+                    frontier.append(neighbor_id)
+
+                if budget_remaining <= 0:
+                    break
+
+        results_by_ft[ft] = results
+
+    return results_by_ft
+
+
 async def retrieve_temporal(
     conn,
     query_emb_str: str,
@@ -748,3 +1110,158 @@ async def _retrieve_parallel_bfs(
             },
             temporal_constraint=None,
         )
+
+
+async def retrieve_all_fact_types_parallel(
+    pool,
+    query_text: str,
+    query_embedding_str: str,
+    bank_id: str,
+    fact_types: list[str],
+    thinking_budget: int,
+    question_date: datetime | None = None,
+    query_analyzer: Optional["QueryAnalyzer"] = None,
+    graph_retriever: GraphRetriever | None = None,
+) -> MultiFactTypeRetrievalResult:
+    """
+    Optimized retrieval for multiple fact types using batched queries.
+
+    This reduces database round-trips by:
+    1. Combining semantic + BM25 into one CTE query for ALL fact types (1 query instead of 2N)
+    2. Running graph retrieval per fact type in parallel (N parallel tasks)
+    3. Running temporal retrieval per fact type in parallel (N parallel tasks)
+
+    Args:
+        pool: Database connection pool
+        query_text: Query text
+        query_embedding_str: Query embedding as string
+        bank_id: Bank ID
+        fact_types: List of fact types to retrieve
+        thinking_budget: Budget for graph traversal and retrieval limits
+        question_date: Optional date when question was asked (for temporal filtering)
+        query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
+        graph_retriever: Graph retrieval strategy (defaults to configured retriever)
+
+    Returns:
+        MultiFactTypeRetrievalResult with results organized by fact type
+    """
+    import time
+
+    retriever = graph_retriever or get_default_graph_retriever()
+    start_time = time.time()
+    timings: dict[str, float] = {}
+
+    # Step 1: Extract temporal constraint first (CPU work, no DB)
+    # Do this before DB queries so we know if we need temporal retrieval
+    temporal_extraction_start = time.time()
+    from .temporal_extraction import extract_temporal_constraint
+
+    temporal_constraint = extract_temporal_constraint(query_text, reference_date=question_date, analyzer=query_analyzer)
+    temporal_extraction_time = time.time() - temporal_extraction_start
+    timings["temporal_extraction"] = temporal_extraction_time
+
+    # Step 2: Run semantic + BM25 + temporal combined in ONE connection!
+    # This reduces connection usage from 2 to 1 for these operations
+    semantic_bm25_start = time.time()
+    temporal_results_by_ft: dict[str, list[RetrievalResult]] = {}
+    temporal_time = 0.0
+
+    async with acquire_with_retry(pool) as conn:
+        conn_wait = time.time() - semantic_bm25_start
+
+        # Semantic + BM25 combined
+        semantic_bm25_results = await retrieve_semantic_bm25_combined(
+            conn, query_embedding_str, query_text, bank_id, fact_types, thinking_budget
+        )
+        semantic_bm25_time = time.time() - semantic_bm25_start
+
+        # Temporal combined (if constraint detected) - same connection!
+        if temporal_constraint:
+            tc_start, tc_end = temporal_constraint
+            temporal_start = time.time()
+            temporal_results_by_ft = await retrieve_temporal_combined(
+                conn,
+                query_embedding_str,
+                bank_id,
+                fact_types,
+                tc_start,
+                tc_end,
+                budget=thinking_budget,
+                semantic_threshold=0.1,
+            )
+            temporal_time = time.time() - temporal_start
+
+    timings["semantic_bm25_combined"] = semantic_bm25_time
+    timings["temporal_combined"] = temporal_time
+
+    # Step 3: Run graph retrieval for each fact type in parallel
+    async def run_graph_for_fact_type(ft: str) -> tuple[str, list[RetrievalResult], float, MPFPTimings | None]:
+        graph_start = time.time()
+        results, mpfp_timing = await retriever.retrieve(
+            pool=pool,
+            query_embedding_str=query_embedding_str,
+            bank_id=bank_id,
+            fact_type=ft,
+            budget=thinking_budget,
+            query_text=query_text,
+            semantic_seeds=None,
+            temporal_seeds=None,
+        )
+        return ft, results, time.time() - graph_start, mpfp_timing
+
+    # Run graph for all fact types in parallel
+    graph_tasks = [run_graph_for_fact_type(ft) for ft in fact_types]
+    graph_results_list = await asyncio.gather(*graph_tasks)
+
+    # Organize results by fact type
+    results_by_fact_type: dict[str, ParallelRetrievalResult] = {}
+    max_conn_wait = conn_wait  # Single connection for semantic+bm25+temporal
+    all_mpfp_timings: list[MPFPTimings] = []
+
+    for ft in fact_types:
+        # Get semantic + bm25 results for this fact type
+        semantic_results, bm25_results = semantic_bm25_results.get(ft, ([], []))
+
+        # Find graph results for this fact type
+        graph_results = []
+        graph_time = 0.0
+        mpfp_timing = None
+        for gr in graph_results_list:
+            if gr[0] == ft:
+                graph_results = gr[1]
+                graph_time = gr[2]
+                mpfp_timing = gr[3]
+                if mpfp_timing:
+                    all_mpfp_timings.append(mpfp_timing)
+                break
+
+        # Get temporal results for this fact type from combined result
+        temporal_results = temporal_results_by_ft.get(ft) if temporal_constraint else None
+        if temporal_results is not None and len(temporal_results) == 0:
+            temporal_results = None
+
+        results_by_fact_type[ft] = ParallelRetrievalResult(
+            semantic=semantic_results,
+            bm25=bm25_results,
+            graph=graph_results,
+            temporal=temporal_results,
+            timings={
+                "semantic": semantic_bm25_time / 2,  # Approximate split
+                "bm25": semantic_bm25_time / 2,
+                "graph": graph_time,
+                "temporal": temporal_time,  # Same for all fact types (single query)
+                "temporal_extraction": temporal_extraction_time,
+            },
+            temporal_constraint=temporal_constraint,
+            mpfp_timings=[mpfp_timing] if mpfp_timing else [],
+            max_conn_wait=max_conn_wait,
+        )
+
+    total_time = time.time() - start_time
+    timings["total"] = total_time
+
+    return MultiFactTypeRetrievalResult(
+        results_by_fact_type=results_by_fact_type,
+        timings=timings,
+        max_conn_wait=max_conn_wait,
+    )

hindsight-api/hindsight_api/main.py

@@ -200,6 +200,7 @@ def main():
             graph_retriever=config.graph_retriever,
             mpfp_top_k_neighbors=config.mpfp_top_k_neighbors,
             recall_max_concurrent=config.recall_max_concurrent,
+            recall_connection_budget=config.recall_connection_budget,
             observation_min_facts=config.observation_min_facts,
             observation_top_entities=config.observation_top_entities,
             retain_max_completion_tokens=config.retain_max_completion_tokens,

hindsight-api/hindsight_api/migrations.py

@@ -22,6 +22,7 @@ from pathlib import Path
 
 from alembic import command
 from alembic.config import Config
+from alembic.script.revision import ResolutionError
 from sqlalchemy import create_engine, text
 
 logger = logging.getLogger(__name__)
@@ -78,7 +79,18 @@ def _run_migrations_internal(database_url: str, script_location: str, schema: st
         alembic_cfg.set_main_option("target_schema", schema)
 
     # Run migrations
+    try:
         command.upgrade(alembic_cfg, "head")
+    except ResolutionError as e:
+        # This happens during rolling deployments when a newer version of the code
+        # has already run migrations, and this older replica doesn't have the new
+        # migration files. The database is already at a newer revision than we know.
+        # This is safe to ignore - the newer code has already applied its migrations.
+        logger.warning(
+            f"Database is at a newer migration revision than this code version knows about. "
+            f"This is expected during rolling deployments. Skipping migrations. Error: {e}"
+        )
+        return
 
     logger.info(f"Database migrations completed successfully for schema '{schema_name}'")