diff --git a/headroom/memory/adapters/hnsw.py b/headroom/memory/adapters/hnsw.py index e7787edbe..c678325ac 100644 --- a/headroom/memory/adapters/hnsw.py +++ b/headroom/memory/adapters/hnsw.py @@ -194,9 +194,10 @@ class HNSWVectorIndex: - Persistence support with save_index/load_index - Thread-safe operations with Lock - Optional auto-save on index modifications + - **Bounded memory with LRU eviction** (when max_entries is set) Usage: - index = HNSWVectorIndex(dimension=384) + index = HNSWVectorIndex(dimension=384, max_entries=10000) await index.index(memory_with_embedding) results = await index.search(VectorFilter( query_vector=query_embedding, @@ -211,6 +212,12 @@ class HNSWVectorIndex: better recall but use more memory. Default: 16 - ef_search: Size of dynamic candidate list during search. Higher values give better recall but slower search. Default: 50 + + Memory Bounding: + - max_entries: Soft limit on number of entries. When reached, lowest + importance entries are evicted to make room. Default: None (unbounded). + - eviction_batch_size: Number of entries to evict at once when limit + is reached. Default: 100. """ def __init__( @@ -222,6 +229,8 @@ class HNSWVectorIndex: ef_search: int = 50, auto_save: bool = False, save_path: str | Path | None = None, + max_entries: int | None = None, + eviction_batch_size: int = 100, ) -> None: """Initialize the HNSW vector index. @@ -238,6 +247,9 @@ class HNSWVectorIndex: auto_save: If True and save_path is set, automatically save index after modifications. save_path: Path for auto-save operations. Required if auto_save=True. + max_entries: Soft limit on number of entries. When reached, + lowest importance entries are evicted. None = unbounded. + eviction_batch_size: Number of entries to evict when limit is reached. Raises: ValueError: If auto_save is True but save_path is not provided. @@ -263,6 +275,11 @@ class HNSWVectorIndex: self._auto_save = auto_save self._save_path = Path(save_path) if save_path else None + # Memory bounding + self._max_entries = max_entries + self._eviction_batch_size = eviction_batch_size + self._eviction_count = 0 # Track total evictions for stats + # Initialize HNSW index with cosine similarity # hnswlib uses 'cosine' space which internally normalizes vectors # Note: hnswlib is guaranteed non-None here due to _check_hnswlib_available() above @@ -302,7 +319,8 @@ class HNSWVectorIndex: async def index(self, memory: Memory) -> None: """Index a memory's embedding for similarity search. - The memory must have an embedding set. + The memory must have an embedding set. If max_entries is set and + the limit is reached, low-importance entries are evicted. Args: memory: The memory to index. @@ -327,7 +345,13 @@ class HNSWVectorIndex: # Update metadata self._metadata[memory.id] = IndexedMemoryMetadata.from_memory(memory) else: - # Resize if needed + # Evict if at capacity (before adding new entry) + if self._max_entries is not None: + current_size = len(self._memory_to_hnsw) + if current_size >= self._max_entries: + self._evict_entries(self._eviction_batch_size) + + # Resize HNSW index if needed (separate from entry limit) if self._next_hnsw_id >= self._max_elements: self._resize_index(self._max_elements * 2) @@ -350,6 +374,55 @@ class HNSWVectorIndex: if self._auto_save and self._save_path: self.save_index(self._save_path) + def _evict_entries(self, count: int) -> int: + """Evict the lowest importance entries from the index. + + Must be called with lock held. + + Eviction strategy: Sort by importance (ascending), then by age + (oldest first for ties). Evict the lowest scoring entries. + + Args: + count: Number of entries to evict. + + Returns: + Number of entries actually evicted. + """ + if not self._metadata: + return 0 + + # Sort entries by importance (ascending), then by created_at (oldest first) + sorted_entries = sorted( + self._metadata.items(), + key=lambda x: (x[1].importance, x[1].created_at), + ) + + # Evict the lowest importance entries + evicted = 0 + for memory_id, _metadata in sorted_entries[:count]: + if memory_id not in self._memory_to_hnsw: + continue + + hnsw_id = self._memory_to_hnsw[memory_id] + + # Mark as deleted in HNSW index + self._index.mark_deleted(hnsw_id) + + # Remove from mappings + del self._memory_to_hnsw[memory_id] + del self._hnsw_to_memory[hnsw_id] + + # Remove metadata and embedding + if memory_id in self._metadata: + del self._metadata[memory_id] + if memory_id in self._embeddings: + del self._embeddings[memory_id] + + evicted += 1 + + self._eviction_count += evicted + return evicted + async def index_batch(self, memories: list[Memory]) -> int: """Index multiple memories' embeddings. @@ -736,6 +809,9 @@ class HNSWVectorIndex: "ef_construction": self._ef_construction, "m": self._m, "ef_search": self._ef_search, + "max_entries": self._max_entries, + "eviction_batch_size": self._eviction_batch_size, + "eviction_count": self._eviction_count, "memory_to_hnsw": self._memory_to_hnsw, "hnsw_to_memory": self._hnsw_to_memory, "next_hnsw_id": self._next_hnsw_id, @@ -786,6 +862,11 @@ class HNSWVectorIndex: self._m = meta_data["m"] self._ef_search = meta_data["ef_search"] + # Restore bounding parameters (with defaults for backward compatibility) + self._max_entries = meta_data.get("max_entries") + self._eviction_batch_size = meta_data.get("eviction_batch_size", 100) + self._eviction_count = meta_data.get("eviction_count", 0) + # Create new index and load from file self._index = hnswlib.Index(space="cosine", dim=self._dimension) # type: ignore[union-attr] self._index.load_index( @@ -829,6 +910,7 @@ class HNSWVectorIndex: self._next_hnsw_id = 0 self._metadata.clear() self._embeddings.clear() + self._eviction_count = 0 if self._auto_save and self._save_path: self.save_index(self._save_path) @@ -840,17 +922,21 @@ class HNSWVectorIndex: Dictionary with index metrics. """ with self._lock: + current_size = len(self._memory_to_hnsw) return { - "size": len(self._memory_to_hnsw), + "size": current_size, "dimension": self._dimension, "max_elements": self._max_elements, + "max_entries": self._max_entries, "ef_construction": self._ef_construction, "m": self._m, "ef_search": self._ef_search, + "eviction_count": self._eviction_count, "utilization": ( - (len(self._memory_to_hnsw) / self._max_elements) * 100 - if self._max_elements > 0 - else 0.0 + (current_size / self._max_elements) * 100 if self._max_elements > 0 else 0.0 + ), + "entry_utilization": ( + (current_size / self._max_entries) * 100 if self._max_entries else None ), } @@ -902,14 +988,22 @@ class HNSWVectorIndex: index_size_estimate = len(self._memory_to_hnsw) * (self._dimension * 4 + self._m * 8) size_bytes += index_size_estimate + # Calculate budget based on max_entries if set + # Budget = estimated size at max capacity + budget_bytes = None + if self._max_entries is not None: + # Estimate: each entry ~= embedding bytes + metadata overhead (~500 bytes) + per_entry_estimate = self._dimension * 4 + self._m * 8 + 500 + budget_bytes = self._max_entries * per_entry_estimate + return ComponentStats( name="vector_index", entry_count=len(self._memory_to_hnsw), size_bytes=size_bytes, - budget_bytes=None, + budget_bytes=budget_bytes, hits=0, misses=0, - evictions=0, + evictions=self._eviction_count, ) def set_ef_search(self, ef_search: int) -> None: diff --git a/tests/test_hnsw_only.py b/tests/test_hnsw_only.py index 3e2c9e179..4c5c60df2 100644 --- a/tests/test_hnsw_only.py +++ b/tests/test_hnsw_only.py @@ -78,3 +78,201 @@ class TestHNSWVectorIndex: 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