diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 03d17f2f..d320f9c3 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -29,6 +29,7 @@ ENV_HOST = "HINDSIGHT_API_HOST" ENV_PORT = "HINDSIGHT_API_PORT" ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL" ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED" +ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER" # Default values DEFAULT_DATABASE_URL = "pg0" @@ -45,6 +46,7 @@ DEFAULT_HOST = "0.0.0.0" DEFAULT_PORT = 8888 DEFAULT_LOG_LEVEL = "info" DEFAULT_MCP_ENABLED = True +DEFAULT_GRAPH_RETRIEVER = "bfs" # Options: "bfs", "mpfp" # Required embedding dimension for database schema EMBEDDING_DIMENSION = 384 @@ -79,6 +81,9 @@ class HindsightConfig: log_level: str mcp_enabled: bool + # Recall + graph_retriever: str + @classmethod def from_env(cls) -> "HindsightConfig": """Create configuration from environment variables.""" @@ -107,6 +112,9 @@ class HindsightConfig: port=int(os.getenv(ENV_PORT, DEFAULT_PORT)), log_level=os.getenv(ENV_LOG_LEVEL, DEFAULT_LOG_LEVEL), mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true", + + # Recall + graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER), ) def get_llm_base_url(self) -> str: @@ -147,6 +155,7 @@ class HindsightConfig: logger.info(f"LLM: provider={self.llm_provider}, model={self.llm_model}") logger.info(f"Embeddings: provider={self.embeddings_provider}") logger.info(f"Reranker: provider={self.reranker_provider}") + logger.info(f"Graph retriever: {self.graph_retriever}") def get_config() -> HindsightConfig: diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 0b88a272..7dcfc093 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -1156,22 +1156,22 @@ class MemoryEngine: aggregated_timings = {"semantic": 0.0, "bm25": 0.0, "graph": 0.0, "temporal": 0.0} detected_temporal_constraint = None - for idx, (ft_semantic, ft_bm25, ft_graph, ft_temporal, ft_timings, ft_temporal_constraint) in enumerate(all_retrievals): + for idx, retrieval_result in enumerate(all_retrievals): # 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(ft_semantic)}, bm25={len(ft_bm25)}, graph={len(ft_graph)}, temporal={len(ft_temporal) if ft_temporal else 0}") + 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}") - semantic_results.extend(ft_semantic) - bm25_results.extend(ft_bm25) - graph_results.extend(ft_graph) - if ft_temporal: - temporal_results.extend(ft_temporal) + semantic_results.extend(retrieval_result.semantic) + bm25_results.extend(retrieval_result.bm25) + graph_results.extend(retrieval_result.graph) + if retrieval_result.temporal: + temporal_results.extend(retrieval_result.temporal) # Track max timing for each method (since they run in parallel across fact types) - for method, duration in ft_timings.items(): - aggregated_timings[method] = max(aggregated_timings[method], duration) + for method, duration in retrieval_result.timings.items(): + aggregated_timings[method] = max(aggregated_timings.get(method, 0.0), duration) # Capture temporal constraint (same across all fact types) - if ft_temporal_constraint: - detected_temporal_constraint = ft_temporal_constraint + if retrieval_result.temporal_constraint: + detected_temporal_constraint = retrieval_result.temporal_constraint # If no temporal results from any fact type, set to None if not temporal_results: @@ -1210,14 +1210,14 @@ class MemoryEngine: return [(r.id, r.__dict__) for r in results] # Add retrieval results per fact type (to show parallel execution in UI) - for idx, (ft_semantic, ft_bm25, ft_graph, ft_temporal, ft_timings, _) in enumerate(all_retrievals): + for idx, rr in enumerate(all_retrievals): ft_name = fact_type[idx] if idx < len(fact_type) else "unknown" # Add semantic retrieval results for this fact type tracer.add_retrieval_results( method_name="semantic", - results=to_tuple_format(ft_semantic), - duration_seconds=ft_timings.get("semantic", 0.0), + results=to_tuple_format(rr.semantic), + duration_seconds=rr.timings.get("semantic", 0.0), score_field="similarity", metadata={"limit": thinking_budget}, fact_type=ft_name @@ -1226,8 +1226,8 @@ class MemoryEngine: # Add BM25 retrieval results for this fact type tracer.add_retrieval_results( method_name="bm25", - results=to_tuple_format(ft_bm25), - duration_seconds=ft_timings.get("bm25", 0.0), + results=to_tuple_format(rr.bm25), + duration_seconds=rr.timings.get("bm25", 0.0), score_field="bm25_score", metadata={"limit": thinking_budget}, fact_type=ft_name @@ -1236,19 +1236,19 @@ class MemoryEngine: # Add graph retrieval results for this fact type tracer.add_retrieval_results( method_name="graph", - results=to_tuple_format(ft_graph), - duration_seconds=ft_timings.get("graph", 0.0), + results=to_tuple_format(rr.graph), + duration_seconds=rr.timings.get("graph", 0.0), score_field="activation", metadata={"budget": thinking_budget}, fact_type=ft_name ) # Add temporal retrieval results for this fact type (even if empty, to show it ran) - if ft_temporal is not None: + if rr.temporal is not None: tracer.add_retrieval_results( method_name="temporal", - results=to_tuple_format(ft_temporal), - duration_seconds=ft_timings.get("temporal", 0.0), + results=to_tuple_format(rr.temporal), + duration_seconds=rr.timings.get("temporal", 0.0), score_field="temporal_score", metadata={"budget": thinking_budget}, fact_type=ft_name diff --git a/hindsight-api/hindsight_api/engine/search/__init__.py b/hindsight-api/hindsight_api/engine/search/__init__.py index 20d3c8a0..5dbe2226 100644 --- a/hindsight-api/hindsight_api/engine/search/__init__.py +++ b/hindsight-api/hindsight_api/engine/search/__init__.py @@ -11,15 +11,19 @@ from .retrieval import ( retrieve_parallel, get_default_graph_retriever, set_default_graph_retriever, + ParallelRetrievalResult, ) from .graph_retrieval import GraphRetriever, BFSGraphRetriever +from .mpfp_retrieval import MPFPGraphRetriever from .reranking import CrossEncoderReranker __all__ = [ "retrieve_parallel", "get_default_graph_retriever", "set_default_graph_retriever", + "ParallelRetrievalResult", "GraphRetriever", "BFSGraphRetriever", + "MPFPGraphRetriever", "CrossEncoderReranker", ] diff --git a/hindsight-api/hindsight_api/engine/search/graph_retrieval.py b/hindsight-api/hindsight_api/engine/search/graph_retrieval.py index 2ebfa3cd..9d332355 100644 --- a/hindsight-api/hindsight_api/engine/search/graph_retrieval.py +++ b/hindsight-api/hindsight_api/engine/search/graph_retrieval.py @@ -29,7 +29,7 @@ class GraphRetriever(ABC): @property @abstractmethod def name(self) -> str: - """Return identifier for this retrieval strategy (e.g., 'bfs', 'ppr').""" + """Return identifier for this retrieval strategy (e.g., 'bfs', 'mpfp').""" pass @abstractmethod @@ -41,6 +41,8 @@ class GraphRetriever(ABC): fact_type: str, budget: int, query_text: Optional[str] = None, + semantic_seeds: Optional[List[RetrievalResult]] = None, + temporal_seeds: Optional[List[RetrievalResult]] = None, ) -> List[RetrievalResult]: """ Retrieve relevant facts via graph traversal. @@ -52,6 +54,8 @@ class GraphRetriever(ABC): fact_type: Fact type to filter ('world', 'experience', 'opinion', 'observation') budget: Maximum number of nodes to explore/return query_text: Original query text (optional, for some strategies) + semantic_seeds: Pre-computed semantic entry points (from semantic retrieval) + temporal_seeds: Pre-computed temporal entry points (from temporal retrieval) Returns: List of RetrievalResult objects with activation scores set @@ -106,6 +110,8 @@ class BFSGraphRetriever(GraphRetriever): fact_type: str, budget: int, query_text: Optional[str] = None, + semantic_seeds: Optional[List[RetrievalResult]] = None, + temporal_seeds: Optional[List[RetrievalResult]] = None, ) -> List[RetrievalResult]: """ Retrieve facts using BFS spreading activation. @@ -115,6 +121,10 @@ class BFSGraphRetriever(GraphRetriever): 2. BFS traversal: visit neighbors, propagate decaying activation 3. Boost causal links (causes, enables, prevents) 4. Return visited nodes up to budget + + Note: BFS finds its own entry points via embedding search. + The semantic_seeds and temporal_seeds parameters are accepted + for interface compatibility but not used. """ async with acquire_with_retry(pool) as conn: return await self._retrieve_with_conn( diff --git a/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py b/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py new file mode 100644 index 00000000..a8f3982d --- /dev/null +++ b/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py @@ -0,0 +1,454 @@ +""" +Meta-Path Forward Push (MPFP) graph retrieval. + +A sublinear graph traversal algorithm for memory retrieval over heterogeneous +graphs with multiple edge types (semantic, temporal, causal, entity). + +Combines meta-path patterns from HIN literature with Forward Push local +propagation from Approximate PPR. + +Key properties: +- Sublinear in graph size (threshold pruning bounds active nodes) +- Predefined patterns capture different retrieval intents +- All patterns run in parallel, results fused via RRF +- No LLM in the loop during traversal +""" + +import asyncio +import logging +from dataclasses import dataclass, field +from typing import List, Dict, Optional, Tuple +from collections import defaultdict + +from .types import RetrievalResult +from .graph_retrieval import GraphRetriever +from ..db_utils import acquire_with_retry + +logger = logging.getLogger(__name__) + + +# ----------------------------------------------------------------------------- +# Data Classes +# ----------------------------------------------------------------------------- + +@dataclass +class EdgeTarget: + """A neighbor node with its edge weight.""" + node_id: str + weight: float + + +@dataclass +class TypedAdjacency: + """Adjacency lists split by edge type.""" + # edge_type -> from_node_id -> list of (to_node_id, weight) + graphs: Dict[str, Dict[str, List[EdgeTarget]]] = field(default_factory=dict) + + def get_neighbors(self, edge_type: str, node_id: str) -> List[EdgeTarget]: + """Get neighbors for a node via a specific edge type.""" + return self.graphs.get(edge_type, {}).get(node_id, []) + + def get_normalized_neighbors( + self, + edge_type: str, + node_id: str, + top_k: int + ) -> List[EdgeTarget]: + """Get top-k neighbors with weights normalized to sum to 1.""" + neighbors = self.get_neighbors(edge_type, node_id)[:top_k] + if not neighbors: + return [] + + total = sum(n.weight for n in neighbors) + if total == 0: + return [] + + return [ + EdgeTarget(node_id=n.node_id, weight=n.weight / total) + for n in neighbors + ] + + +@dataclass +class PatternResult: + """Result from a single pattern traversal.""" + pattern: List[str] + scores: Dict[str, float] # node_id -> accumulated mass + + +@dataclass +class MPFPConfig: + """Configuration for MPFP algorithm.""" + alpha: float = 0.15 # teleport/keep probability + threshold: float = 1e-6 # mass pruning threshold (lower = explore more) + top_k_neighbors: int = 20 # fan-out limit per node + + # Patterns from semantic seeds + patterns_semantic: List[List[str]] = field(default_factory=lambda: [ + ['semantic', 'semantic'], # topic expansion + ['entity', 'temporal'], # entity timeline + ['semantic', 'causes'], # reasoning chains (forward) + ['semantic', 'caused_by'], # reasoning chains (backward) + ['entity', 'semantic'], # entity context + ]) + + # Patterns from temporal seeds + patterns_temporal: List[List[str]] = field(default_factory=lambda: [ + ['temporal', 'semantic'], # what was happening then + ['temporal', 'entity'], # who was involved then + ]) + + +@dataclass +class SeedNode: + """An entry point node with its initial score.""" + node_id: str + score: float # initial mass (e.g., similarity score) + + +# ----------------------------------------------------------------------------- +# Core Algorithm +# ----------------------------------------------------------------------------- + +def mpfp_traverse( + seeds: List[SeedNode], + pattern: List[str], + adjacency: TypedAdjacency, + config: MPFPConfig, +) -> PatternResult: + """ + Forward Push traversal following a meta-path pattern. + + Args: + seeds: Entry point nodes with initial scores + pattern: Sequence of edge types to follow + adjacency: Typed adjacency structure + config: Algorithm parameters + + Returns: + PatternResult with accumulated scores per node + """ + if not seeds: + return PatternResult(pattern=pattern, scores={}) + + scores: Dict[str, float] = {} + + # Initialize frontier with seed masses (normalized) + total_seed_score = sum(s.score for s in seeds) + if total_seed_score == 0: + total_seed_score = len(seeds) # fallback to uniform + + frontier: Dict[str, float] = { + s.node_id: s.score / total_seed_score for s in seeds + } + + # Follow pattern hop by hop + for edge_type in pattern: + next_frontier: Dict[str, float] = {} + + for node_id, mass in frontier.items(): + if mass < config.threshold: + continue + + # Keep α portion for this node + scores[node_id] = scores.get(node_id, 0) + config.alpha * mass + + # Push (1-α) to neighbors + push_mass = (1 - config.alpha) * mass + neighbors = adjacency.get_normalized_neighbors( + edge_type, node_id, config.top_k_neighbors + ) + + for neighbor in neighbors: + next_frontier[neighbor.node_id] = ( + next_frontier.get(neighbor.node_id, 0) + + push_mass * neighbor.weight + ) + + frontier = next_frontier + + # Final frontier nodes get their remaining mass + for node_id, mass in frontier.items(): + if mass >= config.threshold: + scores[node_id] = scores.get(node_id, 0) + mass + + return PatternResult(pattern=pattern, scores=scores) + + +def rrf_fusion( + results: List[PatternResult], + k: int = 60, + top_k: int = 50, +) -> List[Tuple[str, float]]: + """ + Reciprocal Rank Fusion to combine pattern results. + + Args: + results: List of pattern results + k: RRF constant (higher = more uniform weighting) + top_k: Number of results to return + + Returns: + List of (node_id, fused_score) tuples, sorted by score descending + """ + fused: Dict[str, float] = {} + + for result in results: + if not result.scores: + continue + + # Rank nodes by their score in this pattern + ranked = sorted( + result.scores.keys(), + key=lambda n: result.scores[n], + reverse=True + ) + + for rank, node_id in enumerate(ranked): + fused[node_id] = fused.get(node_id, 0) + 1.0 / (k + rank + 1) + + # Sort by fused score and return top-k + sorted_results = sorted( + fused.items(), + key=lambda x: x[1], + reverse=True + ) + + return sorted_results[:top_k] + + +# ----------------------------------------------------------------------------- +# Database Loading +# ----------------------------------------------------------------------------- + +async def load_typed_adjacency(pool, bank_id: str) -> TypedAdjacency: + """ + Load all edges for a bank, split by edge type. + + Single query, then organize in-memory for fast traversal. + """ + async with acquire_with_retry(pool) as conn: + rows = await conn.fetch( + """ + SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight + FROM memory_links ml + JOIN memory_units mu ON ml.from_unit_id = mu.id + WHERE mu.bank_id = $1 + AND ml.weight >= 0.1 + ORDER BY ml.from_unit_id, ml.weight DESC + """, + bank_id + ) + + graphs: Dict[str, Dict[str, List[EdgeTarget]]] = defaultdict( + lambda: defaultdict(list) + ) + + for row in rows: + from_id = str(row['from_unit_id']) + to_id = str(row['to_unit_id']) + link_type = row['link_type'] + weight = row['weight'] + + graphs[link_type][from_id].append( + EdgeTarget(node_id=to_id, weight=weight) + ) + + return TypedAdjacency(graphs=dict(graphs)) + + +async def fetch_memory_units_by_ids( + pool, + node_ids: List[str], + fact_type: str, +) -> List[RetrievalResult]: + """Fetch full memory unit details for a list of node IDs.""" + if not node_ids: + return [] + + async with acquire_with_retry(pool) as conn: + rows = await conn.fetch( + """ + SELECT id, text, context, event_date, occurred_start, occurred_end, + mentioned_at, access_count, embedding, fact_type, document_id, chunk_id + FROM memory_units + WHERE id = ANY($1::uuid[]) + AND fact_type = $2 + """, + node_ids, + fact_type + ) + + return [RetrievalResult.from_db_row(dict(r)) for r in rows] + + +# ----------------------------------------------------------------------------- +# Graph Retriever Implementation +# ----------------------------------------------------------------------------- + +class MPFPGraphRetriever(GraphRetriever): + """ + Graph retrieval using Meta-Path Forward Push. + + Runs predefined patterns in parallel from semantic and temporal seeds, + then fuses results via RRF. + """ + + def __init__(self, config: Optional[MPFPConfig] = None): + """ + Initialize MPFP retriever. + + Args: + config: Algorithm configuration (uses defaults if None) + """ + self.config = config or MPFPConfig() + self._adjacency_cache: Dict[str, TypedAdjacency] = {} + + @property + def name(self) -> str: + return "mpfp" + + async def retrieve( + self, + pool, + query_embedding_str: str, + bank_id: str, + fact_type: str, + budget: int, + query_text: Optional[str] = None, + semantic_seeds: Optional[List[RetrievalResult]] = None, + temporal_seeds: Optional[List[RetrievalResult]] = None, + ) -> List[RetrievalResult]: + """ + Retrieve facts using MPFP algorithm. + + Args: + pool: Database connection pool + query_embedding_str: Query embedding (used for fallback seed finding) + bank_id: Memory bank ID + fact_type: Fact type to filter + budget: Maximum results to return + query_text: Original query text (optional) + semantic_seeds: Pre-computed semantic entry points + temporal_seeds: Pre-computed temporal entry points + + Returns: + List of RetrievalResult with activation scores + """ + # Load typed adjacency (could cache per bank_id with TTL) + adjacency = await load_typed_adjacency(pool, bank_id) + + # Convert seeds to SeedNode format + semantic_seed_nodes = self._convert_seeds(semantic_seeds, 'similarity') + temporal_seed_nodes = self._convert_seeds(temporal_seeds, 'temporal_score') + + # If no semantic seeds provided, fall back to finding our own + if not semantic_seed_nodes: + semantic_seed_nodes = await self._find_semantic_seeds( + pool, query_embedding_str, bank_id, fact_type + ) + + # Run all patterns in parallel + tasks = [] + + # Patterns from semantic seeds + for pattern in self.config.patterns_semantic: + if semantic_seed_nodes: + tasks.append( + asyncio.to_thread( + mpfp_traverse, + semantic_seed_nodes, + pattern, + adjacency, + self.config, + ) + ) + + # Patterns from temporal seeds + for pattern in self.config.patterns_temporal: + if temporal_seed_nodes: + tasks.append( + asyncio.to_thread( + mpfp_traverse, + temporal_seed_nodes, + pattern, + adjacency, + self.config, + ) + ) + + if not tasks: + return [] + + # Gather pattern results + pattern_results = await asyncio.gather(*tasks) + + # Fuse results + fused = rrf_fusion(pattern_results, top_k=budget) + + if not fused: + return [] + + # Get top result IDs (don't exclude seeds - they may be highly relevant) + result_ids = [node_id for node_id, score in fused][:budget] + + # Fetch full details + results = await fetch_memory_units_by_ids(pool, result_ids, fact_type) + + # Add activation scores from fusion + score_map = {node_id: score for node_id, score in fused} + for result in results: + result.activation = score_map.get(result.id, 0.0) + + # Sort by activation + results.sort(key=lambda r: r.activation or 0, reverse=True) + + return results + + def _convert_seeds( + self, + seeds: Optional[List[RetrievalResult]], + score_attr: str, + ) -> List[SeedNode]: + """Convert RetrievalResult seeds to SeedNode format.""" + if not seeds: + return [] + + result = [] + for seed in seeds: + score = getattr(seed, score_attr, None) + if score is None: + score = seed.activation or seed.similarity or 1.0 + result.append(SeedNode(node_id=seed.id, score=score)) + + return result + + async def _find_semantic_seeds( + self, + pool, + query_embedding_str: str, + bank_id: str, + fact_type: str, + limit: int = 20, + threshold: float = 0.3, + ) -> List[SeedNode]: + """Fallback: find semantic seeds via embedding search.""" + async with acquire_with_retry(pool) as conn: + rows = await conn.fetch( + """ + SELECT id, 1 - (embedding <=> $1::vector) AS similarity + FROM memory_units + WHERE bank_id = $2 + AND embedding IS NOT NULL + AND fact_type = $3 + AND (1 - (embedding <=> $1::vector)) >= $4 + ORDER BY embedding <=> $1::vector + LIMIT $5 + """, + query_embedding_str, bank_id, fact_type, threshold, limit + ) + + return [ + SeedNode(node_id=str(r['id']), score=r['similarity']) + for r in rows + ] diff --git a/hindsight-api/hindsight_api/engine/search/retrieval.py b/hindsight-api/hindsight_api/engine/search/retrieval.py index 3fa8e142..2c36da9f 100644 --- a/hindsight-api/hindsight_api/engine/search/retrieval.py +++ b/hindsight-api/hindsight_api/engine/search/retrieval.py @@ -8,22 +8,50 @@ Implements: 4. Temporal retrieval (time-aware search with spreading) """ -from typing import List, Dict, Any, Tuple, Optional +from typing import List, Dict, Optional +from dataclasses import dataclass, field from datetime import datetime import asyncio +import logging from ..db_utils import acquire_with_retry from .types import RetrievalResult from .graph_retrieval import GraphRetriever, BFSGraphRetriever +from .mpfp_retrieval import MPFPGraphRetriever +from ...config import get_config + +logger = logging.getLogger(__name__) + + +@dataclass +class ParallelRetrievalResult: + """Result from parallel retrieval across all methods.""" + semantic: List[RetrievalResult] + bm25: List[RetrievalResult] + graph: List[RetrievalResult] + temporal: Optional[List[RetrievalResult]] + timings: Dict[str, float] = field(default_factory=dict) + temporal_constraint: Optional[tuple] = None # (start_date, end_date) + # Default graph retriever instance (can be overridden) _default_graph_retriever: Optional[GraphRetriever] = None def get_default_graph_retriever() -> GraphRetriever: - """Get or create the default graph retriever.""" + """Get or create the default graph retriever based on config.""" global _default_graph_retriever if _default_graph_retriever is None: - _default_graph_retriever = BFSGraphRetriever() + config = get_config() + retriever_type = config.graph_retriever.lower() + if retriever_type == "mpfp": + _default_graph_retriever = MPFPGraphRetriever() + logger.info("Using MPFP graph retriever") + elif retriever_type == "bfs": + _default_graph_retriever = BFSGraphRetriever() + logger.info("Using BFS graph retriever") + else: + logger.warning(f"Unknown graph retriever '{retriever_type}', falling back to MPFP") + _default_graph_retriever = MPFPGraphRetriever() return _default_graph_retriever @@ -324,7 +352,7 @@ async def retrieve_parallel( question_date: Optional[datetime] = None, query_analyzer: Optional["QueryAnalyzer"] = None, graph_retriever: Optional[GraphRetriever] = None, -) -> Tuple[List[RetrievalResult], List[RetrievalResult], List[RetrievalResult], Optional[List[RetrievalResult]], Dict[str, float], Optional[Tuple[datetime, datetime]]]: +) -> ParallelRetrievalResult: """ Run 3-way or 4-way parallel retrieval (adds temporal if detected). @@ -337,43 +365,259 @@ async def retrieve_parallel( 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 BFSGraphRetriever) + graph_retriever: Graph retrieval strategy (defaults to configured retriever) Returns: - Tuple of (semantic_results, bm25_results, graph_results, temporal_results, timings, temporal_constraint) - Each results list contains RetrievalResult objects - temporal_results is None if no temporal constraint detected - timings is a dict with per-method latencies in seconds - temporal_constraint is the (start_date, end_date) tuple if detected, else None + ParallelRetrievalResult with semantic, bm25, graph, temporal results and timings """ - # Detect temporal constraint from .temporal_extraction import extract_temporal_constraint - import time temporal_constraint = extract_temporal_constraint( query_text, reference_date=question_date, analyzer=query_analyzer ) - # Use provided graph retriever or default retriever = graph_retriever or get_default_graph_retriever() - # Wrapper to track timing for each retrieval method - async def timed_retrieval(name: str, coro): + if retriever.name == "mpfp": + return await _retrieve_parallel_mpfp( + pool, query_text, query_embedding_str, bank_id, fact_type, + thinking_budget, temporal_constraint, retriever + ) + else: + return await _retrieve_parallel_bfs( + pool, query_text, query_embedding_str, bank_id, fact_type, + thinking_budget, temporal_constraint, retriever + ) + + +@dataclass +class _SemanticGraphResult: + """Internal result from semantic→graph chain.""" + semantic: List[RetrievalResult] + graph: List[RetrievalResult] + semantic_time: float + graph_time: float + + +@dataclass +class _TimedResult: + """Internal result with timing.""" + results: List[RetrievalResult] + time: float + + +async def _retrieve_parallel_mpfp( + pool, + query_text: str, + query_embedding_str: str, + bank_id: str, + fact_type: str, + thinking_budget: int, + temporal_constraint: Optional[tuple], + retriever: GraphRetriever, +) -> ParallelRetrievalResult: + """ + MPFP retrieval with optimized parallelization. + + Runs 2-3 parallel task chains: + - Task 1: Semantic → Graph (chained, graph uses semantic seeds) + - Task 2: BM25 (independent) + - Task 3: Temporal (if constraint detected) + """ + import time + + async def run_semantic_then_graph() -> _SemanticGraphResult: + """Chain: semantic retrieval → graph retrieval (using semantic as seeds).""" start = time.time() - result = await coro - duration = time.time() - start - return result, name, duration - - async def run_semantic(): async with acquire_with_retry(pool) as conn: - return await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget) + semantic = await retrieve_semantic( + conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget + ) + semantic_time = time.time() - start - async def run_bm25(): + # Get temporal seeds if needed (quick query, part of this chain) + temporal_seeds = None + if temporal_constraint: + tc_start, tc_end = temporal_constraint + async with acquire_with_retry(pool) as conn: + temporal_seeds = await _get_temporal_entry_points( + conn, query_embedding_str, bank_id, fact_type, + tc_start, tc_end, limit=20 + ) + + # Run graph with seeds + start = time.time() + graph = await retriever.retrieve( + pool=pool, + query_embedding_str=query_embedding_str, + bank_id=bank_id, + fact_type=fact_type, + budget=thinking_budget, + query_text=query_text, + semantic_seeds=semantic, + temporal_seeds=temporal_seeds, + ) + graph_time = time.time() - start + + return _SemanticGraphResult(semantic, graph, semantic_time, graph_time) + + async def run_bm25() -> _TimedResult: + """Independent BM25 retrieval.""" + start = time.time() async with acquire_with_retry(pool) as conn: - return await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget) + results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget) + return _TimedResult(results, time.time() - start) - async def run_graph(): - return await retriever.retrieve( + async def run_temporal(tc_start, tc_end) -> _TimedResult: + """Temporal retrieval (uses its own entry point finding).""" + start = time.time() + async with acquire_with_retry(pool) as conn: + results = await retrieve_temporal( + conn, query_embedding_str, bank_id, fact_type, + tc_start, tc_end, budget=thinking_budget, semantic_threshold=0.1 + ) + return _TimedResult(results, time.time() - start) + + # Run parallel task chains + if temporal_constraint: + tc_start, tc_end = temporal_constraint + sg_result, bm25_result, temporal_result = await asyncio.gather( + run_semantic_then_graph(), + run_bm25(), + run_temporal(tc_start, tc_end), + ) + return ParallelRetrievalResult( + semantic=sg_result.semantic, + bm25=bm25_result.results, + graph=sg_result.graph, + temporal=temporal_result.results, + timings={ + "semantic": sg_result.semantic_time, + "graph": sg_result.graph_time, + "bm25": bm25_result.time, + "temporal": temporal_result.time, + }, + temporal_constraint=temporal_constraint, + ) + else: + sg_result, bm25_result = await asyncio.gather( + run_semantic_then_graph(), + run_bm25(), + ) + return ParallelRetrievalResult( + semantic=sg_result.semantic, + bm25=bm25_result.results, + graph=sg_result.graph, + temporal=None, + timings={ + "semantic": sg_result.semantic_time, + "graph": sg_result.graph_time, + "bm25": bm25_result.time, + }, + temporal_constraint=None, + ) + + +async def _get_temporal_entry_points( + conn, + query_embedding_str: str, + bank_id: str, + fact_type: str, + start_date: datetime, + end_date: datetime, + limit: int = 20, + semantic_threshold: float = 0.1, +) -> List[RetrievalResult]: + """Get temporal entry points (facts in date range with semantic relevance).""" + from datetime import timezone + + if start_date.tzinfo is None: + start_date = start_date.replace(tzinfo=timezone.utc) + if end_date.tzinfo is None: + end_date = end_date.replace(tzinfo=timezone.utc) + + rows = await conn.fetch( + """ + 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 + FROM memory_units + WHERE bank_id = $2 + AND fact_type = $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 + ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, + (embedding <=> $1::vector) ASC + LIMIT $7 + """, + query_embedding_str, bank_id, fact_type, start_date, end_date, semantic_threshold, limit + ) + + results = [] + total_days = max((end_date - start_date).total_seconds() / 86400, 1) + mid_date = start_date + (end_date - start_date) / 2 + + for row in rows: + result = RetrievalResult.from_db_row(dict(row)) + + # Calculate temporal proximity score + best_date = None + if row["occurred_start"] and row["occurred_end"]: + best_date = row["occurred_start"] + (row["occurred_end"] - row["occurred_start"]) / 2 + elif row["occurred_start"]: + best_date = row["occurred_start"] + elif row["occurred_end"]: + best_date = row["occurred_end"] + elif row["mentioned_at"]: + best_date = row["mentioned_at"] + + if best_date: + days_from_mid = abs((best_date - mid_date).total_seconds() / 86400) + result.temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) + else: + result.temporal_proximity = 0.5 + + result.temporal_score = result.temporal_proximity + results.append(result) + + return results + + +async def _retrieve_parallel_bfs( + pool, + query_text: str, + query_embedding_str: str, + bank_id: str, + fact_type: str, + thinking_budget: int, + temporal_constraint: Optional[tuple], + retriever: GraphRetriever, +) -> ParallelRetrievalResult: + """BFS retrieval: all methods run in parallel (original behavior).""" + import time + + async def run_semantic() -> _TimedResult: + start = time.time() + async with acquire_with_retry(pool) as conn: + results = await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget) + return _TimedResult(results, time.time() - start) + + async def run_bm25() -> _TimedResult: + start = time.time() + async with acquire_with_retry(pool) as conn: + results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget) + return _TimedResult(results, time.time() - start) + + async def run_graph() -> _TimedResult: + start = time.time() + results = await retriever.retrieve( pool=pool, query_embedding_str=query_embedding_str, bank_id=bank_id, @@ -381,37 +625,53 @@ async def retrieve_parallel( budget=thinking_budget, query_text=query_text, ) + return _TimedResult(results, time.time() - start) - async def run_temporal(start_date, end_date): + async def run_temporal(tc_start, tc_end) -> _TimedResult: + start = time.time() async with acquire_with_retry(pool) as conn: - return await retrieve_temporal( + results = await retrieve_temporal( conn, query_embedding_str, bank_id, fact_type, - start_date, end_date, budget=thinking_budget, semantic_threshold=0.1 + tc_start, tc_end, budget=thinking_budget, semantic_threshold=0.1 ) + return _TimedResult(results, time.time() - start) - # Run retrievals in parallel with timing - timings = {} if temporal_constraint: - start_date, end_date = temporal_constraint - results = await asyncio.gather( - timed_retrieval("semantic", run_semantic()), - timed_retrieval("bm25", run_bm25()), - timed_retrieval("graph", run_graph()), - timed_retrieval("temporal", run_temporal(start_date, end_date)) + tc_start, tc_end = temporal_constraint + semantic_r, bm25_r, graph_r, temporal_r = await asyncio.gather( + run_semantic(), + run_bm25(), + run_graph(), + run_temporal(tc_start, tc_end), + ) + return ParallelRetrievalResult( + semantic=semantic_r.results, + bm25=bm25_r.results, + graph=graph_r.results, + temporal=temporal_r.results, + timings={ + "semantic": semantic_r.time, + "bm25": bm25_r.time, + "graph": graph_r.time, + "temporal": temporal_r.time, + }, + temporal_constraint=temporal_constraint, ) - semantic_results, _, timings["semantic"] = results[0] - bm25_results, _, timings["bm25"] = results[1] - graph_results, _, timings["graph"] = results[2] - temporal_results, _, timings["temporal"] = results[3] else: - results = await asyncio.gather( - timed_retrieval("semantic", run_semantic()), - timed_retrieval("bm25", run_bm25()), - timed_retrieval("graph", run_graph()) + semantic_r, bm25_r, graph_r = await asyncio.gather( + run_semantic(), + run_bm25(), + run_graph(), + ) + return ParallelRetrievalResult( + semantic=semantic_r.results, + bm25=bm25_r.results, + graph=graph_r.results, + temporal=None, + timings={ + "semantic": semantic_r.time, + "bm25": bm25_r.time, + "graph": graph_r.time, + }, + temporal_constraint=None, ) - semantic_results, _, timings["semantic"] = results[0] - bm25_results, _, timings["bm25"] = results[1] - graph_results, _, timings["graph"] = results[2] - temporal_results = None - - return semantic_results, bm25_results, graph_results, temporal_results, timings, temporal_constraint