mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
f5cea7c51e
Make the ONNX + sqlite-vec memory path truly batched. Batch ONNX embed_batch calls, batch sqlite-vec index/remove work under a single cached connection, and update MCP warm-up to use batch embed/save/index flows. Add focused regression tests for ONNX batching, sqlite-vec single-connection batch behavior, and MCP warm-up batching. Skip the MCP-specific test when optional MCP dependencies are not installed. Refs #240
525 lines
16 KiB
Python
525 lines
16 KiB
Python
"""Tests for SQLiteVectorIndex using sqlite-vec.
|
|
|
|
Tests verify:
|
|
- Vector indexing and search
|
|
- True CRUD operations (real deletes)
|
|
- Filtering by user_id, session_id, etc.
|
|
- Persistence across instances
|
|
- Memory stats and bounding
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from headroom.memory.models import Memory
|
|
from headroom.memory.ports import VectorFilter
|
|
|
|
# Check if sqlite-vec is available
|
|
try:
|
|
from headroom.memory.adapters.sqlite_vector import is_sqlite_vec_available
|
|
|
|
SQLITE_VEC_AVAILABLE = is_sqlite_vec_available()
|
|
except ImportError:
|
|
SQLITE_VEC_AVAILABLE = False
|
|
|
|
|
|
@pytest.mark.skipif(not SQLITE_VEC_AVAILABLE, reason="sqlite-vec not available")
|
|
class TestSQLiteVectorIndex:
|
|
"""Tests for SQLiteVectorIndex."""
|
|
|
|
@pytest.fixture
|
|
def index(self):
|
|
"""Create a temporary SQLite vector index."""
|
|
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
|
|
db_path = f.name
|
|
|
|
from headroom.memory.adapters.sqlite_vector import SQLiteVectorIndex
|
|
|
|
index = SQLiteVectorIndex(dimension=384, db_path=db_path)
|
|
yield index
|
|
|
|
if os.path.exists(db_path):
|
|
os.unlink(db_path)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_index_and_search(self, index):
|
|
"""Test basic indexing and search."""
|
|
np.random.seed(42)
|
|
|
|
# Create memories with random embeddings
|
|
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,
|
|
)
|
|
await index.index(memory)
|
|
memories.append(memory)
|
|
|
|
assert index.size == 10
|
|
|
|
# Search with first memory's embedding - should find itself
|
|
filter = VectorFilter(
|
|
query_vector=memories[0].embedding,
|
|
top_k=3,
|
|
user_id="alice",
|
|
)
|
|
results = await index.search(filter)
|
|
|
|
assert len(results) == 3
|
|
assert results[0].memory.id == memories[0].id
|
|
assert results[0].similarity > 0.99 # Should be ~1.0 for exact match
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_true_delete(self, index):
|
|
"""Test that delete actually removes entries."""
|
|
np.random.seed(42)
|
|
|
|
# Add memories
|
|
memories = []
|
|
for i in range(5):
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content=f"Content {i}",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
await index.index(memory)
|
|
memories.append(memory)
|
|
|
|
assert index.size == 5
|
|
|
|
# Delete one
|
|
result = await index.remove(memories[0].id)
|
|
assert result is True
|
|
assert index.size == 4
|
|
|
|
# Search should not find deleted memory
|
|
filter = VectorFilter(
|
|
query_vector=memories[0].embedding,
|
|
top_k=10,
|
|
user_id="alice",
|
|
)
|
|
results = await index.search(filter)
|
|
|
|
result_ids = {r.memory.id for r in results}
|
|
assert memories[0].id not in result_ids
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_embedding(self, index):
|
|
"""Test updating an existing entry."""
|
|
np.random.seed(42)
|
|
|
|
embedding1 = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content="Original content",
|
|
user_id="alice",
|
|
embedding=embedding1,
|
|
)
|
|
await index.index(memory)
|
|
|
|
# Update with new embedding
|
|
embedding2 = np.random.randn(384).astype(np.float32)
|
|
memory.embedding = embedding2
|
|
memory.content = "Updated content"
|
|
await index.index(memory)
|
|
|
|
# Should still be only 1 entry
|
|
assert index.size == 1
|
|
|
|
# Get stored embedding
|
|
stored = await index.get_embedding(memory.id)
|
|
assert stored is not None
|
|
np.testing.assert_array_almost_equal(stored, embedding2)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filter_by_user(self, index):
|
|
"""Test filtering search results by user_id."""
|
|
np.random.seed(42)
|
|
|
|
# Create memories for different users with similar embeddings
|
|
base_embedding = np.random.randn(384).astype(np.float32)
|
|
|
|
for user in ["alice", "bob", "charlie"]:
|
|
# Slightly perturb embedding for each user
|
|
embedding = base_embedding + np.random.randn(384).astype(np.float32) * 0.1
|
|
memory = Memory(
|
|
content=f"Content for {user}",
|
|
user_id=user,
|
|
embedding=embedding,
|
|
)
|
|
await index.index(memory)
|
|
|
|
# Search filtered by user
|
|
filter = VectorFilter(
|
|
query_vector=base_embedding,
|
|
top_k=10,
|
|
user_id="alice",
|
|
)
|
|
results = await index.search(filter)
|
|
|
|
assert len(results) == 1
|
|
assert results[0].memory.user_id == "alice"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filter_by_session(self, index):
|
|
"""Test filtering by session_id."""
|
|
np.random.seed(42)
|
|
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
|
|
# Same user, different sessions
|
|
for session in ["session1", "session2", None]:
|
|
memory = Memory(
|
|
content=f"Content for {session}",
|
|
user_id="alice",
|
|
session_id=session,
|
|
embedding=embedding + np.random.randn(384).astype(np.float32) * 0.01,
|
|
)
|
|
await index.index(memory)
|
|
|
|
filter = VectorFilter(
|
|
query_vector=embedding,
|
|
top_k=10,
|
|
user_id="alice",
|
|
session_id="session1",
|
|
)
|
|
results = await index.search(filter)
|
|
|
|
assert len(results) == 1
|
|
assert results[0].memory.session_id == "session1"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_min_similarity_filter(self, index):
|
|
"""Test minimum similarity threshold."""
|
|
np.random.seed(42)
|
|
|
|
# Create memories with varying similarity to query
|
|
query = np.random.randn(384).astype(np.float32)
|
|
query = query / np.linalg.norm(query) # Normalize
|
|
|
|
# Very similar
|
|
similar = query + np.random.randn(384).astype(np.float32) * 0.1
|
|
similar = similar / np.linalg.norm(similar)
|
|
|
|
# Less similar
|
|
less_similar = np.random.randn(384).astype(np.float32)
|
|
less_similar = less_similar / np.linalg.norm(less_similar)
|
|
|
|
await index.index(Memory(content="Similar", user_id="alice", embedding=similar))
|
|
await index.index(Memory(content="Less similar", user_id="alice", embedding=less_similar))
|
|
|
|
# High threshold should filter out dissimilar
|
|
filter = VectorFilter(
|
|
query_vector=query,
|
|
top_k=10,
|
|
min_similarity=0.8,
|
|
)
|
|
results = await index.search(filter)
|
|
|
|
# Only the similar one should pass
|
|
assert len(results) <= 1
|
|
if len(results) == 1:
|
|
assert results[0].similarity >= 0.8
|
|
|
|
|
|
@pytest.mark.skipif(not SQLITE_VEC_AVAILABLE, reason="sqlite-vec not available")
|
|
class TestSQLiteVectorIndexPersistence:
|
|
"""Tests for persistence across index instances."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_data_persists_across_instances(self):
|
|
"""Test that data survives index restart."""
|
|
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
|
|
db_path = f.name
|
|
|
|
try:
|
|
from headroom.memory.adapters.sqlite_vector import SQLiteVectorIndex
|
|
|
|
# Create index and add data
|
|
index1 = SQLiteVectorIndex(dimension=384, db_path=db_path)
|
|
|
|
np.random.seed(42)
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content="Persistent content",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
await index1.index(memory)
|
|
|
|
memory_id = memory.id
|
|
|
|
# Create new index instance
|
|
index2 = SQLiteVectorIndex(dimension=384, db_path=db_path)
|
|
|
|
assert index2.size == 1
|
|
|
|
# Should find the memory
|
|
filter = VectorFilter(
|
|
query_vector=embedding,
|
|
top_k=1,
|
|
)
|
|
results = await index2.search(filter)
|
|
|
|
assert len(results) == 1
|
|
assert results[0].memory.id == memory_id
|
|
finally:
|
|
if os.path.exists(db_path):
|
|
os.unlink(db_path)
|
|
|
|
|
|
@pytest.mark.skipif(not SQLITE_VEC_AVAILABLE, reason="sqlite-vec not available")
|
|
class TestSQLiteVectorIndexMemoryStats:
|
|
"""Tests for memory statistics."""
|
|
|
|
@pytest.fixture
|
|
def index(self):
|
|
"""Create a temporary SQLite vector index."""
|
|
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
|
|
db_path = f.name
|
|
|
|
from headroom.memory.adapters.sqlite_vector import SQLiteVectorIndex
|
|
|
|
index = SQLiteVectorIndex(dimension=384, db_path=db_path, page_cache_size_kb=4096)
|
|
yield index
|
|
|
|
if os.path.exists(db_path):
|
|
os.unlink(db_path)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_memory_stats(self, index):
|
|
"""Test memory statistics."""
|
|
stats = index.get_memory_stats()
|
|
|
|
assert stats.name == "sqlite_vector_index"
|
|
assert stats.entry_count == 0
|
|
assert stats.budget_bytes == 4096 * 1024 # 4MB cache
|
|
|
|
# Add some entries
|
|
np.random.seed(42)
|
|
for i in range(10):
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content=f"Content {i}",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
await index.index(memory)
|
|
|
|
stats = index.get_memory_stats()
|
|
assert stats.entry_count == 10
|
|
assert stats.size_bytes > 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stats(self, index):
|
|
"""Test index statistics."""
|
|
np.random.seed(42)
|
|
|
|
for i in range(5):
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content=f"Content {i}",
|
|
user_id="alice" if i < 3 else "bob",
|
|
embedding=embedding,
|
|
)
|
|
await index.index(memory)
|
|
|
|
stats = index.stats()
|
|
|
|
assert stats["size"] == 5
|
|
assert stats["dimension"] == 384
|
|
assert stats["users"] == 2
|
|
assert stats["page_cache_size_kb"] == 4096
|
|
assert stats["db_size_bytes"] > 0
|
|
|
|
|
|
@pytest.mark.skipif(not SQLITE_VEC_AVAILABLE, reason="sqlite-vec not available")
|
|
class TestSQLiteVectorIndexEdgeCases:
|
|
"""Tests for edge cases."""
|
|
|
|
@pytest.fixture
|
|
def index(self):
|
|
"""Create a temporary SQLite vector index."""
|
|
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
|
|
db_path = f.name
|
|
|
|
from headroom.memory.adapters.sqlite_vector import SQLiteVectorIndex
|
|
|
|
index = SQLiteVectorIndex(dimension=384, db_path=db_path)
|
|
yield index
|
|
|
|
if os.path.exists(db_path):
|
|
os.unlink(db_path)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_empty_index(self, index):
|
|
"""Test searching an empty index."""
|
|
np.random.seed(42)
|
|
query = np.random.randn(384).astype(np.float32)
|
|
|
|
filter = VectorFilter(query_vector=query, top_k=10)
|
|
results = await index.search(filter)
|
|
|
|
assert len(results) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_remove_nonexistent(self, index):
|
|
"""Test removing a nonexistent entry."""
|
|
result = await index.remove("nonexistent-id")
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wrong_dimension_raises(self, index):
|
|
"""Test that wrong embedding dimension raises error."""
|
|
wrong_embedding = np.random.randn(128).astype(np.float32) # Wrong dimension
|
|
memory = Memory(
|
|
content="Test",
|
|
user_id="alice",
|
|
embedding=wrong_embedding,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="dimension"):
|
|
await index.index(memory)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_embedding_raises(self, index):
|
|
"""Test that missing embedding raises error."""
|
|
memory = Memory(
|
|
content="Test",
|
|
user_id="alice",
|
|
embedding=None,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="no embedding"):
|
|
await index.index(memory)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_clear(self, index):
|
|
"""Test clearing all entries."""
|
|
np.random.seed(42)
|
|
|
|
for i in range(5):
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content=f"Content {i}",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
await index.index(memory)
|
|
|
|
assert index.size == 5
|
|
|
|
index.clear()
|
|
|
|
assert index.size == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_index(self, index):
|
|
"""Test batch indexing."""
|
|
np.random.seed(42)
|
|
|
|
memories = []
|
|
for i in range(10):
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content=f"Content {i}",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
memories.append(memory)
|
|
|
|
# Add one without embedding
|
|
memories.append(Memory(content="No embedding", user_id="alice"))
|
|
|
|
indexed = await index.index_batch(memories)
|
|
|
|
assert indexed == 10
|
|
assert index.size == 10
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_index_uses_single_connection(self, index, monkeypatch):
|
|
"""Test batch indexing reuses a single sqlite-vec connection."""
|
|
np.random.seed(42)
|
|
memories = [
|
|
Memory(
|
|
content=f"Content {i}",
|
|
user_id="alice",
|
|
embedding=np.random.randn(384).astype(np.float32),
|
|
)
|
|
for i in range(10)
|
|
]
|
|
|
|
original_get_conn = index._get_conn
|
|
conn_calls = 0
|
|
|
|
def counting_get_conn():
|
|
nonlocal conn_calls
|
|
conn_calls += 1
|
|
return original_get_conn()
|
|
|
|
monkeypatch.setattr(index, "_get_conn", counting_get_conn)
|
|
|
|
indexed = await index.index_batch(memories)
|
|
|
|
assert indexed == 10
|
|
assert conn_calls == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_remove(self, index):
|
|
"""Test batch removal."""
|
|
np.random.seed(42)
|
|
|
|
memories = []
|
|
for i in range(5):
|
|
embedding = np.random.randn(384).astype(np.float32)
|
|
memory = Memory(
|
|
content=f"Content {i}",
|
|
user_id="alice",
|
|
embedding=embedding,
|
|
)
|
|
await index.index(memory)
|
|
memories.append(memory)
|
|
|
|
# Remove some
|
|
ids_to_remove = [memories[0].id, memories[2].id, "nonexistent"]
|
|
removed = await index.remove_batch(ids_to_remove)
|
|
|
|
assert removed == 2
|
|
assert index.size == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_remove_uses_single_connection(self, index, monkeypatch):
|
|
"""Test batch removal reuses a single sqlite-vec connection."""
|
|
np.random.seed(42)
|
|
memories = []
|
|
for i in range(5):
|
|
memory = Memory(
|
|
content=f"Content {i}",
|
|
user_id="alice",
|
|
embedding=np.random.randn(384).astype(np.float32),
|
|
)
|
|
await index.index(memory)
|
|
memories.append(memory)
|
|
|
|
original_get_conn = index._get_conn
|
|
conn_calls = 0
|
|
|
|
def counting_get_conn():
|
|
nonlocal conn_calls
|
|
conn_calls += 1
|
|
return original_get_conn()
|
|
|
|
monkeypatch.setattr(index, "_get_conn", counting_get_conn)
|
|
|
|
removed = await index.remove_batch([memories[0].id, memories[2].id, "nonexistent"])
|
|
|
|
assert removed == 2
|
|
assert conn_calls == 1
|