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

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
+        ]