5de793eec7
* feat: add compound tag filtering via tag_groups
Adds tag_groups to RecallRequest and ReflectRequest to express arbitrary
boolean tag predicates: leaf {tags, match}, and/or/not compounds.
Top-level groups are AND-ed. Existing tags/tags_match unchanged.
Examples:
Step filter AND user scope:
tag_groups: [{tags: ["step:5","step:8"], match: "any_strict"},
{tags: ["user:alice"], match: "all_strict"}]
Exclusion:
tag_groups: [{tags: ["user:alice"], match: "all_strict"},
{not: {tags: ["archived"], match: "any_strict"}}]
- Recursive SQL builder (build_tag_groups_where_clause) threads through
all 4 retrieval strategies (semantic/BM25, temporal, graph, MPFP)
- Python-side filter (filter_results_by_tag_groups) for post-traversal
- 22 new unit tests
- OpenAPI spec + all clients regenerated (Rust, Python, TypeScript, Go)
* fix: add tag_groups: None to Rust CLI struct initializers
* fix: add tag_groups: None to Rust client test RecallRequest initializer
* feat: reject tags+tag_groups together, add tag_groups integration tests
- Add model_validator to RecallRequest and ReflectRequest that returns 422
when both `tags` and `tag_groups` are set (mutually exclusive)
- Add 5 integration tests for tag_groups compound filtering:
* validation: 422 when both fields are set
* AND filter: two leaf groups (step scope AND user scope)
* OR compound: user:alice OR user:bob
* NOT compound: user:alice AND NOT archived
* Nested: user:alice AND (step:5 OR step:8)
* ci: trigger CI run
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,
|
|
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
|