diff --git a/src/basic_memory/repository/milvus_repository.py b/src/basic_memory/repository/milvus_repository.py index 0d6a870b8..815f102de 100644 --- a/src/basic_memory/repository/milvus_repository.py +++ b/src/basic_memory/repository/milvus_repository.py @@ -239,6 +239,13 @@ def collection_dimensions(self, collection_name: str) -> int | None: raise RuntimeError( f"Milvus collection '{collection_name}' requires a COSINE index on 'embedding'." ) + # Milvus Lite releases persisted collections when the owning process exits. + # Fresh CLI processes must load a compatible existing collection before any + # search, query, or orphan-reconciliation operation can use it. + self._client.load_collection( + collection_name=collection_name, + timeout=self._timeout, + ) return dimensions def create_collection(self, collection_name: str, dimensions: int) -> bool: diff --git a/test-int/semantic/test_milvus_lite.py b/test-int/semantic/test_milvus_lite.py index c73a31d0a..7f4db5e08 100644 --- a/test-int/semantic/test_milvus_lite.py +++ b/test-int/semantic/test_milvus_lite.py @@ -2,6 +2,7 @@ from __future__ import annotations +import subprocess import sys import pytest @@ -23,6 +24,61 @@ ] +_RESTART_SCRIPT = """ +import asyncio +import sys + +from basic_memory.repository.milvus_config import MilvusSettings +from basic_memory.repository.milvus_index import MilvusVectorIndex +from basic_memory.repository.semantic_vector_index import ( + VectorIndexScope, + VectorKey, + VectorRecord, +) + + +async def main() -> None: + phase, database_path = sys.argv[1:] + scope = VectorIndexScope( + namespace="restart-database", + project_id=7, + embedding_identity="Stub:model", + dimensions=3, + ) + key = VectorKey(entity_id=1, chunk_key="summary:0") + index = MilvusVectorIndex(scope, MilvusSettings(uri=database_path)) + + if phase == "write": + await index.upsert( + [ + VectorRecord( + key=key, + source_hash="auth-v1", + values=(1.0, 0.0, 0.0), + ) + ] + ) + return + + await index.delete_orphans([key]) + matches = await index.search((1.0, 0.0, 0.0), limit=1) + assert [match.key for match in matches] == [key] + + +asyncio.run(main()) +""" + + +def _run_restart_phase(phase: str, database_path: str) -> None: + result = subprocess.run( + [sys.executable, "-c", _RESTART_SCRIPT, phase, database_path], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + @pytest.mark.asyncio async def test_milvus_lite_vector_lifecycle(tmp_path) -> None: scope = VectorIndexScope( @@ -64,3 +120,10 @@ async def test_milvus_lite_vector_lifecycle(tmp_path) -> None: await index.delete_orphans([]) assert await index.search((1.0, 0.0, 0.0), limit=2) == [] + + +def test_milvus_lite_reloads_collection_after_process_restart(tmp_path) -> None: + database_path = str(tmp_path / "restart-vectors.db") + + _run_restart_phase("write", database_path) + _run_restart_phase("read", database_path) diff --git a/tests/repository/test_milvus_repository.py b/tests/repository/test_milvus_repository.py index e48f5bd82..0682d6e32 100644 --- a/tests/repository/test_milvus_repository.py +++ b/tests/repository/test_milvus_repository.py @@ -104,6 +104,7 @@ def __init__(self) -> None: self.index_params = FakeIndexParams() self.created: list[dict[str, object]] = [] self.create_exception: MilvusException | None = None + self.loaded: list[dict[str, object]] = [] self.upserts: list[dict[str, object]] = [] self.deletes: list[dict[str, object]] = [] self.iterator = FakeIterator([[{"id": "one"}, {"id": "two"}], []]) @@ -166,6 +167,15 @@ def create_collection(self, **kwargs: object) -> None: if self.create_exception is not None: raise self.create_exception + def load_collection(self, *, collection_name: str, timeout: float) -> None: + self._record_timeout("load_collection", timeout) + self.loaded.append( + { + "collection_name": collection_name, + "timeout": timeout, + } + ) + def upsert(self, **kwargs: object) -> None: self._record_kwargs_timeout("upsert", kwargs) self.upserts.append(kwargs) @@ -248,6 +258,7 @@ def test_repository_applies_finite_deadline_to_every_remote_operation( "describe_index", "has_collection", "list_indexes", + "load_collection", "query_iterator", "search", "upsert", @@ -258,7 +269,7 @@ def test_repository_applies_finite_deadline_to_every_remote_operation( assert client.searches[0]["consistency_level"] == "Strong" -def test_collection_dimensions_reports_missing_and_valid_collection( +def test_collection_dimensions_reports_missing_and_loads_valid_collection( repository: PyMilvusRepository, client: FakeClient, ) -> None: @@ -266,6 +277,7 @@ def test_collection_dimensions_reports_missing_and_valid_collection( client.has_collection_result = True assert repository.collection_dimensions("vectors") == 4 + assert client.loaded == [{"collection_name": "vectors", "timeout": 12.5}] @pytest.mark.parametrize(