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.
282 lines
11 KiB
Python
282 lines
11 KiB
Python
"""
|
|
Graph retrieval strategies for memory recall.
|
|
|
|
This module provides an abstraction for graph-based memory retrieval,
|
|
allowing different algorithms (BFS spreading activation, PPR, etc.) to be
|
|
swapped without changing the rest of the recall pipeline.
|
|
"""
|
|
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
|
|
from ..db_utils import acquire_with_retry
|
|
from ..memory_engine import fq_table
|
|
from .tags import TagGroup, TagsMatch, filter_results_by_tag_groups, filter_results_by_tags
|
|
from .types import MPFPTimings, RetrievalResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GraphRetriever(ABC):
|
|
"""
|
|
Abstract base class for graph-based memory retrieval.
|
|
|
|
Implementations traverse the memory graph (entity links, temporal links,
|
|
causal links) to find relevant facts that might not be found by
|
|
semantic or keyword search alone.
|
|
"""
|
|
|
|
@property
|
|
@abstractmethod
|
|
def name(self) -> str:
|
|
"""Return identifier for this retrieval strategy (e.g., 'bfs', 'mpfp')."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
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, # TypedAdjacency, optional pre-loaded graph
|
|
tags: list[str] | None = None, # Visibility scope tags for filtering
|
|
tags_match: TagsMatch = "any", # How to match tags: 'any' (OR) or 'all' (AND)
|
|
tag_groups: list[TagGroup] | None = None, # Compound boolean tag filter groups
|
|
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
|
|
"""
|
|
Retrieve relevant facts via graph traversal.
|
|
|
|
Args:
|
|
pool: Database connection pool
|
|
query_embedding_str: Query embedding as string (for finding entry points)
|
|
bank_id: Memory bank identifier
|
|
fact_type: Fact type to filter ('world', 'experience', 'opinion', 'observation')
|
|
budget: Maximum number of nodes to explore/return
|
|
query_text: Original query text (optional, for some strategies)
|
|
semantic_seeds: Pre-computed semantic entry points (from semantic retrieval)
|
|
temporal_seeds: Pre-computed temporal entry points (from temporal retrieval)
|
|
adjacency: Pre-loaded typed adjacency graph (optional, for MPFP)
|
|
tags: Optional list of tags for visibility filtering (OR matching)
|
|
|
|
Returns:
|
|
Tuple of (List of RetrievalResult with activation scores, optional timing info)
|
|
"""
|
|
pass
|
|
|
|
|
|
class BFSGraphRetriever(GraphRetriever):
|
|
"""
|
|
Graph retrieval using BFS-style spreading activation.
|
|
|
|
Starting from semantic entry points, spreads activation through
|
|
the memory graph (entity, temporal, causal links) using breadth-first
|
|
traversal with decaying activation.
|
|
|
|
This is the original Hindsight graph retrieval algorithm.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
entry_point_limit: int = 5,
|
|
entry_point_threshold: float = 0.5,
|
|
activation_decay: float = 0.8,
|
|
min_activation: float = 0.1,
|
|
batch_size: int = 20,
|
|
):
|
|
"""
|
|
Initialize BFS graph retriever.
|
|
|
|
Args:
|
|
entry_point_limit: Maximum number of entry points to start from
|
|
entry_point_threshold: Minimum semantic similarity for entry points
|
|
activation_decay: Decay factor per hop (activation *= decay)
|
|
min_activation: Minimum activation to continue spreading
|
|
batch_size: Number of nodes to process per batch (for neighbor fetching)
|
|
"""
|
|
self.entry_point_limit = entry_point_limit
|
|
self.entry_point_threshold = entry_point_threshold
|
|
self.activation_decay = activation_decay
|
|
self.min_activation = min_activation
|
|
self.batch_size = batch_size
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "bfs"
|
|
|
|
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, # Not used by BFS
|
|
tags: list[str] | None = None,
|
|
tags_match: TagsMatch = "any",
|
|
tag_groups: list[TagGroup] | None = None,
|
|
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
|
|
"""
|
|
Retrieve facts using BFS spreading activation.
|
|
|
|
Algorithm:
|
|
1. Find entry points (top semantic matches above threshold)
|
|
2. BFS traversal: visit neighbors, propagate decaying activation
|
|
3. Boost causal links (causes, enables, prevents)
|
|
4. Return visited nodes up to budget
|
|
|
|
Note: BFS finds its own entry points via embedding search.
|
|
The semantic_seeds, temporal_seeds, and adjacency parameters are accepted
|
|
for interface compatibility but not used.
|
|
"""
|
|
async with acquire_with_retry(pool) as conn:
|
|
results = await self._retrieve_with_conn(
|
|
conn,
|
|
query_embedding_str,
|
|
bank_id,
|
|
fact_type,
|
|
budget,
|
|
tags=tags,
|
|
tags_match=tags_match,
|
|
tag_groups=tag_groups,
|
|
)
|
|
return results, None
|
|
|
|
async def _retrieve_with_conn(
|
|
self,
|
|
conn,
|
|
query_embedding_str: str,
|
|
bank_id: str,
|
|
fact_type: str,
|
|
budget: int,
|
|
tags: list[str] | None = None,
|
|
tags_match: TagsMatch = "any",
|
|
tag_groups: list[TagGroup] | None = None,
|
|
) -> list[RetrievalResult]:
|
|
"""Internal implementation with connection."""
|
|
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, self.entry_point_threshold, self.entry_point_limit]
|
|
if tags:
|
|
params.append(tags)
|
|
params.extend(groups_params)
|
|
|
|
# Step 1: Find entry points
|
|
entry_points = await conn.fetch(
|
|
f"""
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
|
mentioned_at, fact_type, document_id, chunk_id, tags,
|
|
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,
|
|
)
|
|
|
|
if not entry_points:
|
|
logger.debug(
|
|
f"[BFS] No entry points found for fact_type={fact_type} (tags={tags}, tags_match={tags_match})"
|
|
)
|
|
return []
|
|
|
|
logger.debug(
|
|
f"[BFS] Found {len(entry_points)} entry points for fact_type={fact_type} "
|
|
f"(tags={tags}, tags_match={tags_match})"
|
|
)
|
|
|
|
# Step 2: BFS spreading activation
|
|
visited = set()
|
|
results = []
|
|
queue = [(RetrievalResult.from_db_row(dict(r)), r["similarity"]) for r in entry_points]
|
|
budget_remaining = budget
|
|
|
|
while queue and budget_remaining > 0:
|
|
# Collect a batch of nodes to process
|
|
batch_nodes = []
|
|
batch_activations = {}
|
|
|
|
while queue and len(batch_nodes) < self.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
|
|
current.activation = activation
|
|
results.append(current)
|
|
batch_nodes.append(current.id)
|
|
batch_activations[unit_id] = activation
|
|
|
|
# Batch fetch neighbors
|
|
if batch_nodes and budget_remaining > 0:
|
|
max_neighbors = len(batch_nodes) * 20
|
|
neighbors = await conn.fetch(
|
|
f"""
|
|
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end,
|
|
mu.mentioned_at, mu.fact_type,
|
|
mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
|
ml.weight, ml.link_type, ml.from_unit_id
|
|
FROM {fq_table("memory_links")} ml
|
|
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
|
WHERE ml.from_unit_id = ANY($1::uuid[])
|
|
AND ml.weight >= $2
|
|
AND mu.fact_type = $3
|
|
ORDER BY ml.weight DESC
|
|
LIMIT $4
|
|
""",
|
|
batch_nodes,
|
|
self.min_activation,
|
|
fact_type,
|
|
max_neighbors,
|
|
)
|
|
|
|
for n in neighbors:
|
|
neighbor_id = str(n["id"])
|
|
if neighbor_id not in visited:
|
|
parent_id = str(n["from_unit_id"])
|
|
parent_activation = batch_activations.get(parent_id, 0.5)
|
|
|
|
# Boost causal links
|
|
link_type = n["link_type"]
|
|
base_weight = n["weight"]
|
|
|
|
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
|
|
|
|
effective_weight = base_weight * causal_boost
|
|
new_activation = parent_activation * effective_weight * self.activation_decay
|
|
|
|
if new_activation > self.min_activation:
|
|
neighbor_result = RetrievalResult.from_db_row(dict(n))
|
|
queue.append((neighbor_result, new_activation))
|
|
|
|
# Apply tags filtering (BFS may traverse into memories that don't match tags criteria)
|
|
if tags:
|
|
results = filter_results_by_tags(results, tags, match=tags_match)
|
|
|
|
# Apply compound tag group filtering (post-traversal)
|
|
if tag_groups:
|
|
results = filter_results_by_tag_groups(results, tag_groups)
|
|
|
|
return results
|