headroomlabs-ai--headroom/tests/test_sqlite_vector_index.py

"""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