headroom/tests/test_hnsw_only.py

"""Isolated HNSW tests - copy of relevant parts from test_hierarchical.py."""

import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"

import tempfile
from pathlib import Path

import numpy as np
import pytest

from headroom.memory.models import Memory
from headroom.memory.ports import VectorFilter

# Check if hnswlib is available (use lazy check to avoid SIGILL on incompatible CPUs)
try:
    from headroom.memory.adapters.hnsw import _check_hnswlib_available

    HNSW_AVAILABLE = _check_hnswlib_available()
except ImportError:
    HNSW_AVAILABLE = False


@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.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."""
        print("\n[TEST] Starting test_index_and_search")

        # 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)
        print(f"[TEST] Created {len(memories)} memories")

        # Index all memories
        print("[TEST] Indexing...")
        for memory in memories:
            await vector_index.index(memory)
        print("[TEST] All indexed!")

        # Search with first memory's embedding
        filter = VectorFilter(
            query_vector=memories[0].embedding,
            top_k=3,
            user_id="alice",
        )
        print("[TEST] Searching...")
        results = await vector_index.search(filter)
        print(f"[TEST] Found {len(results)} results")

        assert len(results) == 3
        assert results[0].memory.id == memories[0].id
        assert results[0].similarity > 0.99
        print("[TEST] PASSED!")

    @pytest.mark.asyncio
    async def test_bounded_index_eviction(self, temp_db_path):
        """Test that bounded index evicts low-importance entries."""
        from headroom.memory.adapters.hnsw import HNSWVectorIndex

        # Create bounded index with max 5 entries
        index = HNSWVectorIndex(
            dimension=384,
            max_entries=5,
            eviction_batch_size=2,
        )

        np.random.seed(42)

        # Add 5 memories with varying importance
        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,
                importance=0.1 * (i + 1),  # 0.1, 0.2, 0.3, 0.4, 0.5
            )
            await index.index(memory)
            memories.append(memory)

        assert index.size == 5

        # Add one more - should trigger eviction of lowest importance
        new_embedding = np.random.randn(384).astype(np.float32)
        new_memory = Memory(
            content="New high importance",
            user_id="alice",
            embedding=new_embedding,
            importance=0.9,
        )
        await index.index(new_memory)

        # Should have evicted 2 entries (eviction_batch_size) then added 1
        # So size should be 5 - 2 + 1 = 4
        assert index.size == 4

        # The lowest importance entries (0.1, 0.2) should be gone
        stats = index.get_memory_stats()
        assert stats.evictions == 2

        # Search should not find the evicted memories
        filter = VectorFilter(
            query_vector=memories[0].embedding,  # Lowest importance, should be evicted
            top_k=10,
            user_id="alice",
        )
        results = await index.search(filter)

        # memories[0] and memories[1] should be evicted
        result_ids = {r.memory.id for r in results}
        assert memories[0].id not in result_ids
        assert memories[1].id not in result_ids

    @pytest.mark.asyncio
    async def test_bounded_index_stats(self, temp_db_path):
        """Test that bounded index reports correct stats."""
        from headroom.memory.adapters.hnsw import HNSWVectorIndex

        index = HNSWVectorIndex(
            dimension=384,
            max_entries=100,
        )

        stats = index.get_memory_stats()
        assert stats.name == "vector_index"
        assert stats.entry_count == 0
        assert stats.budget_bytes is not None  # Should have budget when max_entries set
        assert stats.evictions == 0

        # 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_unbounded_index_no_eviction(self, temp_db_path):
        """Test that unbounded index doesn't evict."""
        from headroom.memory.adapters.hnsw import HNSWVectorIndex

        # Create unbounded index (max_entries=None)
        index = HNSWVectorIndex(dimension=384)

        np.random.seed(42)

        # Add many memories
        for i in range(20):
            embedding = np.random.randn(384).astype(np.float32)
            memory = Memory(
                content=f"Content {i}",
                user_id="alice",
                embedding=embedding,
                importance=0.1,
            )
            await index.index(memory)

        # All should be present
        assert index.size == 20

        stats = index.get_memory_stats()
        assert stats.budget_bytes is None  # No budget when unbounded
        assert stats.evictions == 0

    @pytest.mark.asyncio
    async def test_eviction_prefers_low_importance_then_old(self, temp_db_path):
        """Test eviction order: lowest importance first, then oldest."""
        import time

        from headroom.memory.adapters.hnsw import HNSWVectorIndex

        index = HNSWVectorIndex(
            dimension=384,
            max_entries=3,
            eviction_batch_size=1,
        )

        np.random.seed(42)

        # Add memories with same importance but different times
        memories = []
        for i in range(3):
            embedding = np.random.randn(384).astype(np.float32)
            memory = Memory(
                content=f"Content {i}",
                user_id="alice",
                embedding=embedding,
                importance=0.5,  # Same importance
            )
            await index.index(memory)
            memories.append(memory)
            time.sleep(0.01)  # Small delay to ensure different created_at

        # Add one more to trigger eviction
        new_embedding = np.random.randn(384).astype(np.float32)
        await index.index(
            Memory(
                content="New",
                user_id="alice",
                embedding=new_embedding,
                importance=0.5,
            )
        )

        # Should have evicted the oldest (first) entry
        assert index.size == 3
        assert memories[0].id not in index._memory_to_hnsw

    @pytest.mark.asyncio
    async def test_save_load_preserves_eviction_settings(self, temp_db_path):
        """Test that save/load preserves eviction settings."""
        from headroom.memory.adapters.hnsw import HNSWVectorIndex

        index = HNSWVectorIndex(
            dimension=384,
            max_entries=50,
            eviction_batch_size=10,
            save_path=temp_db_path,
        )

        np.random.seed(42)

        # Add some entries
        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)

        # Save
        index.save_index(temp_db_path)

        # Create new index and load
        index2 = HNSWVectorIndex(dimension=384)
        index2.load_index(temp_db_path)

        assert index2._max_entries == 50
        assert index2._eviction_batch_size == 10
        assert index2.size == 5