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

hindsight_api/engine/search/mpfp_retrieval.py

@@ -0,0 +1,438 @@
+"""
+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 collections import defaultdict
+from dataclasses import dataclass, field
+
+from ..db_utils import acquire_with_retry
+from .graph_retrieval import GraphRetriever
+from .types import RetrievalResult
+
+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: MPFPConfig | None = 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: str | None = None,
+        semantic_seeds: list[RetrievalResult] | None = None,
+        temporal_seeds: list[RetrievalResult] | None = 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: list[RetrievalResult] | None,
+        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]

hindsight_api/engine/search/observation_utils.py

@@ -6,7 +6,7 @@ about an entity, without personality influence.
 """
 
 import logging
-from typing import List, Dict, Any
+
 from pydantic import BaseModel, Field
 
 from ..response_models import MemoryFact
@@ -16,18 +16,17 @@ logger = logging.getLogger(__name__)
 
 class Observation(BaseModel):
     """An observation about an entity."""
+
     observation: str = Field(description="The observation text - a factual statement about the entity")
 
 
 class ObservationExtractionResponse(BaseModel):
     """Response containing extracted observations."""
-    observations: List[Observation] = Field(
-        default_factory=list,
-        description="List of observations about the entity"
-    )
 
+    observations: list[Observation] = Field(default_factory=list, description="List of observations about the entity")
 
-def format_facts_for_observation_prompt(facts: List[MemoryFact]) -> str:
+
+def format_facts_for_observation_prompt(facts: list[MemoryFact]) -> str:
     """Format facts as text for observation extraction prompt."""
     import json
 
@@ -35,9 +34,7 @@ def format_facts_for_observation_prompt(facts: List[MemoryFact]) -> str:
         return "[]"
     formatted = []
     for fact in facts:
-        fact_obj = {
-            "text": fact.text
-        }
+        fact_obj = {"text": fact.text}
 
         # Add context if available
         if fact.context:
@@ -92,11 +89,7 @@ def get_observation_system_message() -> str:
     return "You are an objective observer synthesizing facts about an entity. Generate clear, factual observations without opinions or personality influence. Be concise and accurate."
 
 
-async def extract_observations_from_facts(
-    llm_config,
-    entity_name: str,
-    facts: List[MemoryFact]
-) -> List[str]:
+async def extract_observations_from_facts(llm_config, entity_name: str, facts: list[MemoryFact]) -> list[str]:
     """
     Extract observations from facts about an entity using LLM.
 
@@ -118,10 +111,10 @@ async def extract_observations_from_facts(
         result = await llm_config.call(
             messages=[
                 {"role": "system", "content": get_observation_system_message()},
-                {"role": "user", "content": prompt}
+                {"role": "user", "content": prompt},
             ],
             response_format=ObservationExtractionResponse,
-            scope="memory_extract_observation"
+            scope="memory_extract_observation",
         )
 
         observations = [op.observation for op in result.observations]

hindsight_api/engine/search/reranking.py

@@ -2,7 +2,6 @@
 Cross-encoder neural reranking for search results.
 """
 
-from typing import List
 from .types import MergedCandidate, ScoredResult
 
 
@@ -24,14 +23,11 @@ class CrossEncoderReranker:
         """
         if cross_encoder is None:
             from hindsight_api.engine.cross_encoder import create_cross_encoder_from_env
+
             cross_encoder = create_cross_encoder_from_env()
         self.cross_encoder = cross_encoder
 
-    def rerank(
-        self,
-        query: str,
-        candidates: List[MergedCandidate]
-    ) -> List[ScoredResult]:
+    def rerank(self, query: str, candidates: list[MergedCandidate]) -> list[ScoredResult]:
         """
         Rerank candidates using cross-encoder scores.
 
@@ -77,6 +73,7 @@ class CrossEncoderReranker:
         # Normalize scores using sigmoid to [0, 1] range
         # Cross-encoder returns logits which can be negative
         import numpy as np
+
         def sigmoid(x):
             return 1 / (1 + np.exp(-x))
 
@@ -89,7 +86,7 @@ class CrossEncoderReranker:
                 candidate=candidate,
                 cross_encoder_score=float(raw_score),
                 cross_encoder_score_normalized=float(norm_score),
-                weight=float(norm_score)  # Initial weight is just cross-encoder score
+                weight=float(norm_score),  # Initial weight is just cross-encoder score
             )
             scored_results.append(scored_result)