454 lines
14 KiB
Python
454 lines
14 KiB
Python
"""
|
||
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
|
||
]
|