hindsight-api 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

hindsight_api/engine/search/retrieval.py

@@ -4,24 +4,68 @@ Retrieval module for 4-way parallel search.
 Implements:
 1. Semantic retrieval (vector similarity)
 2. BM25 retrieval (keyword/full-text search)
-3. Graph retrieval (spreading activation)
+3. Graph retrieval (via pluggable GraphRetriever interface)
 4. Temporal retrieval (time-aware search with spreading)
 """
 
-from typing import List, Dict, Any, Tuple, Optional
-from datetime import datetime
 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 .graph_retrieval import BFSGraphRetriever, GraphRetriever
+from .mpfp_retrieval import MPFPGraphRetriever
 from .types import 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)
+
+
+# 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("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
+
+
+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(
-    conn,
-    query_emb_str: str,
-    bank_id: str,
-    fact_type: str,
-    limit: int
-) -> List[RetrievalResult]:
+    conn, query_emb_str: str, bank_id: str, fact_type: str, limit: int
+) -> list[RetrievalResult]:
     """
     Semantic retrieval via vector similarity.
 
@@ -47,18 +91,15 @@ async def retrieve_semantic(
         ORDER BY embedding <=> $1::vector
         LIMIT $4
         """,
-        query_emb_str, bank_id, fact_type, limit
+        query_emb_str,
+        bank_id,
+        fact_type,
+        limit,
     )
     return [RetrievalResult.from_db_row(dict(r)) for r in results]
 
 
-async def retrieve_bm25(
-    conn,
-    query_text: str,
-    bank_id: str,
-    fact_type: str,
-    limit: int
-) -> List[RetrievalResult]:
+async def retrieve_bm25(conn, query_text: str, bank_id: str, fact_type: str, limit: int) -> list[RetrievalResult]:
     """
     BM25 keyword retrieval via full-text search.
 
@@ -76,7 +117,7 @@ async def retrieve_bm25(
 
     # Sanitize query text: remove special characters that have meaning in tsquery
     # Keep only alphanumeric characters and spaces
-    sanitized_text = re.sub(r'[^\w\s]', ' ', query_text.lower())
+    sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
 
     # Split and filter empty strings
     tokens = [token for token in sanitized_text.split() if token]
@@ -100,126 +141,14 @@ async def retrieve_bm25(
         ORDER BY bm25_score DESC
         LIMIT $4
         """,
-        query_tsquery, bank_id, fact_type, limit
+        query_tsquery,
+        bank_id,
+        fact_type,
+        limit,
     )
     return [RetrievalResult.from_db_row(dict(r)) for r in results]
 
 
-async def retrieve_graph(
-    conn,
-    query_emb_str: str,
-    bank_id: str,
-    fact_type: str,
-    budget: int
-) -> List[RetrievalResult]:
-    """
-    Graph retrieval via spreading activation.
-
-    Args:
-        conn: Database connection
-        query_emb_str: Query embedding as string
-        agent_id: bank ID
-        fact_type: Fact type to filter
-        budget: Node budget for graph traversal
-
-    Returns:
-        List of RetrievalResult objects
-    """
-    # Find entry points
-    entry_points = 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 embedding IS NOT NULL
-          AND fact_type = $3
-          AND (1 - (embedding <=> $1::vector)) >= 0.5
-        ORDER BY embedding <=> $1::vector
-        LIMIT 5
-        """,
-        query_emb_str, bank_id, fact_type
-    )
-
-    if not entry_points:
-        return []
-
-    # BFS-style spreading activation with batched neighbor fetching
-    visited = set()
-    results = []
-    queue = [(RetrievalResult.from_db_row(dict(r)), r["similarity"]) for r in entry_points]
-    budget_remaining = budget
-
-    # Process nodes in batches to reduce DB roundtrips
-    batch_size = 20  # Fetch neighbors for up to 20 nodes at once
-
-    while queue and budget_remaining > 0:
-        # Collect a batch of nodes to process
-        batch_nodes = []
-        batch_activations = {}
-
-        while queue and len(batch_nodes) < batch_size and budget_remaining > 0:
-            current, activation = queue.pop(0)
-            unit_id = current.id
-
-            if unit_id not in visited:
-                visited.add(unit_id)
-                budget_remaining -= 1
-                results.append(current)
-                batch_nodes.append(current.id)
-                batch_activations[unit_id] = activation
-
-        # Batch fetch neighbors for all nodes in this batch
-        # Fetch top weighted neighbors (batch_size * 20 = ~400 for good distribution)
-        if batch_nodes and budget_remaining > 0:
-            max_neighbors = len(batch_nodes) * 20
-            neighbors = await conn.fetch(
-                """
-                SELECT mu.id, mu.text, mu.context, 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
-                FROM memory_links ml
-                JOIN memory_units mu ON ml.to_unit_id = mu.id
-                WHERE ml.from_unit_id = ANY($1::uuid[])
-                  AND ml.weight >= 0.1
-                  AND mu.fact_type = $2
-                ORDER BY ml.weight DESC
-                LIMIT $3
-                """,
-                batch_nodes, fact_type, max_neighbors
-            )
-
-            for n in neighbors:
-                neighbor_id = str(n["id"])
-                if neighbor_id not in visited:
-                    # Get parent activation
-                    parent_id = str(n["from_unit_id"])
-                    activation = batch_activations.get(parent_id, 0.5)
-
-                    # Boost activation for causal links (they're high-value relationships)
-                    link_type = n["link_type"]
-                    base_weight = n["weight"]
-
-                    # Causal links get 1.5-2.0x boost depending on type
-                    if link_type in ("causes", "caused_by"):
-                        # Direct causation - very strong relationship
-                        causal_boost = 2.0
-                    elif link_type in ("enables", "prevents"):
-                        # Conditional causation - strong but not as direct
-                        causal_boost = 1.5
-                    else:
-                        # Temporal, semantic, entity links - standard weight
-                        causal_boost = 1.0
-
-                    effective_weight = base_weight * causal_boost
-                    new_activation = activation * effective_weight * 0.8
-                    if new_activation > 0.1:
-                        neighbor_result = RetrievalResult.from_db_row(dict(n))
-                        queue.append((neighbor_result, new_activation))
-
-    return results
-
-
 async def retrieve_temporal(
     conn,
     query_emb_str: str,
@@ -228,8 +157,8 @@ async def retrieve_temporal(
     start_date: datetime,
     end_date: datetime,
     budget: int,
-    semantic_threshold: float = 0.1
-) -> List[RetrievalResult]:
+    semantic_threshold: float = 0.1,
+) -> list[RetrievalResult]:
     """
     Temporal retrieval with spreading activation.
 
@@ -251,13 +180,12 @@ async def retrieve_temporal(
     Returns:
         List of RetrievalResult objects with temporal scores
     """
-    from datetime import timezone
 
     # Ensure start_date and end_date are timezone-aware (UTC) to match database datetimes
     if start_date.tzinfo is None:
-        start_date = start_date.replace(tzinfo=timezone.utc)
+        start_date = start_date.replace(tzinfo=UTC)
     if end_date.tzinfo is None:
-        end_date = end_date.replace(tzinfo=timezone.utc)
+        end_date = end_date.replace(tzinfo=UTC)
 
     entry_points = await conn.fetch(
         """
@@ -284,7 +212,12 @@ async def retrieve_temporal(
         ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, (embedding <=> $1::vector) ASC
         LIMIT 10
         """,
-        query_emb_str, bank_id, fact_type, start_date, end_date, semantic_threshold
+        query_emb_str,
+        bank_id,
+        fact_type,
+        start_date,
+        end_date,
+        semantic_threshold,
     )
 
     if not entry_points:
@@ -327,7 +260,9 @@ async def retrieve_temporal(
         results.append(ep_result)
 
     # Spread through temporal links
-    queue = [(RetrievalResult.from_db_row(dict(ep)), ep["similarity"], 1.0) for ep in entry_points]  # (unit, semantic_sim, temporal_score)
+    queue = [
+        (RetrievalResult.from_db_row(dict(ep)), ep["similarity"], 1.0) for ep in entry_points
+    ]  # (unit, semantic_sim, temporal_score)
     budget_remaining = budget - len(entry_points)
 
     while queue and budget_remaining > 0:
@@ -352,7 +287,10 @@ async def retrieve_temporal(
                 ORDER BY ml.weight DESC
                 LIMIT 10
                 """,
-                query_emb_str, current.id, fact_type, semantic_threshold
+                query_emb_str,
+                current.id,
+                fact_type,
+                semantic_threshold,
             )
 
             for n in neighbors:
@@ -376,7 +314,9 @@ async def retrieve_temporal(
 
                 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
+                    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  # Lower score if no temporal data
 
@@ -418,9 +358,10 @@ async def retrieve_parallel(
     bank_id: str,
     fact_type: str,
     thinking_budget: int,
-    question_date: Optional[datetime] = None,
-    query_analyzer: Optional["QueryAnalyzer"] = None
-) -> Tuple[List[RetrievalResult], List[RetrievalResult], List[RetrievalResult], Optional[List[RetrievalResult]], Dict[str, float], Optional[Tuple[datetime, datetime]]]:
+    question_date: datetime | None = None,
+    query_analyzer: Optional["QueryAnalyzer"] = None,
+    graph_retriever: GraphRetriever | None = None,
+) -> ParallelRetrievalResult:
     """
     Run 3-way or 4-way parallel retrieval (adds temporal if detected).
 
@@ -428,76 +369,330 @@ async def retrieve_parallel(
         pool: Database connection pool
         query_text: Query text
         query_embedding_str: Query embedding as string
-        agent_id: bank ID
+        bank_id: Bank ID
         fact_type: Fact type to filter
         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:
-        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
+
+    temporal_constraint = extract_temporal_constraint(query_text, reference_date=question_date, analyzer=query_analyzer)
+
+    retriever = graph_retriever or get_default_graph_retriever()
+
+    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: tuple | None,
+    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
 
-    temporal_constraint = extract_temporal_constraint(
-        query_text, reference_date=question_date, analyzer=query_analyzer
-    )
+    async def run_semantic_then_graph() -> _SemanticGraphResult:
+        """Chain: semantic retrieval → graph retrieval (using semantic as seeds)."""
+        start = time.time()
+        async with acquire_with_retry(pool) as conn:
+            semantic = await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget)
+        semantic_time = time.time() - start
+
+        # 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)
 
-    # Wrapper to track timing for each retrieval method
-    async def timed_retrieval(name: str, coro):
+    async def run_bm25() -> _TimedResult:
+        """Independent BM25 retrieval."""
         start = time.time()
-        result = await coro
-        duration = time.time() - start
-        return result, name, duration
+        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_semantic():
+    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:
-            return await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget)
+            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)."""
+
+    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)
+
+    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
 
-    async def run_bm25():
+        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: tuple | None,
+    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:
-            return await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget)
+            results = await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget)
+        return _TimedResult(results, time.time() - start)
 
-    async def run_graph():
+    async def run_bm25() -> _TimedResult:
+        start = time.time()
         async with acquire_with_retry(pool) as conn:
-            return await retrieve_graph(conn, query_embedding_str, bank_id, fact_type, budget=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() -> _TimedResult:
+        start = time.time()
+        results = 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,
+        )
+        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(
-                conn, query_embedding_str, bank_id, fact_type,
-                start_date, end_date, budget=thinking_budget, semantic_threshold=0.1
+            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 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