chore: import upstream snapshot with attribution

tests/test_sqlite_vector_index.py

@@ -0,0 +1,525 @@
+"""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