0bcbf8491b
* fix: return metadata in recall responses (#674) Metadata stored during retain was never retrieved during recall. Add metadata to all SQL SELECT queries, the RetrievalResult dataclass, ScoredResult.to_dict(), and MemoryFact construction in the recall pipeline. * test: add metadata round-trip test for retain→recall Replace placeholder metadata test with one that actually passes metadata via retain_batch_async and asserts it is returned on recall. * fix: parse metadata JSON string from database in MemoryFact asyncpg may return JSONB columns as strings. Add a field_validator to MemoryFact.metadata to handle JSON string deserialization.
702 lines
25 KiB
Python
702 lines
25 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)
|
||
- Lazy edge loading: only loads edges for frontier nodes, not entire graph
|
||
- 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 ..memory_engine import fq_table
|
||
from .graph_retrieval import GraphRetriever
|
||
from .tags import TagGroup, TagsMatch
|
||
from .types import MPFPTimings, RetrievalResult
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Data Classes
|
||
# -----------------------------------------------------------------------------
|
||
|
||
|
||
@dataclass
|
||
class EdgeTarget:
|
||
"""A neighbor node with its edge weight."""
|
||
|
||
node_id: str
|
||
weight: float
|
||
|
||
|
||
@dataclass
|
||
class EdgeCache:
|
||
"""
|
||
Cache for lazily-loaded edges.
|
||
|
||
Grows per-hop as edges are loaded for frontier nodes.
|
||
Shared across patterns to avoid redundant loads.
|
||
Loads ALL edge types at once to minimize DB queries.
|
||
Thread-safe via asyncio lock to prevent redundant concurrent loads.
|
||
"""
|
||
|
||
# edge_type -> from_node_id -> list of EdgeTarget
|
||
graphs: dict[str, dict[str, list[EdgeTarget]]] = field(default_factory=dict)
|
||
# Track which nodes have been fully loaded (all edge types)
|
||
_fully_loaded: set[str] = field(default_factory=set)
|
||
# Timing stats
|
||
db_queries: int = 0
|
||
edge_load_time: float = 0.0
|
||
# Detailed hop timing for debugging
|
||
hop_details: list[dict] = field(default_factory=list)
|
||
# Lock to prevent redundant concurrent loads
|
||
_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
||
|
||
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]
|
||
|
||
def is_fully_loaded(self, node_id: str) -> bool:
|
||
"""Check if all edges for this node have been loaded."""
|
||
return node_id in self._fully_loaded
|
||
|
||
def get_uncached(self, node_ids: list[str]) -> list[str]:
|
||
"""Get node IDs that haven't been fully loaded yet."""
|
||
return [n for n in node_ids if not self.is_fully_loaded(n)]
|
||
|
||
def add_all_edges(self, edges_by_type: dict[str, dict[str, list[EdgeTarget]]], all_queried: list[str]):
|
||
"""
|
||
Add loaded edges to the cache (all edge types at once).
|
||
|
||
Args:
|
||
edges_by_type: Dict mapping edge_type -> from_node_id -> list of EdgeTarget
|
||
all_queried: All node IDs that were queried (marks them as fully loaded)
|
||
"""
|
||
for edge_type, edges in edges_by_type.items():
|
||
if edge_type not in self.graphs:
|
||
self.graphs[edge_type] = {}
|
||
for node_id, neighbors in edges.items():
|
||
self.graphs[edge_type][node_id] = neighbors
|
||
|
||
# Mark all queried nodes as fully loaded (even if they have no edges)
|
||
self._fully_loaded.update(all_queried)
|
||
|
||
|
||
@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)
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Lazy Edge Loading
|
||
# -----------------------------------------------------------------------------
|
||
|
||
|
||
async def load_all_edges_for_frontier(
|
||
pool,
|
||
node_ids: list[str],
|
||
top_k_per_type: int = 20,
|
||
) -> dict[str, dict[str, list[EdgeTarget]]]:
|
||
"""
|
||
Load top-k edges per (node, edge_type) for frontier nodes.
|
||
|
||
Uses a LATERAL join to efficiently fetch only the top-k edges per type,
|
||
avoiding loading hundreds of entity edges when only 20 are needed.
|
||
|
||
Requires composite index: (from_unit_id, link_type, weight DESC)
|
||
|
||
Args:
|
||
pool: Database connection pool
|
||
node_ids: Frontier node IDs to load edges for
|
||
top_k_per_type: Max edges to load per (node, link_type) pair
|
||
|
||
Returns:
|
||
Dict mapping edge_type -> from_node_id -> list of EdgeTarget
|
||
"""
|
||
if not node_ids:
|
||
return {}
|
||
|
||
async with acquire_with_retry(pool) as conn:
|
||
# Use LATERAL join to get top-k per (from_node, link_type)
|
||
# This leverages the composite index for efficient early termination
|
||
rows = await conn.fetch(
|
||
f"""
|
||
WITH frontier(node_id) AS (SELECT unnest($1::uuid[]))
|
||
SELECT f.node_id as from_unit_id, lt.link_type, edges.to_unit_id, edges.weight
|
||
FROM frontier f
|
||
CROSS JOIN (VALUES ('semantic'), ('temporal'), ('entity'), ('causes'), ('caused_by')) AS lt(link_type)
|
||
CROSS JOIN LATERAL (
|
||
SELECT ml.to_unit_id, ml.weight
|
||
FROM {fq_table("memory_links")} ml
|
||
WHERE ml.from_unit_id = f.node_id
|
||
AND ml.link_type = lt.link_type
|
||
AND ml.weight >= 0.1
|
||
ORDER BY ml.weight DESC
|
||
LIMIT $2
|
||
) edges
|
||
""",
|
||
node_ids,
|
||
top_k_per_type,
|
||
)
|
||
|
||
# Group by edge_type -> from_node -> neighbors
|
||
result: dict[str, dict[str, list[EdgeTarget]]] = defaultdict(lambda: defaultdict(list))
|
||
for row in rows:
|
||
edge_type = row["link_type"]
|
||
from_id = str(row["from_unit_id"])
|
||
to_id = str(row["to_unit_id"])
|
||
weight = row["weight"]
|
||
result[edge_type][from_id].append(EdgeTarget(node_id=to_id, weight=weight))
|
||
|
||
# Convert nested defaultdicts to regular dicts
|
||
return {edge_type: dict(edges) for edge_type, edges in result.items()}
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Core Algorithm (Async with Lazy Loading)
|
||
# -----------------------------------------------------------------------------
|
||
|
||
|
||
@dataclass
|
||
class PatternState:
|
||
"""State for a pattern traversal between hops."""
|
||
|
||
pattern: list[str]
|
||
hop_index: int
|
||
scores: dict[str, float]
|
||
frontier: dict[str, float]
|
||
|
||
|
||
def _init_pattern_state(seeds: list[SeedNode], pattern: list[str]) -> PatternState:
|
||
"""Initialize pattern state from seeds."""
|
||
if not seeds:
|
||
return PatternState(pattern=pattern, hop_index=0, scores={}, frontier={})
|
||
|
||
total_seed_score = sum(s.score for s in seeds)
|
||
if total_seed_score == 0:
|
||
total_seed_score = len(seeds)
|
||
|
||
frontier = {s.node_id: s.score / total_seed_score for s in seeds}
|
||
return PatternState(pattern=pattern, hop_index=0, scores={}, frontier=frontier)
|
||
|
||
|
||
def _execute_hop(state: PatternState, cache: EdgeCache, config: MPFPConfig) -> set[str]:
|
||
"""
|
||
Execute ONE hop of traversal, return frontier nodes for next hop.
|
||
|
||
This is a pure function that uses cached edges (no DB access).
|
||
Returns set of uncached nodes needed for next hop.
|
||
"""
|
||
if state.hop_index >= len(state.pattern):
|
||
return set()
|
||
|
||
edge_type = state.pattern[state.hop_index]
|
||
|
||
# Collect active nodes above threshold
|
||
active_nodes = [node_id for node_id, mass in state.frontier.items() if mass >= config.threshold]
|
||
if not active_nodes:
|
||
state.frontier = {}
|
||
return set()
|
||
|
||
# Propagate mass using cached edges
|
||
next_frontier: dict[str, float] = {}
|
||
uncached_for_next: set[str] = set()
|
||
|
||
for node_id, mass in state.frontier.items():
|
||
if mass < config.threshold:
|
||
continue
|
||
|
||
# Keep α portion for this node
|
||
state.scores[node_id] = state.scores.get(node_id, 0) + config.alpha * mass
|
||
|
||
# Push (1-α) to neighbors
|
||
push_mass = (1 - config.alpha) * mass
|
||
neighbors = cache.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
|
||
# Track if we'll need edges for this node in the next hop
|
||
if not cache.is_fully_loaded(neighbor.node_id):
|
||
uncached_for_next.add(neighbor.node_id)
|
||
|
||
state.frontier = next_frontier
|
||
state.hop_index += 1
|
||
|
||
return uncached_for_next
|
||
|
||
|
||
def _finalize_pattern(state: PatternState, config: MPFPConfig) -> PatternResult:
|
||
"""Finalize pattern by adding remaining frontier mass to scores."""
|
||
for node_id, mass in state.frontier.items():
|
||
if mass >= config.threshold:
|
||
state.scores[node_id] = state.scores.get(node_id, 0) + mass
|
||
|
||
return PatternResult(pattern=state.pattern, scores=state.scores)
|
||
|
||
|
||
async def mpfp_traverse_hop_synchronized(
|
||
pool,
|
||
pattern_jobs: list[tuple[list[SeedNode], list[str]]],
|
||
config: MPFPConfig,
|
||
cache: EdgeCache,
|
||
) -> list[PatternResult]:
|
||
"""
|
||
Execute ALL patterns with hop-synchronized edge loading.
|
||
|
||
Instead of running each pattern independently (causing multiple DB queries),
|
||
this function:
|
||
1. Runs hop 1 for ALL patterns (using pre-warmed seed edges)
|
||
2. Collects ALL unique hop-2 frontier nodes across patterns
|
||
3. Pre-warms hop-2 edges in ONE query
|
||
4. Runs hop 2 for ALL patterns
|
||
|
||
This reduces DB queries from O(patterns * hops) to O(hops).
|
||
|
||
Args:
|
||
pool: Database connection pool
|
||
pattern_jobs: List of (seeds, pattern) tuples
|
||
config: Algorithm parameters
|
||
cache: Shared edge cache (should be pre-warmed with seed edges)
|
||
|
||
Returns:
|
||
List of PatternResult for each pattern
|
||
"""
|
||
import time
|
||
|
||
# Initialize all pattern states
|
||
states = [_init_pattern_state(seeds, pattern) for seeds, pattern in pattern_jobs]
|
||
|
||
# Determine max hops (all patterns should be same length, but be safe)
|
||
max_hops = max((len(p) for _, p in pattern_jobs), default=0)
|
||
|
||
# Detailed timing for debugging
|
||
hop_times: list[dict] = []
|
||
|
||
# Execute hop-by-hop across ALL patterns
|
||
for hop in range(max_hops):
|
||
hop_start = time.time()
|
||
hop_timing = {"hop": hop, "patterns_executed": 0, "uncached_count": 0, "load_time": 0.0}
|
||
|
||
# Execute this hop for all patterns, collect uncached nodes for next hop
|
||
all_uncached: set[str] = set()
|
||
exec_start = time.time()
|
||
for state in states:
|
||
if state.hop_index < len(state.pattern):
|
||
uncached = _execute_hop(state, cache, config)
|
||
all_uncached.update(uncached)
|
||
hop_timing["patterns_executed"] += 1
|
||
hop_timing["exec_time"] = time.time() - exec_start
|
||
|
||
# Pre-warm edges for ALL uncached nodes before next hop
|
||
hop_timing["uncached_count"] = len(all_uncached)
|
||
if all_uncached:
|
||
uncached_list = list(all_uncached - cache._fully_loaded)
|
||
hop_timing["uncached_after_filter"] = len(uncached_list)
|
||
if uncached_list:
|
||
load_start = time.time()
|
||
edges_by_type = await load_all_edges_for_frontier(pool, uncached_list, config.top_k_neighbors)
|
||
hop_timing["load_time"] = time.time() - load_start
|
||
cache.edge_load_time += hop_timing["load_time"]
|
||
cache.db_queries += 1
|
||
cache.add_all_edges(edges_by_type, uncached_list)
|
||
hop_timing["edges_loaded"] = sum(
|
||
len(neighbors) for edges in edges_by_type.values() for neighbors in edges.values()
|
||
)
|
||
|
||
hop_timing["total_time"] = time.time() - hop_start
|
||
hop_times.append(hop_timing)
|
||
|
||
# Store hop timing details in cache for logging
|
||
cache.hop_details = hop_times
|
||
|
||
# Finalize all patterns
|
||
return [_finalize_pattern(state, config) for state in states]
|
||
|
||
|
||
async def mpfp_traverse_async(
|
||
pool,
|
||
seeds: list[SeedNode],
|
||
pattern: list[str],
|
||
config: MPFPConfig,
|
||
cache: EdgeCache,
|
||
) -> PatternResult:
|
||
"""
|
||
Async Forward Push traversal with lazy edge loading.
|
||
|
||
NOTE: For better performance with multiple patterns, use mpfp_traverse_hop_synchronized().
|
||
This function is kept for single-pattern use cases.
|
||
"""
|
||
if not seeds:
|
||
return PatternResult(pattern=pattern, scores={})
|
||
|
||
results = await mpfp_traverse_hop_synchronized(pool, [(seeds, pattern)], config, cache)
|
||
return results[0] if results else PatternResult(pattern=pattern, 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 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(
|
||
f"""
|
||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||
mentioned_at, fact_type, document_id, chunk_id, tags, metadata
|
||
FROM {fq_table("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 with lazy edge loading.
|
||
|
||
Runs predefined patterns in parallel from semantic and temporal seeds,
|
||
loading edges on-demand per hop instead of loading entire graph upfront.
|
||
"""
|
||
|
||
def __init__(self, config: MPFPConfig | None = None):
|
||
"""
|
||
Initialize MPFP retriever.
|
||
|
||
Args:
|
||
config: Algorithm configuration (uses defaults if None)
|
||
"""
|
||
if config is None:
|
||
# Read top_k_neighbors from global config
|
||
from ...config import get_config
|
||
|
||
global_config = get_config()
|
||
config = MPFPConfig(top_k_neighbors=global_config.mpfp_top_k_neighbors)
|
||
self.config = config
|
||
|
||
@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,
|
||
adjacency=None, # Ignored - kept for interface compatibility
|
||
tags: list[str] | None = None,
|
||
tags_match: TagsMatch = "any",
|
||
tag_groups: list[TagGroup] | None = None,
|
||
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
|
||
"""
|
||
Retrieve facts using MPFP algorithm with lazy edge loading.
|
||
|
||
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
|
||
adjacency: Ignored (kept for interface compatibility)
|
||
tags: Optional list of tags for visibility filtering (OR matching)
|
||
|
||
Returns:
|
||
Tuple of (List of RetrievalResult with activation scores, MPFPTimings)
|
||
"""
|
||
import time
|
||
|
||
timings = MPFPTimings(fact_type=fact_type)
|
||
|
||
# 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:
|
||
seeds_start = time.time()
|
||
semantic_seed_nodes = await self._find_semantic_seeds(
|
||
pool,
|
||
query_embedding_str,
|
||
bank_id,
|
||
fact_type,
|
||
tags=tags,
|
||
tags_match=tags_match,
|
||
tag_groups=tag_groups,
|
||
)
|
||
timings.seeds_time = time.time() - seeds_start
|
||
logger.debug(
|
||
f"[MPFP] Found {len(semantic_seed_nodes)} semantic seeds for fact_type={fact_type} (tags={tags}, tags_match={tags_match})"
|
||
)
|
||
|
||
# Collect all pattern jobs
|
||
pattern_jobs = []
|
||
|
||
# Patterns from semantic seeds
|
||
for pattern in self.config.patterns_semantic:
|
||
if semantic_seed_nodes:
|
||
pattern_jobs.append((semantic_seed_nodes, pattern))
|
||
|
||
# Patterns from temporal seeds
|
||
for pattern in self.config.patterns_temporal:
|
||
if temporal_seed_nodes:
|
||
pattern_jobs.append((temporal_seed_nodes, pattern))
|
||
|
||
if not pattern_jobs:
|
||
logger.debug(
|
||
f"[MPFP] No pattern jobs (semantic_seeds={len(semantic_seed_nodes)}, temporal_seeds={len(temporal_seed_nodes)})"
|
||
)
|
||
return [], timings
|
||
|
||
timings.pattern_count = len(pattern_jobs)
|
||
|
||
# Shared edge cache across all patterns
|
||
cache = EdgeCache()
|
||
|
||
# Pre-warm cache with ALL seed node edges BEFORE running patterns
|
||
# This prevents redundant DB queries at hop 1
|
||
all_seed_ids = list({s.node_id for seeds, _ in pattern_jobs for s in seeds})
|
||
if all_seed_ids:
|
||
import time as time_module
|
||
|
||
prewarm_start = time_module.time()
|
||
edges_by_type = await load_all_edges_for_frontier(pool, all_seed_ids, self.config.top_k_neighbors)
|
||
cache.edge_load_time += time_module.time() - prewarm_start
|
||
cache.db_queries += 1
|
||
cache.add_all_edges(edges_by_type, all_seed_ids)
|
||
|
||
# Run all patterns with HOP-SYNCHRONIZED edge loading
|
||
# This batches hop-2 edge loads across ALL patterns into ONE query
|
||
# Reduces DB queries from O(patterns * hops) to O(hops)
|
||
step_start = time.time()
|
||
pattern_results = await mpfp_traverse_hop_synchronized(pool, pattern_jobs, self.config, cache)
|
||
timings.traverse = time.time() - step_start
|
||
|
||
# Record edge loading stats from cache
|
||
timings.edge_count = sum(len(neighbors) for g in cache.graphs.values() for neighbors in g.values())
|
||
timings.db_queries = cache.db_queries
|
||
timings.edge_load_time = cache.edge_load_time
|
||
timings.hop_details = cache.hop_details
|
||
|
||
# Fuse results
|
||
step_start = time.time()
|
||
fused = rrf_fusion(pattern_results, top_k=budget)
|
||
timings.fusion = time.time() - step_start
|
||
|
||
if not fused:
|
||
logger.debug(f"[MPFP] No fused results after RRF fusion (pattern_count={len(pattern_results)})")
|
||
return [], timings
|
||
|
||
# Get top result IDs
|
||
result_ids = [node_id for node_id, score in fused][:budget]
|
||
|
||
# Fetch full details
|
||
step_start = time.time()
|
||
results = await fetch_memory_units_by_ids(pool, result_ids, fact_type)
|
||
timings.fetch = time.time() - step_start
|
||
|
||
# Filter results by tags (graph traversal may have picked up unfiltered memories)
|
||
if tags:
|
||
from .tags import filter_results_by_tags
|
||
|
||
results = filter_results_by_tags(results, tags, match=tags_match)
|
||
|
||
# Apply compound tag group filtering (post-traversal)
|
||
if tag_groups:
|
||
from .tags import filter_results_by_tag_groups
|
||
|
||
results = filter_results_by_tag_groups(results, tag_groups)
|
||
|
||
timings.result_count = len(results)
|
||
|
||
# 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, timings
|
||
|
||
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,
|
||
tags: list[str] | None = None,
|
||
tags_match: TagsMatch = "any",
|
||
tag_groups: list[TagGroup] | None = None,
|
||
) -> list[SeedNode]:
|
||
"""Fallback: find semantic seeds via embedding search."""
|
||
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
|
||
|
||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
|
||
if tags:
|
||
params.append(tags)
|
||
params.extend(groups_params)
|
||
|
||
async with acquire_with_retry(pool) as conn:
|
||
rows = await conn.fetch(
|
||
f"""
|
||
SELECT id, 1 - (embedding <=> $1::vector) AS similarity
|
||
FROM {fq_table("memory_units")}
|
||
WHERE bank_id = $2
|
||
AND embedding IS NOT NULL
|
||
AND fact_type = $3
|
||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||
{tags_clause}
|
||
{groups_clause}
|
||
ORDER BY embedding <=> $1::vector
|
||
LIMIT $5
|
||
""",
|
||
*params,
|
||
)
|
||
|
||
return [SeedNode(node_id=str(r["id"]), score=r["similarity"]) for r in rows]
|