""" Retrieval module for 4-way parallel search. Implements: 1. Semantic retrieval (vector similarity) 2. BM25 retrieval (keyword/full-text search) 3. Graph retrieval (via pluggable GraphRetriever interface) 4. Temporal retrieval (time-aware search with spreading) """ import asyncio import logging from dataclasses import dataclass, field from datetime import UTC, datetime from typing import Optional from ...config import get_config from ..db_utils import acquire_with_retry from ..memory_engine import fq_table from .graph_retrieval import BFSGraphRetriever, GraphRetriever from .link_expansion_retrieval import LinkExpansionRetriever from .mpfp_retrieval import MPFPGraphRetriever from .tags import TagsMatch, build_tags_where_clause_simple from .types import MPFPTimings, RetrievalResult logger = logging.getLogger(__name__) @dataclass class ParallelRetrievalResult: """Result from parallel retrieval across all methods.""" semantic: list[RetrievalResult] bm25: list[RetrievalResult] graph: list[RetrievalResult] temporal: list[RetrievalResult] | None timings: dict[str, float] = field(default_factory=dict) temporal_constraint: tuple | None = None # (start_date, end_date) mpfp_timings: list[MPFPTimings] = field(default_factory=list) # MPFP sub-step timings per fact type 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 def get_default_graph_retriever() -> GraphRetriever: """Get or create the default graph retriever based on config.""" global _default_graph_retriever if _default_graph_retriever is None: config = get_config() retriever_type = config.graph_retriever.lower() if retriever_type == "mpfp": _default_graph_retriever = MPFPGraphRetriever() logger.info( f"Using MPFP graph retriever (top_k_neighbors={_default_graph_retriever.config.top_k_neighbors})" ) elif retriever_type == "bfs": _default_graph_retriever = BFSGraphRetriever() logger.info("Using BFS graph retriever") elif retriever_type == "link_expansion": _default_graph_retriever = LinkExpansionRetriever() logger.info("Using LinkExpansion graph retriever") else: logger.warning(f"Unknown graph retriever '{retriever_type}', falling back to link_expansion") _default_graph_retriever = LinkExpansionRetriever() return _default_graph_retriever def set_default_graph_retriever(retriever: GraphRetriever) -> None: """Set the default graph retriever (for configuration/testing).""" global _default_graph_retriever _default_graph_retriever = retriever async def retrieve_semantic_bm25_combined( conn, query_emb_str: str, query_text: str, bank_id: str, fact_types: list[str], limit: int, tags: list[str] | None = None, tags_match: TagsMatch = "any", ) -> 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: tags_clause = build_tags_where_clause_simple(tags, 5, match=tags_match) params = [query_emb_str, bank_id, fact_types, limit] if tags: params.append(tags) results = await conn.fetch( f""" WITH semantic_ranked AS ( SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, 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 {tags_clause} ) SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity, bm25_score, source FROM semantic_ranked WHERE rn <= $4 """, *params, ) # 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 # Build BM25 query based on text search backend config = get_config() # Build tags clause - param 6 if tags provided tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match) # Build backend-specific BM25 parts if config.text_search_extension == "vchord": # VectorChord BM25: use <&> operator with to_bm25query and tokenize # Note: VectorChord scores are negative (higher = better, so -1 > -10) bm25_score_expr = "search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($5, 'llmlingua2'))" bm25_order_by = f"{bm25_score_expr} DESC" bm25_where_filter = "" # No additional WHERE filter for vchord params = [query_emb_str, bank_id, fact_types, limit, query_text] # Pass raw query_text for tokenization elif config.text_search_extension == "pg_textsearch": # Timescale pg_textsearch: use <@> operator with to_bm25query # Note: pg_textsearch scores are negative (lower/more negative = better, so -10 > -1) # We negate the score to maintain API consistency (higher = better) bm25_score_expr = "-(text <@> to_bm25query($5, 'idx_memory_units_text_search'))" bm25_order_by = "text <@> to_bm25query($5, 'idx_memory_units_text_search') ASC" bm25_where_filter = "" # No additional WHERE filter for pg_textsearch params = [query_emb_str, bank_id, fact_types, limit, query_text] else: # native # Native PostgreSQL: use ts_rank_cd with to_tsquery query_tsquery = " | ".join(tokens) bm25_score_expr = "ts_rank_cd(search_vector, to_tsquery('english', $5))" bm25_order_by = f"{bm25_score_expr} DESC" bm25_where_filter = "AND search_vector @@ to_tsquery('english', $5)" params = [query_emb_str, bank_id, fact_types, limit, query_tsquery] if tags: params.append(tags) # Single query template with backend-specific parts injected query = f""" WITH semantic_ranked AS ( SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, 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 {tags_clause} ), bm25_ranked AS ( SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, NULL::float AS similarity, {bm25_score_expr} AS bm25_score, 'bm25' AS source, ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY {bm25_order_by}) AS rn FROM {fq_table("memory_units")} WHERE bank_id = $2 AND fact_type = ANY($3) {bm25_where_filter} {tags_clause} ), semantic AS ( SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity, bm25_score, source FROM semantic_ranked WHERE rn <= $4 ), bm25 AS ( SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity, bm25_score, source FROM bm25_ranked WHERE rn <= $4 ) SELECT * FROM semantic UNION ALL SELECT * FROM bm25 """ # 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(query, *params) # 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, tags: list[str] | None = None, tags_match: TagsMatch = "any", ) -> 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) # Build tags clause tags_clause = build_tags_where_clause_simple(tags, 7, match=tags_match) params = [query_emb_str, bank_id, fact_types, start_date, end_date, semantic_threshold] if tags: params.append(tags) # Two-phase entry point query: # Phase 1 (date_ranked): rank by date only — no embedding computation — for all units in # the temporal window. This lets the planner use date indexes for filtering. # Phase 2 (sim_ranked): join back to memory_units for only the top-50-per-type candidates # and compute embedding similarity for that small set (≤ 50 × len(fact_types) rows). # This avoids computing embedding distances for potentially thousands of date-range rows. entry_points = await conn.fetch( f""" WITH date_ranked AS MATERIALIZED ( SELECT id, fact_type, ROW_NUMBER() OVER ( PARTITION BY fact_type ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC NULLS LAST ) 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) ) {tags_clause} ), sim_ranked AS ( SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, 1 - (mu.embedding <=> $1::vector) AS similarity, ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn FROM date_ranked dr JOIN {fq_table("memory_units")} mu ON mu.id = dr.id WHERE dr.rn <= 50 AND (1 - (mu.embedding <=> $1::vector)) >= $6 ) SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity FROM sim_ranked WHERE sim_rn <= 10 """, *params, ) 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 # Per-source neighbor limit: lets the planner use the composite index # (from_unit_id, link_type, weight DESC) with early termination, avoiding # a full scan of all links from all source nodes before sorting. per_source_limit = 10 # Safety cap on BFS iterations to prevent runaway spreading in dense graphs. max_iterations = 5 iteration = 0 # Build tags clause for spreading (use param 7 since 1-6 are used) spreading_tags_clause = build_tags_where_clause_simple(tags, 7, table_alias="mu.", match=tags_match) while frontier and budget_remaining > 0 and iteration < max_iterations: iteration += 1 batch_ids = frontier[:batch_size] frontier = frontier[batch_size:] # $1=query_emb, $2=batch_ids, $3=fact_type, $4=threshold, $5=per_source_limit, $6=bank_id, $7=tags spreading_params = [query_emb_str, batch_ids, ft, semantic_threshold, per_source_limit, bank_id] if tags: spreading_params.append(tags) # LATERAL join: for each source node, fetch top-K neighbors by weight using # the existing idx_memory_links_from_type_weight index with early-exit semantics. # This avoids scanning all temporal links from all source nodes before sorting. # bank_id on memory_units lets the planner use idx_memory_units_bank_fact_type. neighbors = await conn.fetch( f""" SELECT src.from_unit_id, mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, l.weight, l.link_type, 1 - (mu.embedding <=> $1::vector) AS similarity FROM unnest($2::uuid[]) AS src(from_unit_id) CROSS JOIN LATERAL ( SELECT ml.to_unit_id, ml.weight, ml.link_type FROM {fq_table("memory_links")} ml WHERE ml.from_unit_id = src.from_unit_id AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents') AND ml.weight >= 0.1 ORDER BY ml.weight DESC LIMIT $5 ) l JOIN {fq_table("memory_units")} mu ON mu.id = l.to_unit_id WHERE mu.bank_id = $6 AND mu.fact_type = $3 AND mu.embedding IS NOT NULL AND (1 - (mu.embedding <=> $1::vector)) >= $4 {spreading_tags_clause} """, *spreading_params, ) 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_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, tags: list[str] | None = None, tags_match: TagsMatch = "any", ) -> 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, tags=tags, tags_match=tags_match, ) 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, tags=tags, tags_match=tags_match, ) 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, tags=tags, tags_match=tags_match, ) 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, )