555 lines
19 KiB
Python
555 lines
19 KiB
Python
"""
|
|
Tests for MPFP (Meta-Path Forward Push) graph retrieval.
|
|
|
|
Tests cover:
|
|
1. EdgeCache - lazy caching behavior
|
|
2. mpfp_traverse_async - core traversal algorithm
|
|
3. load_edges_for_frontier - lazy edge loading
|
|
4. rrf_fusion - result fusion
|
|
5. MPFPGraphRetriever - full integration
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
from datetime import datetime, timezone
|
|
|
|
from hindsight_api.engine.search.mpfp_retrieval import (
|
|
EdgeCache,
|
|
EdgeTarget,
|
|
MPFPConfig,
|
|
MPFPGraphRetriever,
|
|
PatternResult,
|
|
SeedNode,
|
|
load_all_edges_for_frontier,
|
|
mpfp_traverse_async,
|
|
rrf_fusion,
|
|
)
|
|
from hindsight_api.engine.search.types import RetrievalResult
|
|
|
|
|
|
class TestEdgeCache:
|
|
"""Tests for the EdgeCache lazy loading cache."""
|
|
|
|
def test_empty_cache_returns_empty_neighbors(self):
|
|
"""Empty cache should return empty list for any node."""
|
|
cache = EdgeCache()
|
|
neighbors = cache.get_neighbors("semantic", "node-1")
|
|
assert neighbors == []
|
|
|
|
def test_is_fully_loaded_false_for_uncached(self):
|
|
"""is_fully_loaded should return False for nodes not yet loaded."""
|
|
cache = EdgeCache()
|
|
assert cache.is_fully_loaded("node-1") is False
|
|
|
|
def test_add_all_edges_marks_as_fully_loaded(self):
|
|
"""Adding edges should mark nodes as fully loaded."""
|
|
cache = EdgeCache()
|
|
|
|
edges_by_type = {
|
|
"semantic": {"node-1": [EdgeTarget("node-2", 0.8), EdgeTarget("node-3", 0.6)]},
|
|
}
|
|
cache.add_all_edges(edges_by_type, ["node-1", "node-4"]) # node-4 has no edges
|
|
|
|
assert cache.is_fully_loaded("node-1") is True
|
|
assert cache.is_fully_loaded("node-4") is True # Marked even with no edges
|
|
assert cache.is_fully_loaded("node-2") is False # Target, not source
|
|
|
|
def test_get_neighbors_returns_added_edges(self):
|
|
"""get_neighbors should return edges after add_all_edges."""
|
|
cache = EdgeCache()
|
|
|
|
edges_by_type = {
|
|
"semantic": {"node-1": [EdgeTarget("node-2", 0.8), EdgeTarget("node-3", 0.6)]},
|
|
}
|
|
cache.add_all_edges(edges_by_type, ["node-1"])
|
|
|
|
neighbors = cache.get_neighbors("semantic", "node-1")
|
|
assert len(neighbors) == 2
|
|
assert neighbors[0].node_id == "node-2"
|
|
assert neighbors[0].weight == 0.8
|
|
|
|
def test_get_uncached_filters_loaded_nodes(self):
|
|
"""get_uncached should only return nodes not yet fully loaded."""
|
|
cache = EdgeCache()
|
|
|
|
# Load some nodes (all edge types)
|
|
cache.add_all_edges({"semantic": {"node-1": []}}, ["node-1", "node-2"])
|
|
|
|
# Check uncached
|
|
uncached = cache.get_uncached(["node-1", "node-2", "node-3", "node-4"])
|
|
assert set(uncached) == {"node-3", "node-4"}
|
|
|
|
def test_get_normalized_neighbors_normalizes_weights(self):
|
|
"""get_normalized_neighbors should normalize weights to sum to 1."""
|
|
cache = EdgeCache()
|
|
|
|
edges_by_type = {
|
|
"semantic": {
|
|
"node-1": [
|
|
EdgeTarget("node-2", 0.8),
|
|
EdgeTarget("node-3", 0.4),
|
|
EdgeTarget("node-4", 0.2),
|
|
],
|
|
},
|
|
}
|
|
cache.add_all_edges(edges_by_type, ["node-1"])
|
|
|
|
# Get top 2, normalized
|
|
neighbors = cache.get_normalized_neighbors("semantic", "node-1", top_k=2)
|
|
assert len(neighbors) == 2
|
|
|
|
# Weights should sum to 1
|
|
total = sum(n.weight for n in neighbors)
|
|
assert abs(total - 1.0) < 0.001
|
|
|
|
# node-2 should have higher normalized weight than node-3
|
|
assert neighbors[0].node_id == "node-2"
|
|
assert neighbors[1].node_id == "node-3"
|
|
# Original: 0.8 and 0.4, so normalized: 0.8/1.2 and 0.4/1.2
|
|
assert abs(neighbors[0].weight - 0.8 / 1.2) < 0.001
|
|
assert abs(neighbors[1].weight - 0.4 / 1.2) < 0.001
|
|
|
|
def test_different_edge_types_are_separate(self):
|
|
"""Different edge types should be stored separately."""
|
|
cache = EdgeCache()
|
|
|
|
edges_by_type = {
|
|
"semantic": {"node-1": [EdgeTarget("node-2", 0.8)]},
|
|
"temporal": {"node-1": [EdgeTarget("node-3", 0.5)]},
|
|
}
|
|
cache.add_all_edges(edges_by_type, ["node-1"])
|
|
|
|
semantic_neighbors = cache.get_neighbors("semantic", "node-1")
|
|
temporal_neighbors = cache.get_neighbors("temporal", "node-1")
|
|
|
|
assert len(semantic_neighbors) == 1
|
|
assert semantic_neighbors[0].node_id == "node-2"
|
|
|
|
assert len(temporal_neighbors) == 1
|
|
assert temporal_neighbors[0].node_id == "node-3"
|
|
|
|
|
|
class TestRRFFusion:
|
|
"""Tests for RRF (Reciprocal Rank Fusion)."""
|
|
|
|
def test_empty_results(self):
|
|
"""Empty results should return empty fusion."""
|
|
fused = rrf_fusion([])
|
|
assert fused == []
|
|
|
|
def test_single_pattern_ranking(self):
|
|
"""Single pattern should preserve ranking order."""
|
|
result = PatternResult(
|
|
pattern=["semantic"],
|
|
scores={"node-1": 0.9, "node-2": 0.7, "node-3": 0.5},
|
|
)
|
|
|
|
fused = rrf_fusion([result], top_k=3)
|
|
assert len(fused) == 3
|
|
# node-1 should be first (highest score)
|
|
assert fused[0][0] == "node-1"
|
|
assert fused[1][0] == "node-2"
|
|
assert fused[2][0] == "node-3"
|
|
|
|
def test_multiple_patterns_boost_common_nodes(self):
|
|
"""Nodes appearing in multiple patterns should get boosted."""
|
|
result1 = PatternResult(
|
|
pattern=["semantic", "semantic"],
|
|
scores={"node-1": 0.9, "node-2": 0.7},
|
|
)
|
|
result2 = PatternResult(
|
|
pattern=["entity", "temporal"],
|
|
scores={"node-1": 0.8, "node-3": 0.6}, # node-1 in both
|
|
)
|
|
|
|
fused = rrf_fusion([result1, result2], top_k=3)
|
|
|
|
# node-1 should be first (appears in both patterns)
|
|
assert fused[0][0] == "node-1"
|
|
# Its score should be higher than others
|
|
assert fused[0][1] > fused[1][1]
|
|
|
|
def test_top_k_limits_results(self):
|
|
"""top_k should limit the number of results."""
|
|
result = PatternResult(
|
|
pattern=["semantic"],
|
|
scores={f"node-{i}": 1.0 / (i + 1) for i in range(10)},
|
|
)
|
|
|
|
fused = rrf_fusion([result], top_k=3)
|
|
assert len(fused) == 3
|
|
|
|
def test_empty_pattern_scores_ignored(self):
|
|
"""Patterns with empty scores should be ignored."""
|
|
result1 = PatternResult(pattern=["semantic"], scores={})
|
|
result2 = PatternResult(
|
|
pattern=["entity"],
|
|
scores={"node-1": 0.5},
|
|
)
|
|
|
|
fused = rrf_fusion([result1, result2], top_k=3)
|
|
assert len(fused) == 1
|
|
assert fused[0][0] == "node-1"
|
|
|
|
|
|
class TestMPFPTraverseAsync:
|
|
"""Tests for the async MPFP traversal algorithm."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_seeds_returns_empty(self):
|
|
"""Empty seeds should return empty result."""
|
|
cache = EdgeCache()
|
|
config = MPFPConfig()
|
|
|
|
result = await mpfp_traverse_async(
|
|
pool=None, # Not used when no seeds
|
|
seeds=[],
|
|
pattern=["semantic"],
|
|
config=config,
|
|
cache=cache,
|
|
)
|
|
|
|
assert result.scores == {}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_hop_no_edges(self):
|
|
"""Single hop with no edges should deposit mass at seeds."""
|
|
cache = EdgeCache()
|
|
config = MPFPConfig(alpha=0.15, threshold=1e-6)
|
|
|
|
# Pre-populate cache with empty edges for seed (marks as fully loaded)
|
|
cache.add_all_edges({}, ["seed-1"])
|
|
|
|
seeds = [SeedNode("seed-1", 1.0)]
|
|
|
|
with patch(
|
|
"hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier",
|
|
new_callable=AsyncMock,
|
|
return_value={},
|
|
):
|
|
result = await mpfp_traverse_async(
|
|
pool=MagicMock(),
|
|
seeds=seeds,
|
|
pattern=["semantic"],
|
|
config=config,
|
|
cache=cache,
|
|
)
|
|
|
|
# Seed should have alpha portion of its mass
|
|
assert "seed-1" in result.scores
|
|
assert result.scores["seed-1"] == pytest.approx(config.alpha, rel=0.01)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_hop_with_edges(self):
|
|
"""Single hop should spread mass to neighbors."""
|
|
cache = EdgeCache()
|
|
config = MPFPConfig(alpha=0.15, threshold=1e-6, top_k_neighbors=10)
|
|
|
|
seeds = [SeedNode("seed-1", 1.0)]
|
|
|
|
# Mock edge loading (returns all edge types at once)
|
|
async def mock_load_all_edges(pool, node_ids):
|
|
if "seed-1" in node_ids:
|
|
return {
|
|
"semantic": {
|
|
"seed-1": [
|
|
EdgeTarget("neighbor-1", 0.8),
|
|
EdgeTarget("neighbor-2", 0.4),
|
|
]
|
|
}
|
|
}
|
|
return {}
|
|
|
|
with patch(
|
|
"hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier",
|
|
side_effect=mock_load_all_edges,
|
|
):
|
|
result = await mpfp_traverse_async(
|
|
pool=MagicMock(),
|
|
seeds=seeds,
|
|
pattern=["semantic"],
|
|
config=config,
|
|
cache=cache,
|
|
)
|
|
|
|
# Seed keeps alpha portion
|
|
assert "seed-1" in result.scores
|
|
assert result.scores["seed-1"] == pytest.approx(config.alpha, rel=0.01)
|
|
|
|
# Neighbors get remaining mass (normalized)
|
|
assert "neighbor-1" in result.scores
|
|
assert "neighbor-2" in result.scores
|
|
|
|
# neighbor-1 should get more (higher weight)
|
|
assert result.scores["neighbor-1"] > result.scores["neighbor-2"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_two_hops(self):
|
|
"""Two-hop pattern should traverse through neighbors."""
|
|
cache = EdgeCache()
|
|
config = MPFPConfig(alpha=0.15, threshold=1e-6, top_k_neighbors=10)
|
|
|
|
seeds = [SeedNode("seed-1", 1.0)]
|
|
|
|
# Mock edge loading for two hops (returns all edge types at once)
|
|
async def mock_load_all_edges(pool, node_ids):
|
|
edges: dict[str, dict[str, list[EdgeTarget]]] = {"semantic": {}}
|
|
if "seed-1" in node_ids:
|
|
edges["semantic"]["seed-1"] = [EdgeTarget("hop1-node", 1.0)]
|
|
if "hop1-node" in node_ids:
|
|
edges["semantic"]["hop1-node"] = [EdgeTarget("hop2-node", 1.0)]
|
|
return edges
|
|
|
|
with patch(
|
|
"hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier",
|
|
side_effect=mock_load_all_edges,
|
|
):
|
|
result = await mpfp_traverse_async(
|
|
pool=MagicMock(),
|
|
seeds=seeds,
|
|
pattern=["semantic", "semantic"], # Two hops
|
|
config=config,
|
|
cache=cache,
|
|
)
|
|
|
|
# Should have scores for all three nodes
|
|
assert "seed-1" in result.scores
|
|
assert "hop1-node" in result.scores
|
|
assert "hop2-node" in result.scores
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_reuse(self):
|
|
"""Cache should prevent redundant edge loading."""
|
|
cache = EdgeCache()
|
|
config = MPFPConfig(alpha=0.15, threshold=1e-6)
|
|
|
|
# Pre-load cache (marks seed-1 as fully loaded)
|
|
cache.add_all_edges({"semantic": {"seed-1": [EdgeTarget("neighbor-1", 1.0)]}}, ["seed-1"])
|
|
|
|
seeds = [SeedNode("seed-1", 1.0)]
|
|
|
|
load_mock = AsyncMock(return_value={})
|
|
|
|
with patch(
|
|
"hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier",
|
|
load_mock,
|
|
):
|
|
await mpfp_traverse_async(
|
|
pool=MagicMock(),
|
|
seeds=seeds,
|
|
pattern=["semantic"],
|
|
config=config,
|
|
cache=cache,
|
|
)
|
|
|
|
# Should not call load_all_edges_for_frontier since seed-1 is already cached
|
|
load_mock.assert_not_called()
|
|
|
|
|
|
class TestMPFPGraphRetriever:
|
|
"""Tests for the MPFPGraphRetriever class."""
|
|
|
|
def test_name_is_mpfp(self):
|
|
"""Retriever name should be 'mpfp'."""
|
|
retriever = MPFPGraphRetriever()
|
|
assert retriever.name == "mpfp"
|
|
|
|
def test_default_config(self):
|
|
"""Default config should have expected patterns."""
|
|
retriever = MPFPGraphRetriever()
|
|
|
|
assert len(retriever.config.patterns_semantic) > 0
|
|
assert len(retriever.config.patterns_temporal) > 0
|
|
assert retriever.config.alpha == 0.15
|
|
assert retriever.config.top_k_neighbors == 20
|
|
|
|
def test_custom_config(self):
|
|
"""Custom config should be used."""
|
|
config = MPFPConfig(alpha=0.3, top_k_neighbors=10)
|
|
retriever = MPFPGraphRetriever(config=config)
|
|
|
|
assert retriever.config.alpha == 0.3
|
|
assert retriever.config.top_k_neighbors == 10
|
|
|
|
def test_convert_seeds_from_retrieval_results(self):
|
|
"""_convert_seeds should extract scores from RetrievalResult."""
|
|
retriever = MPFPGraphRetriever()
|
|
|
|
results = [
|
|
RetrievalResult(id="id-1", text="text1", fact_type="world", similarity=0.9),
|
|
RetrievalResult(id="id-2", text="text2", fact_type="world", similarity=0.7),
|
|
]
|
|
|
|
seeds = retriever._convert_seeds(results, "similarity")
|
|
|
|
assert len(seeds) == 2
|
|
assert seeds[0].node_id == "id-1"
|
|
assert seeds[0].score == 0.9
|
|
assert seeds[1].node_id == "id-2"
|
|
assert seeds[1].score == 0.7
|
|
|
|
def test_convert_seeds_empty(self):
|
|
"""_convert_seeds should handle empty/None input."""
|
|
retriever = MPFPGraphRetriever()
|
|
|
|
assert retriever._convert_seeds(None, "similarity") == []
|
|
assert retriever._convert_seeds([], "similarity") == []
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve_no_seeds_returns_empty(self):
|
|
"""Retrieve with no seeds should return empty results."""
|
|
retriever = MPFPGraphRetriever()
|
|
|
|
# Mock _find_semantic_seeds to return empty
|
|
with patch.object(retriever, "_find_semantic_seeds", new_callable=AsyncMock, return_value=[]):
|
|
results, timings = await retriever.retrieve(
|
|
pool=MagicMock(),
|
|
query_embedding_str="[0.1, 0.2]",
|
|
bank_id="test",
|
|
fact_type="world",
|
|
budget=10,
|
|
)
|
|
|
|
assert results == []
|
|
assert timings is not None
|
|
assert timings.pattern_count == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve_with_semantic_seeds(self):
|
|
"""Retrieve with semantic seeds should run patterns and return results."""
|
|
retriever = MPFPGraphRetriever()
|
|
|
|
semantic_seeds = [
|
|
RetrievalResult(id="seed-1", text="seed text", fact_type="world", similarity=0.9),
|
|
]
|
|
|
|
# Mock the internal functions
|
|
async def mock_traverse(*args, **kwargs):
|
|
return PatternResult(pattern=["semantic"], scores={"seed-1": 0.5, "result-1": 0.3})
|
|
|
|
async def mock_fetch(pool, node_ids, fact_type):
|
|
return [
|
|
RetrievalResult(id="seed-1", text="seed text", fact_type="world"),
|
|
RetrievalResult(id="result-1", text="result text", fact_type="world"),
|
|
]
|
|
|
|
with (
|
|
patch(
|
|
"hindsight_api.engine.search.mpfp_retrieval.mpfp_traverse_async",
|
|
side_effect=mock_traverse,
|
|
),
|
|
patch(
|
|
"hindsight_api.engine.search.mpfp_retrieval.fetch_memory_units_by_ids",
|
|
side_effect=mock_fetch,
|
|
),
|
|
):
|
|
results, timings = await retriever.retrieve(
|
|
pool=MagicMock(),
|
|
query_embedding_str="[0.1, 0.2]",
|
|
bank_id="test",
|
|
fact_type="world",
|
|
budget=10,
|
|
semantic_seeds=semantic_seeds,
|
|
)
|
|
|
|
assert len(results) == 2
|
|
assert timings is not None
|
|
assert timings.pattern_count > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mpfp_integration(memory, request_context):
|
|
"""Integration test: MPFP retrieval with real database."""
|
|
bank_id = f"test_mpfp_{datetime.now(timezone.utc).timestamp()}"
|
|
|
|
try:
|
|
# Store memories with entity relationships
|
|
await memory.retain_async(
|
|
bank_id=bank_id,
|
|
content="Alice works at TechCorp as a software engineer",
|
|
context="employee info",
|
|
request_context=request_context,
|
|
)
|
|
await memory.retain_async(
|
|
bank_id=bank_id,
|
|
content="TechCorp is located in San Francisco",
|
|
context="company info",
|
|
request_context=request_context,
|
|
)
|
|
await memory.retain_async(
|
|
bank_id=bank_id,
|
|
content="Bob is Alice's manager at TechCorp",
|
|
context="employee info",
|
|
request_context=request_context,
|
|
)
|
|
await memory.retain_async(
|
|
bank_id=bank_id,
|
|
content="San Francisco has many tech companies",
|
|
context="city info",
|
|
request_context=request_context,
|
|
)
|
|
|
|
# Query should find related facts via graph traversal
|
|
from hindsight_api.engine.memory_engine import Budget
|
|
|
|
result = await memory.recall_async(
|
|
bank_id=bank_id,
|
|
query="Tell me about Alice",
|
|
fact_type=["world"],
|
|
budget=Budget.MID,
|
|
max_tokens=2048,
|
|
request_context=request_context,
|
|
)
|
|
|
|
# Should return results
|
|
assert result.results is not None
|
|
assert len(result.results) > 0
|
|
|
|
# Should find Alice-related facts
|
|
fact_texts = [f.text for f in result.results]
|
|
alice_facts = [t for t in fact_texts if "Alice" in t or "TechCorp" in t]
|
|
assert len(alice_facts) > 0, f"Should find Alice-related facts, got: {fact_texts}"
|
|
|
|
print(f"\n✓ MPFP integration test passed! Found {len(result.results)} facts")
|
|
|
|
finally:
|
|
await memory.delete_bank(bank_id, request_context=request_context)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mpfp_lazy_loading_efficiency(memory, request_context):
|
|
"""Test that MPFP loads edges lazily, not upfront."""
|
|
bank_id = f"test_mpfp_lazy_{datetime.now(timezone.utc).timestamp()}"
|
|
|
|
try:
|
|
# Store many memories to create a larger graph
|
|
for i in range(20):
|
|
await memory.retain_async(
|
|
bank_id=bank_id,
|
|
content=f"Fact number {i} about topic {i % 5}",
|
|
context=f"context {i}",
|
|
request_context=request_context,
|
|
)
|
|
|
|
from hindsight_api.engine.memory_engine import Budget
|
|
|
|
# Query - MPFP should only load edges for relevant frontier nodes
|
|
result = await memory.recall_async(
|
|
bank_id=bank_id,
|
|
query="topic 0",
|
|
fact_type=["world"],
|
|
budget=Budget.LOW,
|
|
max_tokens=1024,
|
|
enable_trace=True,
|
|
request_context=request_context,
|
|
)
|
|
|
|
assert result.results is not None
|
|
|
|
# Check trace for timing info
|
|
if result.trace:
|
|
print(f"\n✓ MPFP lazy loading test passed!")
|
|
print(f" - Facts returned: {len(result.results)}")
|
|
|
|
finally:
|
|
await memory.delete_bank(bank_id, request_context=request_context)
|