fix(memory): persist resolved index backend so auto+usearch snapshots reload correctly

Author: Piyussh01

livekit-memory/livekit/memory/_index.py

@@ -43,6 +43,10 @@
 class VectorIndex(Protocol):
     """Minimal index interface used by `MemoryStore`."""
 
+    # Concrete backend tag ("bruteforce" | "usearch"); persisted in snapshots so a
+    # store created with backend="auto" reloads the same index type it resolved to.
+    kind: str
+
     @property
     def dims(self) -> int: ...
 
@@ -66,6 +70,8 @@ class BruteForceIndex:
     partial top-k selection. Removal is O(1) swap-with-last to keep the matrix dense.
     """
 
+    kind = "bruteforce"
+
     def __init__(self, dims: int, *, initial_capacity: int = 1024) -> None:
         self._dims = int(dims)
         self._count = 0
@@ -157,6 +163,8 @@ class UsearchIndex:
     as similarity (1 - distance) to match `BruteForceIndex`.
     """
 
+    kind = "usearch"
+
     def __init__(
         self,
         dims: int,

livekit-memory/livekit/memory/store.py

@@ -75,10 +75,13 @@ def __init__(
         """
         self._embedder = coerce_embedder(embedder, dims)
         self._dims = self._embedder.dims
-        self._backend_name = backend
         self._lock = threading.RLock()
 
         self._collection: VectorIndex = make_index(self._dims, backend, expected_size=expected_size)
+        # Record the *resolved* concrete backend ("bruteforce"|"usearch"), not the
+        # configured one — "auto" depends on expected_size, which we don't persist, so
+        # snapshots and clear() must rebuild from what was actually selected.
+        self._backend_name = self._collection.kind
         # Facts are always exact and fully scanned — keep them brute-force regardless.
         self._facts: VectorIndex = BruteForceIndex(self._dims)
 

tests/memory/test_store.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+import json
+
 import numpy as np
 import pytest
 
@@ -11,6 +13,7 @@
     BruteForceIndex,
     HashingEmbedder,
     MemoryStore,
+    UsearchIndex,
 )
 
 
@@ -116,6 +119,38 @@ def test_save_and_load(tmp_path):
     assert hits[0].metadata == {"k": 1}
 
 
+def test_snapshot_records_resolved_backend_not_auto(tmp_path):
+    # backend="auto" must persist the concrete backend it resolved to, never "auto",
+    # so load() rebuilds the right index type.
+    store = make_store(backend="auto")
+    store.save(tmp_path / "snap")
+    with open(tmp_path / "snap" / "meta.json") as f:
+        meta = json.load(f)
+    assert meta["backend"] in ("bruteforce", "usearch")
+    assert meta["backend"] != "auto"
+
+
+def test_usearch_auto_save_load_roundtrip(tmp_path):
+    # Regression: backend="auto" + large expected_size selects UsearchIndex; the
+    # snapshot must reload as UsearchIndex rather than crashing in BruteForceIndex.load.
+    pytest.importorskip("usearch")
+    emb = HashingEmbedder(dims=64)
+    store = MemoryStore(embedder=emb, backend="auto", expected_size=200_000)
+    assert isinstance(store._collection, UsearchIndex)
+
+    store.upsert("name", "user is Ada")  # fact -> brute force
+    for i in range(20):
+        store.add(f"memory number {i} about topic {i % 5}")
+    store.save(tmp_path / "snap")
+
+    loaded = MemoryStore.load(tmp_path / "snap", embedder=HashingEmbedder(dims=64))
+    assert isinstance(loaded._collection, UsearchIndex)
+    assert len(loaded) == 21
+    assert loaded.all(namespace=FACTS_NAMESPACE)[0].text == "user is Ada"
+    hits = loaded.search("memory number 7 about topic 2", limit=1)
+    assert hits and "number 7" in hits[0].text
+
+
 def test_dimension_mismatch_rejected():
     store = make_store()  # dims=64
     with pytest.raises(ValueError):