diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 2788e1f4..18264998 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/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", diff --git a/hindsight-api/hindsight_api/engine/db_budget.py b/hindsight-api/hindsight_api/engine/db_budget.py new file mode 100644 index 00000000..5f50e5f9 --- /dev/null +++ b/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) diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index b02addbb..6577b6fa 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/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,38 +1610,40 @@ 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 - # This prevents slow dateparser from blocking semantic/BM25/graph retrieval - tc_duration = 0.0 # Will be tracked inside temporal retrieval timing - - # Run retrieval for each fact type in parallel - # Each retrieve_parallel uses ~4 connections, so 3 fact types = ~12 concurrent connections - retrieval_tasks = [ - retrieve_parallel( - pool, + # Run optimized retrieval with connection budget + config = get_config() + async with budgeted_operation( + max_connections=config.recall_connection_budget, + operation_id=f"recall-{recall_id}", + ) as op: + budgeted_pool = op.wrap_pool(pool) + parallel_start = time.time() + 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 + parallel_duration = time.time() - parallel_start # Combine all results from all fact types and aggregate timings semantic_results = [] @@ -1657,12 +1660,15 @@ class MemoryEngine(MemoryEngineInterface): all_mpfp_timings = [] detected_temporal_constraint = None - max_conn_wait = 0.0 - for idx, retrieval_result in enumerate(all_retrievals): + max_conn_wait = multi_result.max_conn_wait + 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): - ft_name = fact_type[idx] if idx < len(fact_type) else "unknown" + for ft_name in fact_type: + 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( diff --git a/hindsight-api/hindsight_api/engine/search/retrieval.py b/hindsight-api/hindsight_api/engine/search/retrieval.py index dc379744..32f69e43 100644 --- a/hindsight-api/hindsight_api/engine/search/retrieval.py +++ b/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, + ) diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 9b99b463..30264cbe 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/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, diff --git a/hindsight-api/hindsight_api/migrations.py b/hindsight-api/hindsight_api/migrations.py index 2d444598..3931e06b 100644 --- a/hindsight-api/hindsight_api/migrations.py +++ b/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 - command.upgrade(alembic_cfg, "head") + 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}'")