hindsight-api 0.0.13__py3-none-any.whl

hindsight_api/engine/search/retrieval.py

@@ -0,0 +1,503 @@
+"""
+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)
+4. Temporal retrieval (time-aware search with spreading)
+"""
+
+from typing import List, Dict, Any, Tuple, Optional
+from datetime import datetime
+import asyncio
+from ..db_utils import acquire_with_retry
+from .types import RetrievalResult
+
+
+async def retrieve_semantic(
+    conn,
+    query_emb_str: str,
+    bank_id: str,
+    fact_type: str,
+    limit: int
+) -> List[RetrievalResult]:
+    """
+    Semantic retrieval via vector similarity.
+
+    Args:
+        conn: Database connection
+        query_emb_str: Query embedding as string
+        agent_id: bank ID
+        fact_type: Fact type to filter
+        limit: Maximum results to return
+
+    Returns:
+        List of RetrievalResult objects
+    """
+    results = 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.3
+        ORDER BY embedding <=> $1::vector
+        LIMIT $4
+        """,
+        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]:
+    """
+    BM25 keyword retrieval via full-text search.
+
+    Args:
+        conn: Database connection
+        query_text: Query text
+        agent_id: bank ID
+        fact_type: Fact type to filter
+        limit: Maximum results to return
+
+    Returns:
+        List of RetrievalResult objects
+    """
+    import re
+
+    # 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())
+
+    # Split and filter empty strings
+    tokens = [token for token in sanitized_text.split() if token]
+
+    if not tokens:
+        # If no valid tokens, return empty results
+        return []
+
+    # Convert query to tsquery using OR for more flexible matching
+    # This prevents empty results when some terms are missing
+    query_tsquery = " | ".join(tokens)
+
+    results = await conn.fetch(
+        """
+        SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+               ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
+        FROM memory_units
+        WHERE bank_id = $2
+          AND fact_type = $3
+          AND search_vector @@ to_tsquery('english', $1)
+        ORDER BY bm25_score DESC
+        LIMIT $4
+        """,
+        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 * 10 = ~200 for good distribution)
+        if batch_nodes and budget_remaining > 0:
+            max_neighbors = len(batch_nodes) * 10
+            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,
+    bank_id: str,
+    fact_type: str,
+    start_date: datetime,
+    end_date: datetime,
+    budget: int,
+    semantic_threshold: float = 0.1
+) -> List[RetrievalResult]:
+    """
+    Temporal retrieval with spreading activation.
+
+    Strategy:
+    1. Find entry points (facts in date range with semantic relevance)
+    2. Spread through temporal links to related facts
+    3. Score by temporal proximity + semantic similarity + link weight
+
+    Args:
+        conn: Database connection
+        query_emb_str: Query embedding as string
+        agent_id: bank ID
+        fact_type: Fact type to filter
+        start_date: Start of time range
+        end_date: End of time range
+        budget: Node budget for spreading
+        semantic_threshold: Minimum semantic similarity to include
+
+    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)
+    if end_date.tzinfo is None:
+        end_date = end_date.replace(tzinfo=timezone.utc)
+
+    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 fact_type = $3
+          AND embedding IS NOT NULL
+          AND (
+              -- Match if occurred range overlaps with query range
+              (occurred_start IS NOT NULL AND occurred_end IS NOT NULL
+               AND occurred_start <= $5 AND occurred_end >= $4)
+              OR
+              -- Match if mentioned_at falls within query range
+              (mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
+              OR
+              -- Match if any occurred date is set and overlaps (even if only start or end is set)
+              (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 10
+        """,
+        query_emb_str, bank_id, fact_type, start_date, end_date, semantic_threshold
+    )
+
+    if not entry_points:
+        return []
+
+    # Calculate temporal scores for entry points
+    total_days = (end_date - start_date).total_seconds() / 86400
+    mid_date = start_date + (end_date - start_date) / 2  # Calculate once for all comparisons
+    results = []
+    visited = set()
+
+    for ep in entry_points:
+        unit_id = str(ep["id"])
+        visited.add(unit_id)
+
+        # Calculate temporal proximity using the most relevant date
+        # Priority: occurred_start/end (event time) > mentioned_at (mention time)
+        best_date = None
+        if ep["occurred_start"] is not None and ep["occurred_end"] is not None:
+            # Use midpoint of occurred range
+            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"]
+
+        # Temporal proximity score (closer to range center = higher score)
+        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  # Fallback if no dates (shouldn't happen due to WHERE clause)
+
+        # Create RetrievalResult with temporal scores
+        ep_result = RetrievalResult.from_db_row(dict(ep))
+        ep_result.temporal_score = temporal_proximity
+        ep_result.temporal_proximity = temporal_proximity
+        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)
+    budget_remaining = budget - len(entry_points)
+
+    while queue and budget_remaining > 0:
+        current, semantic_sim, temporal_score = queue.pop(0)
+        current_id = current.id
+
+        # Get neighbors via temporal and causal links
+        if budget_remaining > 0:
+            neighbors = await conn.fetch(
+                """
+                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,
+                       1 - (mu.embedding <=> $1::vector) AS similarity
+                FROM memory_links ml
+                JOIN memory_units mu ON ml.to_unit_id = mu.id
+                WHERE ml.from_unit_id = $2
+                  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 10
+                """,
+                query_emb_str, current.id, fact_type, semantic_threshold
+            )
+
+            for n in neighbors:
+                neighbor_id = str(n["id"])
+                if neighbor_id in visited:
+                    continue
+
+                visited.add(neighbor_id)
+                budget_remaining -= 1
+
+                # Calculate temporal score for neighbor using best available date
+                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  # Lower score if no temporal data
+
+                # Boost causal links (same as graph retrieval)
+                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
+
+                # Propagate temporal score through links (decay, with causal boost)
+                propagated_temporal = temporal_score * n["weight"] * causal_boost * 0.7
+
+                # Combined temporal score
+                combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)
+
+                # Create RetrievalResult with temporal scores
+                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)
+
+                # Add to queue for further spreading
+                if budget_remaining > 0 and combined_temporal > 0.2:
+                    queue.append((neighbor_result, n["similarity"], combined_temporal))
+
+                if budget_remaining <= 0:
+                    break
+
+    return results
+
+
+async def retrieve_parallel(
+    pool,
+    query_text: str,
+    query_embedding_str: str,
+    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]]]:
+    """
+    Run 3-way or 4-way parallel retrieval (adds temporal if detected).
+
+    Args:
+        pool: Database connection pool
+        query_text: Query text
+        query_embedding_str: Query embedding as string
+        agent_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)
+
+    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
+    """
+    # 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
+    )
+
+    # Wrapper to track timing for each retrieval method
+    async def timed_retrieval(name: str, coro):
+        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)
+
+    async def run_bm25():
+        async with acquire_with_retry(pool) as conn:
+            return await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget)
+
+    async def run_graph():
+        async with acquire_with_retry(pool) as conn:
+            return await retrieve_graph(conn, query_embedding_str, bank_id, fact_type, budget=thinking_budget)
+
+    async def run_temporal(start_date, end_date):
+        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
+            )
+
+    # 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))
+        )
+        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_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

hindsight_api/engine/search/scoring.py

@@ -0,0 +1,161 @@
+"""
+Scoring functions for memory search and retrieval.
+
+Includes recency weighting, frequency weighting, temporal proximity,
+and similarity calculations used in memory activation and ranking.
+"""
+from datetime import datetime
+from typing import List
+
+
+def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
+    """
+    Calculate cosine similarity between two vectors.
+
+    Args:
+        vec1: First vector
+        vec2: Second vector
+
+    Returns:
+        Similarity score between 0 and 1
+    """
+    if len(vec1) != len(vec2):
+        raise ValueError("Vectors must have same dimension")
+
+    dot_product = sum(a * b for a, b in zip(vec1, vec2))
+    magnitude1 = sum(a * a for a in vec1) ** 0.5
+    magnitude2 = sum(b * b for b in vec2) ** 0.5
+
+    if magnitude1 == 0 or magnitude2 == 0:
+        return 0.0
+
+    return dot_product / (magnitude1 * magnitude2)
+
+
+def calculate_recency_weight(days_since: float, half_life_days: float = 365.0) -> float:
+    """
+    Calculate recency weight using logarithmic decay.
+
+    This provides much better differentiation over long time periods compared to
+    exponential decay. Uses a log-based decay where the half-life parameter controls
+    when memories reach 50% weight.
+
+    Examples:
+        - Today (0 days): 1.0
+        - 1 year (365 days): ~0.5 (with default half_life=365)
+        - 2 years (730 days): ~0.33
+        - 5 years (1825 days): ~0.17
+        - 10 years (3650 days): ~0.09
+
+    This ensures that 2-year-old and 5-year-old memories have meaningfully
+    different weights, unlike exponential decay which makes them both ~0.
+
+    Args:
+        days_since: Number of days since the memory was created
+        half_life_days: Number of days for weight to reach 0.5 (default: 1 year)
+
+    Returns:
+        Weight between 0 and 1
+    """
+    import math
+    # Logarithmic decay: 1 / (1 + log(1 + days_since/half_life))
+    # This decays much slower than exponential, giving better long-term differentiation
+    normalized_age = days_since / half_life_days
+    return 1.0 / (1.0 + math.log1p(normalized_age))
+
+
+def calculate_frequency_weight(access_count: int, max_boost: float = 2.0) -> float:
+    """
+    Calculate frequency weight based on access count.
+
+    Frequently accessed memories are weighted higher.
+    Uses logarithmic scaling to avoid over-weighting.
+
+    Args:
+        access_count: Number of times the memory was accessed
+        max_boost: Maximum multiplier for frequently accessed memories
+
+    Returns:
+        Weight between 1.0 and max_boost
+    """
+    import math
+    if access_count <= 0:
+        return 1.0
+
+    # Logarithmic scaling: log(access_count + 1) / log(10)
+    # This gives: 0 accesses = 1.0, 9 accesses ~= 1.5, 99 accesses ~= 2.0
+    normalized = math.log(access_count + 1) / math.log(10)
+    return 1.0 + min(normalized, max_boost - 1.0)
+
+
+def calculate_temporal_anchor(occurred_start: datetime, occurred_end: datetime) -> datetime:
+    """
+    Calculate a single temporal anchor point from a temporal range.
+
+    Used for spreading activation - we need a single representative date
+    to calculate temporal proximity between facts. This simplifies the
+    range-to-range distance problem.
+
+    Strategy: Use midpoint of the range for balanced representation.
+
+    Args:
+        occurred_start: Start of temporal range
+        occurred_end: End of temporal range
+
+    Returns:
+        Single datetime representing the temporal anchor (midpoint)
+
+    Examples:
+        - Point event (July 14): start=July 14, end=July 14 → anchor=July 14
+        - Month range (February): start=Feb 1, end=Feb 28 → anchor=Feb 14
+        - Year range (2023): start=Jan 1, end=Dec 31 → anchor=July 1
+    """
+    # Calculate midpoint
+    time_delta = occurred_end - occurred_start
+    midpoint = occurred_start + (time_delta / 2)
+    return midpoint
+
+
+def calculate_temporal_proximity(
+    anchor_a: datetime,
+    anchor_b: datetime,
+    half_life_days: float = 30.0
+) -> float:
+    """
+    Calculate temporal proximity between two temporal anchors.
+
+    Used for spreading activation to determine how "close" two facts are
+    in time. Uses logarithmic decay so that temporal similarity doesn't
+    drop off too quickly.
+
+    Args:
+        anchor_a: Temporal anchor of first fact
+        anchor_b: Temporal anchor of second fact
+        half_life_days: Number of days for proximity to reach 0.5
+                       (default: 30 days = 1 month)
+
+    Returns:
+        Proximity score in [0, 1] where:
+        - 1.0 = same day
+        - 0.5 = ~half_life days apart
+        - 0.0 = very distant in time
+
+    Examples:
+        - Same day: 1.0
+        - 1 week apart (half_life=30): ~0.7
+        - 1 month apart (half_life=30): ~0.5
+        - 1 year apart (half_life=30): ~0.2
+    """
+    import math
+
+    days_apart = abs((anchor_a - anchor_b).days)
+
+    if days_apart == 0:
+        return 1.0
+
+    # Logarithmic decay: 1 / (1 + log(1 + days_apart/half_life))
+    # Similar to calculate_recency_weight but for proximity between events
+    normalized_distance = days_apart / half_life_days
+    proximity = 1.0 / (1.0 + math.log1p(normalized_distance))
+
+    return proximity