mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
b0afee85b3
## Description `HNSWVectorIndex.index_batch()` used the live memory map size to decide whether to resize hnswlib before adding new labels. hnswlib does not reclaim capacity slots when labels are removed with `mark_deleted`, so after delete/evict churn the live count can be much lower than the assigned-id high-water mark. That lets a batch add skip resizing and then fail in `add_items` with `number of elements exceeds the specified limit`. ## Fix - Size the batch resize check from `self._next_hnsw_id`, which has already been incremented for the new batch labels. - Match the single-item `index()` path's high-water-mark capacity behavior. - Add a regression test that deletes most entries from a small index and then batch-adds enough new memories to require a resize. - Merge current `main` to refresh mergeability and stale lint results. ## Testing ```text uvx ruff@0.15.17 check headroom/memory/adapters/hnsw.py tests/test_memory/test_hnsw_batch_capacity.py headroom/memory/factory.py All checks passed! uvx ruff@0.15.17 format --check headroom/memory/adapters/hnsw.py tests/test_memory/test_hnsw_batch_capacity.py headroom/memory/factory.py 3 files already formatted git diff --check headroomlabs/main...HEAD # no output uv run --extra dev python -m pytest tests/test_memory/test_hnsw_batch_capacity.py -q 1 passed, 18 warnings ``` ## Review Readiness - [x] Ready for review - [x] Regression test added - [x] CHANGELOG updated Co-authored-by: JerrettDavis <mxjerrett@gmail.com> Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
"""HNSW index_batch must resize based on the assigned-id high-water mark, not
|
|
the live entry count, so batch adds after eviction/deletion churn don't overflow
|
|
hnswlib's max_elements."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from headroom.memory.models import Memory
|
|
|
|
try:
|
|
from headroom.memory.adapters.hnsw import _check_hnswlib_available
|
|
|
|
HNSW_AVAILABLE = _check_hnswlib_available()
|
|
except ImportError:
|
|
HNSW_AVAILABLE = False
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_hnsw_path():
|
|
with tempfile.NamedTemporaryFile(suffix=".hnsw", delete=False) as f:
|
|
yield Path(f.name)
|
|
|
|
|
|
def _mem(i: int, dim: int = 8) -> Memory:
|
|
rng = np.random.default_rng(i)
|
|
return Memory(
|
|
content=f"m{i}",
|
|
user_id="u",
|
|
embedding=rng.standard_normal(dim).astype(np.float32),
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not HNSW_AVAILABLE, reason="hnswlib not installed")
|
|
@pytest.mark.asyncio
|
|
async def test_index_batch_after_deletion_churn_does_not_overflow(temp_hnsw_path):
|
|
from headroom.memory.adapters.hnsw import HNSWVectorIndex
|
|
|
|
# Small ceiling so we hit it quickly. mark_deleted (remove) never frees a
|
|
# slot, so the assigned-id counter climbs toward max_elements while the live
|
|
# count stays low.
|
|
index = HNSWVectorIndex(dimension=8, max_elements=8, save_path=temp_hnsw_path)
|
|
|
|
singles = [_mem(i) for i in range(6)]
|
|
for m in singles:
|
|
await index.index(m) # assigned ids 0..5; next id high-water = 6
|
|
|
|
# Delete 5 of them (mark_deleted; the 5 hnswlib slots are NOT reclaimed).
|
|
for m in singles[:5]:
|
|
await index.remove(m.id)
|
|
|
|
# A batch of 3 now needs slots 6,7,8 -> hnswlib must hold 9 labels. The old
|
|
# check used the live count (1) + 3 = 4 <= 8 and skipped the resize, so
|
|
# add_items raised "number of elements exceeds the specified limit".
|
|
added = await index.index_batch([_mem(100), _mem(101), _mem(102)])
|
|
|
|
assert added == 3
|