mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
da74341858
Implement comprehensive memory system supporting: - Local backend (SQLite + FTS5 + HNSW) for zero-dependency operation - Mem0 backends (Neo4j + Qdrant) for production graph memory - DirectMem0Adapter for optimized pre-extracted data (bypasses LLM) - Memory extraction with facts, entities, and relationships - Proxy integration with --memory flag for automatic memory injection Key components: - headroom/memory/backends/: LocalBackend, Mem0Backend, DirectMem0Adapter - headroom/memory/system.py: MemorySystem with tool-based interface - headroom/memory/extraction.py: Entity and relationship extraction - headroom/proxy/memory_handler.py: Proxy integration layer - headroom/prediction/feature_extractor.py: Content analysis features Testing: - 217 new memory system tests covering all backends - LoCoMo evaluation framework for memory quality assessment - Integration tests for proxy memory functionality Also removes deprecated example files in favor of focused test coverage.
804 lines
28 KiB
Python
804 lines
28 KiB
Python
"""Tests for the hierarchical memory system.
|
|
|
|
Tests cover:
|
|
- Memory models (Memory, ScopeLevel)
|
|
- SQLite memory store
|
|
- HNSW vector index
|
|
- FTS5 text index
|
|
- LRU cache
|
|
- HierarchicalMemory orchestrator
|
|
- Memory bubbling
|
|
- Temporal versioning (supersession)
|
|
"""
|
|
|
|
# CRITICAL: Must set TOKENIZERS_PARALLELISM before any imports that might
|
|
# trigger sentence_transformers/transformers loading. The Rust tokenizers
|
|
# use parallelism that conflicts with Python's forking model, causing
|
|
# deadlocks when combined with asyncio/pytest.
|
|
# See: https://github.com/huggingface/transformers/issues/5486
|
|
import os
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
import asyncio
|
|
import tempfile
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from headroom.memory.adapters.cache import LRUMemoryCache
|
|
from headroom.memory.adapters.fts5 import FTS5TextIndex
|
|
from headroom.memory.adapters.sqlite import SQLiteMemoryStore
|
|
from headroom.memory.models import Memory, ScopeLevel
|
|
from headroom.memory.ports import MemoryFilter, TextFilter, VectorFilter
|
|
|
|
# =============================================================================
|
|
# Fixtures
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_db_path():
|
|
"""Create a temporary database path."""
|
|
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
|
|
yield Path(f.name)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_memory():
|
|
"""Create a sample memory for testing."""
|
|
return Memory(
|
|
content="User prefers Python over JavaScript",
|
|
user_id="alice",
|
|
session_id="session-123",
|
|
importance=0.8,
|
|
entity_refs=["Python", "JavaScript"],
|
|
metadata={"source": "conversation"},
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_embedding():
|
|
"""Create a sample embedding vector."""
|
|
return np.random.randn(384).astype(np.float32)
|
|
|
|
|
|
# =============================================================================
|
|
# Memory Model Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestMemoryModel:
|
|
"""Tests for the Memory dataclass."""
|
|
|
|
def test_memory_creation(self):
|
|
"""Test basic memory creation."""
|
|
memory = Memory(
|
|
content="Test content",
|
|
user_id="test-user",
|
|
)
|
|
assert memory.content == "Test content"
|
|
assert memory.user_id == "test-user"
|
|
assert memory.id is not None # Auto-generated UUID
|
|
assert memory.importance == 0.5 # Default
|
|
|
|
def test_scope_level_computation(self):
|
|
"""Test scope level is correctly computed from hierarchy fields."""
|
|
# USER level - only user_id
|
|
user_mem = Memory(content="test", user_id="alice")
|
|
assert user_mem.scope_level == ScopeLevel.USER
|
|
|
|
# SESSION level - user_id + session_id
|
|
session_mem = Memory(content="test", user_id="alice", session_id="sess-1")
|
|
assert session_mem.scope_level == ScopeLevel.SESSION
|
|
|
|
# AGENT level - user_id + session_id + agent_id
|
|
agent_mem = Memory(content="test", user_id="alice", session_id="sess-1", agent_id="agent-1")
|
|
assert agent_mem.scope_level == ScopeLevel.AGENT
|
|
|
|
# TURN level - all four
|
|
turn_mem = Memory(
|
|
content="test",
|
|
user_id="alice",
|
|
session_id="sess-1",
|
|
agent_id="agent-1",
|
|
turn_id="turn-1",
|
|
)
|
|
assert turn_mem.scope_level == ScopeLevel.TURN
|
|
|
|
def test_is_current_property(self):
|
|
"""Test is_current property for supersession detection."""
|
|
current = Memory(content="test", user_id="alice")
|
|
assert current.is_current is True
|
|
|
|
superseded = Memory(content="test", user_id="alice", valid_until=datetime.utcnow())
|
|
assert superseded.is_current is False
|
|
|
|
def test_memory_serialization(self, sample_embedding):
|
|
"""Test Memory to_dict and from_dict."""
|
|
memory = Memory(
|
|
content="Test content",
|
|
user_id="alice",
|
|
session_id="sess-1",
|
|
importance=0.9,
|
|
entity_refs=["entity1"],
|
|
metadata={"key": "value"},
|
|
embedding=sample_embedding,
|
|
)
|
|
|
|
# Serialize
|
|
data = memory.to_dict()
|
|
assert data["content"] == "Test content"
|
|
assert data["user_id"] == "alice"
|
|
assert data["embedding"] is not None
|
|
|
|
# Deserialize
|
|
restored = Memory.from_dict(data)
|
|
assert restored.content == memory.content
|
|
assert restored.user_id == memory.user_id
|
|
assert restored.importance == memory.importance
|
|
assert np.allclose(restored.embedding, memory.embedding)
|
|
|
|
|
|
# =============================================================================
|
|
# SQLite Store Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestSQLiteMemoryStore:
|
|
"""Tests for SQLiteMemoryStore."""
|
|
|
|
@pytest.fixture
|
|
def store(self, temp_db_path):
|
|
"""Create a SQLite store for testing."""
|
|
return SQLiteMemoryStore(temp_db_path)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_save_and_get(self, store, sample_memory):
|
|
"""Test saving and retrieving a memory."""
|
|
await store.save(sample_memory)
|
|
|
|
retrieved = await store.get(sample_memory.id)
|
|
assert retrieved is not None
|
|
assert retrieved.id == sample_memory.id
|
|
assert retrieved.content == sample_memory.content
|
|
assert retrieved.user_id == sample_memory.user_id
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_save_batch(self, store):
|
|
"""Test batch saving memories."""
|
|
memories = [Memory(content=f"Memory {i}", user_id="alice") for i in range(10)]
|
|
|
|
await store.save_batch(memories)
|
|
|
|
for memory in memories:
|
|
retrieved = await store.get(memory.id)
|
|
assert retrieved is not None
|
|
assert retrieved.content == memory.content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete(self, store, sample_memory):
|
|
"""Test deleting a memory."""
|
|
await store.save(sample_memory)
|
|
|
|
deleted = await store.delete(sample_memory.id)
|
|
assert deleted is True
|
|
|
|
retrieved = await store.get(sample_memory.id)
|
|
assert retrieved is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_by_user(self, store):
|
|
"""Test querying memories by user_id."""
|
|
# Create memories for different users
|
|
alice_memories = [Memory(content=f"Alice {i}", user_id="alice") for i in range(5)]
|
|
bob_memories = [Memory(content=f"Bob {i}", user_id="bob") for i in range(3)]
|
|
|
|
await store.save_batch(alice_memories + bob_memories)
|
|
|
|
# Query Alice's memories
|
|
results = await store.query(MemoryFilter(user_id="alice"))
|
|
assert len(results) == 5
|
|
|
|
# Query Bob's memories
|
|
results = await store.query(MemoryFilter(user_id="bob"))
|
|
assert len(results) == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_by_importance_range(self, store):
|
|
"""Test querying memories by importance range."""
|
|
memories = [
|
|
Memory(content="Low importance", user_id="alice", importance=0.2),
|
|
Memory(content="Medium importance", user_id="alice", importance=0.5),
|
|
Memory(content="High importance", user_id="alice", importance=0.9),
|
|
]
|
|
|
|
await store.save_batch(memories)
|
|
|
|
# Query high importance only
|
|
results = await store.query(MemoryFilter(user_id="alice", min_importance=0.8))
|
|
assert len(results) == 1
|
|
assert results[0].content == "High importance"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_by_importance(self, store):
|
|
"""Test querying memories by importance range."""
|
|
memories = [
|
|
Memory(content="Low", user_id="alice", importance=0.3),
|
|
Memory(content="Medium", user_id="alice", importance=0.5),
|
|
Memory(content="High", user_id="alice", importance=0.9),
|
|
]
|
|
|
|
await store.save_batch(memories)
|
|
|
|
# Query high importance only
|
|
results = await store.query(MemoryFilter(user_id="alice", min_importance=0.8))
|
|
assert len(results) == 1
|
|
assert results[0].content == "High"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_by_scope_level(self, store):
|
|
"""Test querying by explicit scope level."""
|
|
memories = [
|
|
Memory(content="User level", user_id="alice"),
|
|
Memory(content="Session level", user_id="alice", session_id="sess-1"),
|
|
Memory(content="Agent level", user_id="alice", session_id="sess-1", agent_id="agent-1"),
|
|
]
|
|
|
|
await store.save_batch(memories)
|
|
|
|
# Query only USER level
|
|
results = await store.query(MemoryFilter(user_id="alice", scope_levels=[ScopeLevel.USER]))
|
|
assert len(results) == 1
|
|
assert results[0].content == "User level"
|
|
|
|
# Query SESSION level
|
|
results = await store.query(
|
|
MemoryFilter(user_id="alice", scope_levels=[ScopeLevel.SESSION])
|
|
)
|
|
assert len(results) == 1
|
|
assert results[0].content == "Session level"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_supersession(self, store):
|
|
"""Test memory supersession."""
|
|
original = Memory(
|
|
content="User prefers Python",
|
|
user_id="alice",
|
|
)
|
|
await store.save(original)
|
|
|
|
# Supersede with new preference
|
|
new_memory = Memory(
|
|
content="User now prefers Rust",
|
|
user_id="alice",
|
|
)
|
|
|
|
superseded = await store.supersede(original.id, new_memory)
|
|
|
|
# New memory should be linked to old
|
|
assert superseded.supersedes == original.id
|
|
|
|
# Old memory should be marked as superseded
|
|
old_retrieved = await store.get(original.id)
|
|
assert old_retrieved.superseded_by == superseded.id
|
|
assert old_retrieved.valid_until is not None
|
|
assert old_retrieved.is_current is False
|
|
|
|
# New memory should be current
|
|
assert superseded.is_current is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_history(self, store):
|
|
"""Test getting supersession chain history."""
|
|
# Create a chain: v1 -> v2 -> v3
|
|
v1 = Memory(content="Version 1", user_id="alice")
|
|
await store.save(v1)
|
|
|
|
v2 = Memory(content="Version 2", user_id="alice")
|
|
v2 = await store.supersede(v1.id, v2)
|
|
|
|
v3 = Memory(content="Version 3", user_id="alice")
|
|
v3 = await store.supersede(v2.id, v3)
|
|
|
|
# Get history from middle
|
|
history = await store.get_history(v2.id, include_future=True)
|
|
assert len(history) == 3
|
|
assert history[0].content == "Version 1"
|
|
assert history[1].content == "Version 2"
|
|
assert history[2].content == "Version 3"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_clear_scope(self, store):
|
|
"""Test clearing memories at a scope level."""
|
|
# Create memories at different scopes
|
|
memories = [
|
|
Memory(content="User 1", user_id="alice"),
|
|
Memory(content="User 2", user_id="alice"),
|
|
Memory(content="Session 1", user_id="alice", session_id="sess-1"),
|
|
Memory(content="Other user", user_id="bob"),
|
|
]
|
|
await store.save_batch(memories)
|
|
|
|
# Clear Alice's session
|
|
deleted = await store.clear_scope("alice", session_id="sess-1")
|
|
assert deleted == 1
|
|
|
|
# Alice's user-level memories should remain
|
|
remaining = await store.query(MemoryFilter(user_id="alice"))
|
|
assert len(remaining) == 2
|
|
|
|
|
|
# =============================================================================
|
|
# LRU Cache Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestLRUMemoryCache:
|
|
"""Tests for LRUMemoryCache."""
|
|
|
|
@pytest.fixture
|
|
def cache(self):
|
|
"""Create a cache for testing."""
|
|
return LRUMemoryCache(max_size=5)
|
|
|
|
async def test_set_and_get(self, cache, sample_memory):
|
|
"""Test basic cache put and get."""
|
|
await cache.put(sample_memory)
|
|
|
|
retrieved = await cache.get(sample_memory.id)
|
|
assert retrieved is not None
|
|
assert retrieved.id == sample_memory.id
|
|
|
|
async def test_lru_eviction(self, cache):
|
|
"""Test LRU eviction when cache is full."""
|
|
# Fill cache with 5 memories
|
|
memories = [Memory(content=f"Mem {i}", user_id="alice") for i in range(5)]
|
|
for m in memories:
|
|
await cache.put(m)
|
|
|
|
assert cache.size == 5
|
|
|
|
# Add one more - should evict the first
|
|
new_mem = Memory(content="New", user_id="alice")
|
|
await cache.put(new_mem)
|
|
|
|
assert cache.size == 5
|
|
assert await cache.get(memories[0].id) is None # First was evicted
|
|
assert await cache.get(new_mem.id) is not None
|
|
|
|
async def test_access_updates_lru_order(self, cache):
|
|
"""Test that accessing a key moves it to end of LRU."""
|
|
memories = [Memory(content=f"Mem {i}", user_id="alice") for i in range(5)]
|
|
for m in memories:
|
|
await cache.put(m)
|
|
|
|
# Access the first memory (makes it most recently used)
|
|
await cache.get(memories[0].id)
|
|
|
|
# Add new memory - should evict second (now oldest)
|
|
new_mem = Memory(content="New", user_id="alice")
|
|
await cache.put(new_mem)
|
|
|
|
assert await cache.get(memories[0].id) is not None # Still present
|
|
assert await cache.get(memories[1].id) is None # Evicted
|
|
|
|
async def test_delete(self, cache, sample_memory):
|
|
"""Test deleting from cache."""
|
|
await cache.put(sample_memory)
|
|
assert cache.size == 1
|
|
|
|
deleted = await cache.invalidate(sample_memory.id)
|
|
assert deleted is True
|
|
assert cache.size == 0
|
|
assert await cache.get(sample_memory.id) is None
|
|
|
|
async def test_clear(self, cache):
|
|
"""Test clearing the cache."""
|
|
memories = [Memory(content=f"Mem {i}", user_id="alice") for i in range(3)]
|
|
for m in memories:
|
|
await cache.put(m)
|
|
|
|
await cache.clear()
|
|
assert cache.size == 0
|
|
|
|
|
|
# =============================================================================
|
|
# FTS5 Text Index Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestFTS5TextIndex:
|
|
"""Tests for FTS5TextIndex."""
|
|
|
|
@pytest.fixture
|
|
def text_index(self, temp_db_path):
|
|
"""Create a FTS5 text index for testing."""
|
|
return FTS5TextIndex(temp_db_path)
|
|
|
|
def test_index_and_search(self, text_index):
|
|
"""Test indexing and searching text."""
|
|
# Index some memories
|
|
text_index.index("mem-1", "User prefers Python programming", {"user_id": "alice"})
|
|
text_index.index("mem-2", "JavaScript is also popular", {"user_id": "alice"})
|
|
text_index.index("mem-3", "Python is great for data science", {"user_id": "alice"})
|
|
|
|
# Search for Python
|
|
results = text_index.search("Python", k=10)
|
|
assert len(results) == 2
|
|
|
|
# Results should include memory IDs
|
|
result_ids = [r.memory_id for r in results]
|
|
assert "mem-1" in result_ids
|
|
assert "mem-3" in result_ids
|
|
|
|
def test_search_with_user_filter(self, text_index):
|
|
"""Test searching with user filter."""
|
|
text_index.index("mem-1", "Python programming", {"user_id": "alice"})
|
|
text_index.index("mem-2", "Python scripting", {"user_id": "bob"})
|
|
|
|
# Search only Alice's memories
|
|
filter = TextFilter(user_id="alice")
|
|
results = text_index.search("Python", k=10, filter=filter)
|
|
|
|
assert len(results) == 1
|
|
assert results[0].memory_id == "mem-1"
|
|
|
|
def test_search_with_session_filter(self, text_index):
|
|
"""Test searching with session filter."""
|
|
text_index.index("mem-1", "Prefers Python", {"user_id": "alice", "session_id": "sess-1"})
|
|
text_index.index(
|
|
"mem-2", "Python is installed", {"user_id": "alice", "session_id": "sess-2"}
|
|
)
|
|
|
|
# Search only session-1
|
|
filter = TextFilter(user_id="alice", session_id="sess-1")
|
|
results = text_index.search("Python", k=10, filter=filter)
|
|
|
|
assert len(results) == 1
|
|
assert results[0].memory_id == "mem-1"
|
|
|
|
def test_delete(self, text_index):
|
|
"""Test deleting from text index."""
|
|
text_index.index("mem-1", "Test content", {"user_id": "alice"})
|
|
|
|
deleted = text_index.delete("mem-1")
|
|
assert deleted is True
|
|
|
|
results = text_index.search("Test", k=10)
|
|
assert len(results) == 0
|
|
|
|
def test_batch_index(self, text_index):
|
|
"""Test batch indexing."""
|
|
memory_ids = ["mem-1", "mem-2", "mem-3"]
|
|
texts = ["Python code", "JavaScript code", "Rust code"]
|
|
metadata = [{"user_id": "alice"} for _ in range(3)]
|
|
|
|
text_index.index_batch(memory_ids, texts, metadata)
|
|
|
|
assert text_index.count() == 3
|
|
|
|
|
|
# =============================================================================
|
|
# Memory Config Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestMemoryConfig:
|
|
"""Tests for MemoryConfig validation."""
|
|
|
|
def test_default_config(self):
|
|
"""Test default configuration."""
|
|
from headroom.memory.config import MemoryConfig
|
|
|
|
config = MemoryConfig()
|
|
assert config.vector_dimension == 384
|
|
assert config.cache_enabled is True
|
|
assert config.auto_bubble is True
|
|
|
|
def test_invalid_dimension(self):
|
|
"""Test that invalid dimension raises error."""
|
|
from headroom.memory.config import MemoryConfig
|
|
|
|
with pytest.raises(ValueError):
|
|
MemoryConfig(vector_dimension=0)
|
|
|
|
def test_openai_requires_api_key(self):
|
|
"""Test that OpenAI backend requires API key."""
|
|
from headroom.memory.config import EmbedderBackend, MemoryConfig
|
|
|
|
with pytest.raises(ValueError, match="openai_api_key"):
|
|
MemoryConfig(embedder_backend=EmbedderBackend.OPENAI)
|
|
|
|
|
|
# =============================================================================
|
|
# Integration Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestIntegration:
|
|
"""Integration tests that test multiple components together."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_with_embeddings(self, temp_db_path, sample_embedding):
|
|
"""Test storing and retrieving memories with embeddings."""
|
|
store = SQLiteMemoryStore(temp_db_path)
|
|
|
|
memory = Memory(
|
|
content="Test content",
|
|
user_id="alice",
|
|
embedding=sample_embedding,
|
|
)
|
|
|
|
await store.save(memory)
|
|
|
|
retrieved = await store.get(memory.id)
|
|
assert retrieved.embedding is not None
|
|
assert np.allclose(retrieved.embedding, sample_embedding)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_temporal_query(self, temp_db_path):
|
|
"""Test point-in-time temporal queries."""
|
|
store = SQLiteMemoryStore(temp_db_path)
|
|
|
|
# Create a supersession chain
|
|
original = Memory(content="Original preference", user_id="alice")
|
|
await store.save(original)
|
|
|
|
# Capture time after original was created (valid_from is set at Memory creation)
|
|
time_when_original_valid = original.valid_from + timedelta(milliseconds=1)
|
|
|
|
# Wait a bit for time difference
|
|
await asyncio.sleep(0.01)
|
|
|
|
# Supersede
|
|
new_memory = Memory(content="New preference", user_id="alice")
|
|
supersede_time = datetime.utcnow()
|
|
await store.supersede(original.id, new_memory, supersede_time)
|
|
|
|
# Query at a point when original was valid (after its valid_from, before supersession)
|
|
# The past_time must be >= original.valid_from and < supersede_time
|
|
results = await store.query(
|
|
MemoryFilter(
|
|
user_id="alice", valid_at=time_when_original_valid, include_superseded=True
|
|
)
|
|
)
|
|
assert len(results) == 1
|
|
assert results[0].content == "Original preference"
|
|
|
|
# Query current - should return new
|
|
results = await store.query(MemoryFilter(user_id="alice"))
|
|
assert len(results) == 1
|
|
assert results[0].content == "New preference"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hierarchical_scope_query(self, temp_db_path):
|
|
"""Test hierarchical scope filtering."""
|
|
store = SQLiteMemoryStore(temp_db_path)
|
|
|
|
# Create memories at different scopes
|
|
user_mem = Memory(content="User pref", user_id="alice")
|
|
session_mem = Memory(content="Session context", user_id="alice", session_id="sess-1")
|
|
agent_mem = Memory(
|
|
content="Agent decision",
|
|
user_id="alice",
|
|
session_id="sess-1",
|
|
agent_id="agent-1",
|
|
)
|
|
|
|
await store.save_batch([user_mem, session_mem, agent_mem])
|
|
|
|
# Query user scope only - should get just user_mem
|
|
user_only = await store.query(MemoryFilter(user_id="alice", scope_levels=[ScopeLevel.USER]))
|
|
assert len(user_only) == 1
|
|
assert user_only[0].content == "User pref"
|
|
|
|
# Query all scopes for this user
|
|
all_memories = await store.query(MemoryFilter(user_id="alice"))
|
|
assert len(all_memories) == 3
|
|
|
|
# Query specific session
|
|
session_memories = await store.query(MemoryFilter(user_id="alice", session_id="sess-1"))
|
|
assert len(session_memories) == 2 # session and agent level
|
|
|
|
|
|
# =============================================================================
|
|
# HNSW Vector Index Tests
|
|
# =============================================================================
|
|
|
|
# Check if hnswlib is available
|
|
try:
|
|
from headroom.memory.adapters.hnsw import HNSW_AVAILABLE
|
|
except ImportError:
|
|
HNSW_AVAILABLE = False
|
|
|
|
|
|
@pytest.mark.skipif(not HNSW_AVAILABLE, reason="hnswlib not installed")
|
|
class TestHNSWVectorIndex:
|
|
"""Tests for HNSWVectorIndex."""
|
|
|
|
@pytest.fixture
|
|
def vector_index(self, temp_db_path):
|
|
"""Create an HNSW vector index for testing."""
|
|
from headroom.memory.adapters.hnsw import HNSWVectorIndex
|
|
|
|
return HNSWVectorIndex(dimension=384, save_path=temp_db_path.with_suffix(".hnsw"))
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_index_and_search(self, vector_index):
|
|
"""Test indexing and searching vectors."""
|
|
|
|
# Create memories with random embeddings
|
|
np.random.seed(42)
|
|
memories = []
|
|
for i in range(10):
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content=f"Test content {i}",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
memories.append(memory)
|
|
|
|
# Index all memories
|
|
for memory in memories:
|
|
await vector_index.index(memory)
|
|
|
|
# Search with first memory's embedding - should find itself as most similar
|
|
filter = VectorFilter(
|
|
query_vector=memories[0].embedding,
|
|
top_k=3,
|
|
user_id="alice",
|
|
)
|
|
results = await vector_index.search(filter)
|
|
assert len(results) == 3
|
|
assert results[0].memory.id == memories[0].id
|
|
assert results[0].similarity > 0.99 # Should be very close to 1.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_index(self, vector_index):
|
|
"""Test batch indexing."""
|
|
|
|
np.random.seed(42)
|
|
memories = []
|
|
for i in range(100):
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content=f"Test content {i}",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
memories.append(memory)
|
|
|
|
count = await vector_index.index_batch(memories)
|
|
|
|
# Verify count
|
|
assert count == 100
|
|
assert vector_index.size == 100
|
|
|
|
# Search should work
|
|
filter = VectorFilter(
|
|
query_vector=memories[50].embedding,
|
|
top_k=5,
|
|
user_id="alice",
|
|
)
|
|
results = await vector_index.search(filter)
|
|
assert len(results) == 5
|
|
assert results[0].memory.id == memories[50].id
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_remove(self, vector_index):
|
|
"""Test removing from index."""
|
|
np.random.seed(42)
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content="Test content",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
await vector_index.index(memory)
|
|
|
|
# HNSW doesn't support true deletion, but marks as deleted
|
|
removed = await vector_index.remove(memory.id)
|
|
assert removed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persistence(self, temp_db_path):
|
|
"""Test that index persists to disk."""
|
|
from headroom.memory.adapters.hnsw import HNSWVectorIndex
|
|
|
|
save_path = temp_db_path.with_suffix(".hnsw")
|
|
np.random.seed(42)
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content="Test content",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
|
|
# Create and populate index
|
|
index1 = HNSWVectorIndex(dimension=384, save_path=save_path)
|
|
await index1.index(memory)
|
|
index1.save_index(save_path)
|
|
|
|
# Create new index and load from same path
|
|
index2 = HNSWVectorIndex(dimension=384, save_path=save_path)
|
|
index2.load_index(save_path)
|
|
assert index2.size == 1
|
|
|
|
filter = VectorFilter(
|
|
query_vector=embedding,
|
|
top_k=1,
|
|
user_id="alice",
|
|
)
|
|
results = await index2.search(filter)
|
|
assert results[0].memory.id == memory.id
|
|
|
|
|
|
# =============================================================================
|
|
# LocalEmbedder Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestLocalEmbedder:
|
|
"""Tests for LocalEmbedder (sentence-transformers)."""
|
|
|
|
@pytest.fixture
|
|
def embedder(self):
|
|
"""Create a local embedder for testing."""
|
|
from headroom.memory.adapters.embedders import LocalEmbedder
|
|
|
|
return LocalEmbedder()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embed_single(self, embedder):
|
|
"""Test embedding a single text."""
|
|
text = "User prefers Python programming"
|
|
embedding = await embedder.embed(text)
|
|
|
|
assert embedding is not None
|
|
assert embedding.shape == (384,)
|
|
assert embedding.dtype == np.float32
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embed_batch(self, embedder):
|
|
"""Test embedding multiple texts."""
|
|
texts = [
|
|
"Python programming",
|
|
"JavaScript development",
|
|
"Rust systems programming",
|
|
]
|
|
embeddings = await embedder.embed_batch(texts)
|
|
|
|
assert len(embeddings) == 3
|
|
for emb in embeddings:
|
|
assert emb.shape == (384,)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_similar_texts_have_high_similarity(self, embedder):
|
|
"""Test that semantically similar texts have similar embeddings."""
|
|
text1 = "The user prefers Python for data analysis"
|
|
text2 = "Python is the user's preferred language for data science"
|
|
text3 = "The weather is sunny today"
|
|
|
|
emb1 = await embedder.embed(text1)
|
|
emb2 = await embedder.embed(text2)
|
|
emb3 = await embedder.embed(text3)
|
|
|
|
# Cosine similarity
|
|
def cosine_sim(a, b):
|
|
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
|
|
|
|
# Similar texts should have high similarity
|
|
sim_related = cosine_sim(emb1, emb2)
|
|
sim_unrelated = cosine_sim(emb1, emb3)
|
|
|
|
assert sim_related > 0.7 # Related texts
|
|
assert sim_unrelated < 0.5 # Unrelated texts
|
|
assert sim_related > sim_unrelated
|
|
|
|
def test_dimension_property(self, embedder):
|
|
"""Test that dimension property returns correct value."""
|
|
assert embedder.dimension == 384
|