diff --git a/.env.example b/.env.example index d47c6083..737599d4 100644 --- a/.env.example +++ b/.env.example @@ -26,3 +26,11 @@ EMBEDDING_DIMENSION=2560 # Data directory inside container (default set by Dockerfile/compose; override if needed) # DATA_DIR=/app/MCP/data + +# Optional Milvus backend settings for applications that inject +# MilvusVectorStoreBackend. LanceDB remains the default. +MILVUS_URI=./milvus.db +MILVUS_TOKEN= +MILVUS_DB_NAME= +MILVUS_COLLECTION_NAME=memory_entries +MILVUS_CONSISTENCY_LEVEL=Session diff --git a/README.md b/README.md index 7fbd7d8e..78749c82 100644 --- a/README.md +++ b/README.md @@ -344,6 +344,7 @@ pip install -r requirements.txt # — OR — install as an editable package pip install -e . # default: text + multimodal + evolver +pip install -e ".[milvus]" # + optional Milvus vector store backend pip install -e ".[server]" # + MCP / HTTP server (mcp, fastapi, ...) pip install -e ".[all]" # everything, including dev tools @@ -376,6 +377,51 @@ EMBEDDING_MODEL = "Qwen/Qwen3-Embedding-0.6B" > `deepseek-ai/deepseek-v4-pro` is a reasoning model, so leave enough output budget; if responses come back empty, raise the token limit (>= 512). Any other OpenAI-compatible chat model on Atlas Cloud (e.g. `Qwen/Qwen3-Next-80B-A3B-Instruct`, `zai-org/glm-5`, `moonshotai/kimi-k2.6`) works the same way. +### Optional Milvus vector store + +LanceDB remains the default vector store. Install the `milvus` extra and inject +`MilvusVectorStoreBackend` when you want the same semantic, BM25 keyword, and +structured retrieval paths backed by Milvus: + +```bash +pip install -e ".[milvus]" +``` + +```python +from simplemem.core.database import MilvusVectorStoreBackend, VectorStore +from simplemem.core.settings import settings + +store = VectorStore( + backend_factory=lambda dimension: MilvusVectorStoreBackend( + collection_name=settings.MILVUS_COLLECTION_NAME, + vector_dimension=dimension, + uri=settings.MILVUS_URI, + token=settings.MILVUS_TOKEN, + db_name=settings.MILVUS_DB_NAME, + consistency_level=settings.MILVUS_CONSISTENCY_LEVEL, + ) +) +``` + +The default `MILVUS_URI=./milvus.db` starts Milvus Lite with a local database +file. The same backend connects to Milvus Server or Zilliz Cloud without code +changes: + +```bash +# Milvus Server +export MILVUS_URI=http://localhost:19530 + +# Zilliz Cloud +export MILVUS_URI=https://YOUR-ENDPOINT.api.gcp-us-west1.zillizcloud.com +export MILVUS_TOKEN=YOUR_API_KEY +``` + +Set `MILVUS_DB_NAME`, `MILVUS_COLLECTION_NAME`, and +`MILVUS_CONSISTENCY_LEVEL` when the deployment requires non-default database, +collection, or consistency settings. Milvus Lite 3.x uses a storage format that +is not compatible with databases created by Milvus Lite 2.x; create a new local +database or migrate the old data before upgrading. + --- ## 🐳 Run with Docker @@ -580,5 +626,5 @@ This project is licensed under the **MIT License** - see the [LICENSE](LICENSE) We would like to thank the following projects and teams: - 🔍 **Embedding Model**: [Qwen3-Embedding](https://github.com/QwenLM/Qwen) - State-of-the-art retrieval performance -- 🗄️ **Vector Database**: [LanceDB](https://lancedb.com/) - High-performance columnar storage +- 🗄️ **Vector Databases**: [LanceDB](https://lancedb.com/) by default, with optional [Milvus](https://milvus.io/) support - 📊 **Benchmark**: [LoCoMo](https://github.com/snap-research/locomo) - Long-context memory evaluation framework diff --git a/config.py.example b/config.py.example index 78c8e6a4..187ce59d 100644 --- a/config.py.example +++ b/config.py.example @@ -89,6 +89,15 @@ LANCEDB_PATH = "./lancedb_data" # Memory table name MEMORY_TABLE_NAME = "memory_entries" +# Optional Milvus backend settings. LanceDB remains the default backend. +# Use a local file for Milvus Lite, http://localhost:19530 for Milvus Server, +# or a Zilliz Cloud public endpoint with MILVUS_TOKEN. +MILVUS_URI = "./milvus.db" +MILVUS_TOKEN = "" +MILVUS_DB_NAME = "" +MILVUS_COLLECTION_NAME = "memory_entries" +MILVUS_CONSISTENCY_LEVEL = "Session" + # ============================================================================ diff --git a/setup.py b/setup.py index e51428de..4ec9bcef 100644 --- a/setup.py +++ b/setup.py @@ -51,6 +51,9 @@ def _read_version() -> str: EXTRAS = { + "milvus": [ + "pymilvus[milvus-lite]>=3.0.0", + ], # MCP / HTTP server integration — bounds + members from MCP/requirements.txt # (authoritative for the MCP/server/ subproject). Members already in defaults # (lancedb, pyarrow, httpx, pydantic) intentionally omitted here — pip will pick diff --git a/simplemem/core/config_default.py b/simplemem/core/config_default.py index edeab553..5c25b1d2 100644 --- a/simplemem/core/config_default.py +++ b/simplemem/core/config_default.py @@ -87,6 +87,13 @@ # Memory table name MEMORY_TABLE_NAME = "memory_entries" +# Optional Milvus backend settings. LanceDB remains the default backend. +MILVUS_URI = "./milvus.db" +MILVUS_TOKEN = "" +MILVUS_DB_NAME = "" +MILVUS_COLLECTION_NAME = "memory_entries" +MILVUS_CONSISTENCY_LEVEL = "Session" + # ============================================================================ diff --git a/simplemem/core/database/__init__.py b/simplemem/core/database/__init__.py index fa10871a..38643af5 100644 --- a/simplemem/core/database/__init__.py +++ b/simplemem/core/database/__init__.py @@ -1,3 +1,6 @@ +from simplemem.core.database.milvus_vector_store_backend import ( + MilvusVectorStoreBackend, +) from simplemem.core.database.vector_store import VectorStore from simplemem.core.database.vector_store_backend import ( LanceDBVectorStoreBackend, @@ -9,6 +12,7 @@ __all__ = [ "LanceDBVectorStoreBackend", + "MilvusVectorStoreBackend", "ScoreOrder", "VectorStore", "VectorStoreBackend", diff --git a/simplemem/core/database/milvus_vector_store_backend.py b/simplemem/core/database/milvus_vector_store_backend.py new file mode 100644 index 00000000..b9b71ba9 --- /dev/null +++ b/simplemem/core/database/milvus_vector_store_backend.py @@ -0,0 +1,623 @@ +"""Optional Milvus implementation of the SimpleMem vector store contract.""" + +import json +import math +import os +import re +from importlib.metadata import PackageNotFoundError, version +from typing import Any, Dict, List, Optional, Sequence + +from simplemem.core.database.vector_store_backend import ( + ScoreOrder, + VectorStoreRecord, + VectorStoreSearchResult, +) + + +class MilvusVectorStoreBackend: + """Store SimpleMem records in Milvus Lite, Milvus, or Zilliz Cloud.""" + + semantic_score_order = ScoreOrder.ASCENDING + keyword_score_order = ScoreOrder.DESCENDING + + _ENTRY_ID_MAX_LENGTH = 4096 + _TEXT_MAX_LENGTH = 65535 + _METADATA_MAX_LENGTH = 8192 + _ARRAY_MAX_CAPACITY = 1024 + _ARRAY_ITEM_MAX_LENGTH = 4096 + _SCALAR_FILTER_FIELDS = { + "entry_id", + "lossless_restatement", + "timestamp", + "location", + "topic", + } + _ARRAY_FILTER_FIELDS = {"keywords", "persons", "entities"} + _METADATA_FIELDS = ( + "lossless_restatement", + "keywords", + "timestamp", + "location", + "persons", + "entities", + "topic", + ) + _OUTPUT_FIELDS = ["entry_id", *_METADATA_FIELDS] + _REMOTE_URI_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*://") + _WINDOWS_PATH_PATTERN = re.compile(r"^[A-Za-z]:[\\/]") + + def __init__( + self, + collection_name: Optional[str] = None, + vector_dimension: int = 0, + uri: Optional[str] = None, + token: Optional[str] = None, + db_name: Optional[str] = None, + consistency_level: Optional[str] = None, + ): + if vector_dimension <= 0: + raise ValueError("vector_dimension must be greater than zero") + + MilvusClient, DataType, Function, FunctionType = self._load_pymilvus() + self._data_type = DataType + self._function = Function + self._function_type = FunctionType + self.collection_name = collection_name or os.getenv( + "MILVUS_COLLECTION_NAME", "memory_entries" + ) + self.vector_dimension = vector_dimension + self.uri = uri if uri is not None else os.getenv("MILVUS_URI", "./milvus.db") + self.token = token if token is not None else os.getenv("MILVUS_TOKEN", "") + self.db_name = ( + db_name if db_name is not None else os.getenv("MILVUS_DB_NAME", "") + ) + self.consistency_level = consistency_level or os.getenv( + "MILVUS_CONSISTENCY_LEVEL", "Session" + ) + self._use_raw_cosine_distance = self._is_lite_3_0_cosine_distance( + self.uri, + self._package_version("milvus-lite"), + ) + + client_options = {"uri": self.uri} + if self.token: + client_options["token"] = self.token + if self.db_name: + client_options["db_name"] = self.db_name + self.client = MilvusClient(**client_options) + try: + self._init_collection() + except Exception: + self.client.close() + raise + + @staticmethod + def _load_pymilvus(): + try: + from pymilvus import DataType, Function, FunctionType, MilvusClient + except ImportError as error: + raise ImportError( + "Milvus support requires the optional dependency. " + 'Install SimpleMem with `pip install -e ".[milvus]"`.' + ) from error + return MilvusClient, DataType, Function, FunctionType + + @staticmethod + def _package_version(package_name: str) -> Optional[str]: + try: + return version(package_name) + except PackageNotFoundError: + return None + + @classmethod + def _is_lite_3_0_cosine_distance( + cls, + uri: str, + milvus_lite_version: Optional[str], + ) -> bool: + return cls._is_local_path_uri(uri) and milvus_lite_version in {"3.0", "3.0.0"} + + @classmethod + def _is_local_path_uri(cls, uri: str) -> bool: + if cls._REMOTE_URI_PATTERN.match(uri): + return False + if cls._WINDOWS_PATH_PATTERN.match(uri): + return True + return ":" not in uri + + def _init_collection(self) -> None: + if self.client.has_collection(collection_name=self.collection_name): + self._validate_collection_schema() + return + self._create_collection() + + def _create_collection(self) -> None: + schema = self.client.create_schema(auto_id=False, enable_dynamic_field=False) + schema.add_field( + field_name="entry_id", + datatype=self._data_type.VARCHAR, + is_primary=True, + max_length=self._ENTRY_ID_MAX_LENGTH, + ) + schema.add_field( + field_name="lossless_restatement", + datatype=self._data_type.VARCHAR, + max_length=self._TEXT_MAX_LENGTH, + enable_analyzer=True, + enable_match=True, + analyzer_params={"type": "standard"}, + ) + for field_name in ("keywords", "persons", "entities"): + schema.add_field( + field_name=field_name, + datatype=self._data_type.ARRAY, + element_type=self._data_type.VARCHAR, + max_capacity=self._ARRAY_MAX_CAPACITY, + max_length=self._ARRAY_ITEM_MAX_LENGTH, + ) + for field_name in ("timestamp", "location", "topic"): + schema.add_field( + field_name=field_name, + datatype=self._data_type.VARCHAR, + max_length=self._METADATA_MAX_LENGTH, + ) + schema.add_field( + field_name="vector", + datatype=self._data_type.FLOAT_VECTOR, + dim=self.vector_dimension, + ) + schema.add_field( + field_name="sparse", + datatype=self._data_type.SPARSE_FLOAT_VECTOR, + ) + schema.add_function( + self._function( + name="lossless_restatement_bm25", + function_type=self._function_type.BM25, + input_field_names=["lossless_restatement"], + output_field_names=["sparse"], + ) + ) + + index_params = self.client.prepare_index_params() + index_params.add_index( + field_name="vector", + index_type="AUTOINDEX", + metric_type="COSINE", + ) + index_params.add_index( + field_name="sparse", + index_type="AUTOINDEX", + metric_type="BM25", + ) + self.client.create_collection( + collection_name=self.collection_name, + schema=schema, + index_params=index_params, + consistency_level=self.consistency_level, + ) + + def _validate_collection_schema(self) -> None: + description = self.client.describe_collection( + collection_name=self.collection_name + ) + fields = {field["name"]: field for field in description.get("fields", [])} + expected_types = { + "entry_id": self._data_type.VARCHAR, + "lossless_restatement": self._data_type.VARCHAR, + "keywords": self._data_type.ARRAY, + "timestamp": self._data_type.VARCHAR, + "location": self._data_type.VARCHAR, + "persons": self._data_type.ARRAY, + "entities": self._data_type.ARRAY, + "topic": self._data_type.VARCHAR, + "vector": self._data_type.FLOAT_VECTOR, + "sparse": self._data_type.SPARSE_FLOAT_VECTOR, + } + for field_name, expected_type in expected_types.items(): + field = fields.get(field_name) + if field is None: + raise ValueError( + f"Milvus collection {self.collection_name!r} is missing required " + f"field {field_name!r}" + ) + if field.get("type") != expected_type: + raise ValueError( + f"Milvus collection {self.collection_name!r} field " + f"{field_name!r} has type {field.get('type')!r}; " + f"expected {expected_type!r}" + ) + + primary_field = fields["entry_id"] + if not primary_field.get("is_primary"): + raise ValueError( + f"Milvus collection {self.collection_name!r} must use entry_id as " + "its primary key" + ) + vector_dimension = int(fields["vector"].get("params", {}).get("dim", 0)) + if vector_dimension != self.vector_dimension: + raise ValueError( + f"Milvus collection {self.collection_name!r} has vector dimension " + f"{vector_dimension}; expected {self.vector_dimension}" + ) + for field_name in self._ARRAY_FILTER_FIELDS: + if fields[field_name].get("element_type") != self._data_type.VARCHAR: + raise ValueError( + f"Milvus collection {self.collection_name!r} field " + f"{field_name!r} must contain VARCHAR values" + ) + + functions = description.get("functions", []) + has_bm25_function = any( + function.get("type") == self._function_type.BM25 + and function.get("input_field_names") == ["lossless_restatement"] + and function.get("output_field_names") == ["sparse"] + for function in functions + ) + if not has_bm25_function: + raise ValueError( + f"Milvus collection {self.collection_name!r} is missing the " + "lossless_restatement BM25 function" + ) + + def insert(self, records: Sequence[VectorStoreRecord]) -> None: + if not records: + return + rows = [self._record_to_row(record) for record in records] + self.client.insert(collection_name=self.collection_name, data=rows) + + def _record_to_row(self, record: VectorStoreRecord) -> Dict[str, Any]: + if not isinstance(record.entry_id, str): + raise TypeError("Milvus entry_id values must be strings") + self._validate_string_length( + "entry_id", record.entry_id, self._ENTRY_ID_MAX_LENGTH + ) + vector = [float(value) for value in record.vector] + if len(vector) != self.vector_dimension: + raise ValueError( + f"Record {record.entry_id!r} has vector dimension {len(vector)}; " + f"expected {self.vector_dimension}" + ) + if not all(math.isfinite(value) for value in vector): + raise ValueError(f"Record {record.entry_id!r} contains a non-finite vector") + + unknown_fields = set(record.metadata) - set(self._METADATA_FIELDS) + if unknown_fields: + raise ValueError( + "Milvus records contain unsupported metadata fields: " + + ", ".join(sorted(unknown_fields)) + ) + row = { + "entry_id": record.entry_id, + "lossless_restatement": self._metadata_string( + record.metadata, "lossless_restatement", self._TEXT_MAX_LENGTH + ), + "keywords": self._metadata_string_array(record.metadata, "keywords"), + "timestamp": self._metadata_string( + record.metadata, "timestamp", self._METADATA_MAX_LENGTH + ), + "location": self._metadata_string( + record.metadata, "location", self._METADATA_MAX_LENGTH + ), + "persons": self._metadata_string_array(record.metadata, "persons"), + "entities": self._metadata_string_array(record.metadata, "entities"), + "topic": self._metadata_string( + record.metadata, "topic", self._METADATA_MAX_LENGTH + ), + "vector": vector, + } + return row + + @classmethod + def _metadata_string( + cls, + metadata: Dict[str, Any], + field_name: str, + max_length: int, + ) -> str: + value = metadata.get(field_name, "") + if value is None: + value = "" + if not isinstance(value, str): + raise TypeError(f"Milvus metadata field {field_name!r} must be a string") + cls._validate_string_length(field_name, value, max_length) + return value + + @classmethod + def _metadata_string_array( + cls, + metadata: Dict[str, Any], + field_name: str, + ) -> List[str]: + value = metadata.get(field_name, []) + if value is None: + value = [] + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError( + f"Milvus metadata field {field_name!r} must be a sequence of strings" + ) + values = list(value) + if len(values) > cls._ARRAY_MAX_CAPACITY: + raise ValueError( + f"Milvus metadata field {field_name!r} exceeds the maximum of " + f"{cls._ARRAY_MAX_CAPACITY} values" + ) + for item in values: + if not isinstance(item, str): + raise TypeError( + f"Milvus metadata field {field_name!r} must contain only strings" + ) + cls._validate_string_length(field_name, item, cls._ARRAY_ITEM_MAX_LENGTH) + return values + + @staticmethod + def _validate_string_length( + field_name: str, + value: str, + max_length: int, + ) -> None: + if len(value.encode("utf-8")) > max_length: + raise ValueError( + f"Milvus field {field_name!r} exceeds its {max_length}-byte limit" + ) + + def semantic_search( + self, + query_vector: Sequence[float], + top_k: int, + filters: Optional[Dict[str, Any]] = None, + ) -> List[VectorStoreSearchResult]: + if top_k <= 0 or self.count() == 0: + return [] + vector = [float(value) for value in query_vector] + if len(vector) != self.vector_dimension: + raise ValueError( + f"Query vector has dimension {len(vector)}; " + f"expected {self.vector_dimension}" + ) + if not all(math.isfinite(value) for value in vector): + raise ValueError("Query vector contains a non-finite value") + + expression = self._build_filter_expression(filters or {}) + hits = self.client.search( + collection_name=self.collection_name, + data=[vector], + anns_field="vector", + filter=expression, + limit=top_k, + output_fields=self._OUTPUT_FIELDS, + search_params={"metric_type": "COSINE", "params": {}}, + consistency_level=self.consistency_level, + )[0] + results = [ + self._hit_to_result( + hit, + score=self._semantic_distance(float(hit["distance"])), + ) + for hit in hits + ] + results.sort(key=lambda result: result.score) + return results + + def _semantic_distance(self, raw_score: float) -> float: + if self._use_raw_cosine_distance: + return raw_score + return 1.0 - raw_score + + def keyword_search( + self, + keywords: Sequence[str], + top_k: int, + ) -> List[VectorStoreSearchResult]: + if top_k <= 0 or not keywords or self.count() == 0: + return [] + if isinstance(keywords, (str, bytes)) or not all( + isinstance(keyword, str) for keyword in keywords + ): + raise TypeError("Milvus keyword queries must be a sequence of strings") + query = " ".join(keyword for keyword in keywords if keyword.strip()).strip() + if not query: + return [] + + hits = self.client.search( + collection_name=self.collection_name, + data=[query], + anns_field="sparse", + limit=top_k, + output_fields=self._OUTPUT_FIELDS, + search_params={"metric_type": "BM25", "params": {}}, + consistency_level=self.consistency_level, + )[0] + results = [ + self._hit_to_result( + hit, + score=self._bm25_relevance(float(hit["distance"])), + ) + for hit in hits + ] + results.sort(key=lambda result: result.score, reverse=True) + return results + + @staticmethod + def _bm25_relevance(raw_score: float) -> float: + return -raw_score if raw_score < 0 else raw_score + + def structured_search( + self, + persons: Optional[Sequence[str]] = None, + timestamp_range: Optional[tuple] = None, + location: Optional[str] = None, + entities: Optional[Sequence[str]] = None, + top_k: Optional[int] = None, + ) -> List[VectorStoreSearchResult]: + if self.count() == 0: + return [] + if top_k is not None and top_k <= 0: + return [] + if not any([persons, timestamp_range, location, entities]): + return [] + + conditions = [] + if persons: + conditions.append( + "ARRAY_CONTAINS_ANY(persons, " + f"{self._format_string_list(persons, 'persons')})" + ) + if location: + if not isinstance(location, str): + raise TypeError("Milvus location filters must be strings") + if "%" in location or "_" in location: + raise ValueError( + "Milvus location filters do not accept LIKE wildcard characters" + ) + conditions.append(f"location like {self._quote(f'%{location}%')}") + if entities: + conditions.append( + "ARRAY_CONTAINS_ANY(entities, " + f"{self._format_string_list(entities, 'entities')})" + ) + if timestamp_range: + if not isinstance(timestamp_range, tuple) or len(timestamp_range) != 2: + raise TypeError("timestamp_range must be a two-item tuple") + start_time, end_time = timestamp_range + if not isinstance(start_time, str) or not isinstance(end_time, str): + raise TypeError("Milvus timestamp bounds must be strings") + conditions.append( + f"timestamp >= {self._quote(start_time)} " + f"and timestamp <= {self._quote(end_time)}" + ) + + expression = " and ".join(conditions) + if top_k is not None: + rows = self.client.query( + collection_name=self.collection_name, + filter=expression, + output_fields=self._OUTPUT_FIELDS, + limit=top_k, + consistency_level=self.consistency_level, + ) + else: + rows = self._query_all(filter_expression=expression) + return [self._row_to_result(row) for row in rows] + + @classmethod + def _build_filter_expression(cls, filters: Dict[str, Any]) -> str: + conditions = [] + for field_name, value in filters.items(): + if field_name in cls._SCALAR_FILTER_FIELDS: + if cls._is_filter_sequence(value): + conditions.append( + f"{field_name} in {cls._format_string_list(value, field_name)}" + ) + elif isinstance(value, str): + conditions.append(f"{field_name} == {cls._quote(value)}") + else: + raise TypeError( + f"Milvus scalar filter {field_name!r} supports only strings " + "or sequences of strings" + ) + elif field_name in cls._ARRAY_FILTER_FIELDS: + if cls._is_filter_sequence(value): + values = cls._format_string_list(value, field_name) + conditions.append(f"ARRAY_CONTAINS_ANY({field_name}, {values})") + elif isinstance(value, str): + conditions.append( + f"ARRAY_CONTAINS({field_name}, {cls._quote(value)})" + ) + else: + raise TypeError( + f"Milvus array filter {field_name!r} supports only strings " + "or sequences of strings" + ) + else: + raise ValueError(f"Invalid semantic filter field: {field_name!r}") + return " and ".join(conditions) + + @staticmethod + def _is_filter_sequence(value: Any) -> bool: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes)) + + @classmethod + def _format_string_list( + cls, + values: Sequence[Any], + field_name: str, + ) -> str: + if isinstance(values, (str, bytes)): + raise TypeError( + f"Milvus filter {field_name!r} must be a sequence of strings" + ) + values = list(values) + if not values: + raise ValueError(f"Milvus filter {field_name!r} cannot be empty") + if not all(isinstance(value, str) for value in values): + raise TypeError(f"Milvus filter {field_name!r} must contain only strings") + return "[" + ", ".join(cls._quote(value) for value in values) + "]" + + @staticmethod + def _quote(value: str) -> str: + return json.dumps(value, ensure_ascii=False) + + def count(self) -> int: + rows = self.client.query( + collection_name=self.collection_name, + output_fields=["count(*)"], + consistency_level=self.consistency_level, + ) + return int(rows[0]["count(*)"]) if rows else 0 + + def get_all(self) -> List[VectorStoreSearchResult]: + return [self._row_to_result(row) for row in self._query_all()] + + def _query_all(self, filter_expression: str = "") -> List[Dict[str, Any]]: + iterator = self.client.query_iterator( + collection_name=self.collection_name, + filter=filter_expression, + output_fields=self._OUTPUT_FIELDS, + batch_size=1000, + consistency_level=self.consistency_level, + ) + rows_by_id = {} + try: + while True: + batch = iterator.next() + if not batch: + break + for row in batch: + rows_by_id[row["entry_id"]] = row + finally: + iterator.close() + return list(rows_by_id.values()) + + def optimize(self) -> None: + # Milvus maintains AUTOINDEX indexes automatically. Explicit compaction is + # asynchronous and is not required for the backend contract. + return None + + def clear(self) -> None: + if self.client.has_collection(collection_name=self.collection_name): + self.client.drop_collection(collection_name=self.collection_name) + self._create_collection() + + def close(self) -> None: + """Close the underlying Milvus client connection.""" + self.client.close() + + def _hit_to_result( + self, + hit: Dict[str, Any], + score: float, + ) -> VectorStoreSearchResult: + entity = dict(hit.get("entity") or {}) + entry_id = str(entity.pop("entry_id", hit.get("id", ""))) + return VectorStoreSearchResult( + entry_id=entry_id, + metadata={field: entity.get(field) for field in self._METADATA_FIELDS}, + score=score, + ) + + def _row_to_result(self, row: Dict[str, Any]) -> VectorStoreSearchResult: + return VectorStoreSearchResult( + entry_id=str(row["entry_id"]), + metadata={field: row.get(field) for field in self._METADATA_FIELDS}, + ) diff --git a/simplemem/core/settings.py b/simplemem/core/settings.py index f1ece2f3..a174e4a1 100644 --- a/simplemem/core/settings.py +++ b/simplemem/core/settings.py @@ -30,6 +30,11 @@ "STRUCTURED_TOP_K": 5, "LANCEDB_PATH": "./lancedb_data", "MEMORY_TABLE_NAME": "memory_entries", + "MILVUS_URI": "./milvus.db", + "MILVUS_TOKEN": "", + "MILVUS_DB_NAME": "", + "MILVUS_COLLECTION_NAME": "memory_entries", + "MILVUS_CONSISTENCY_LEVEL": "Session", "ENABLE_PARALLEL_PROCESSING": True, "MAX_PARALLEL_WORKERS": 16, "ENABLE_PARALLEL_RETRIEVAL": True, diff --git a/tests/test_milvus_vector_store_backend.py b/tests/test_milvus_vector_store_backend.py new file mode 100644 index 00000000..70e047ef --- /dev/null +++ b/tests/test_milvus_vector_store_backend.py @@ -0,0 +1,423 @@ +import json + +import numpy as np +import pytest + +pytest.importorskip("pymilvus") +from pymilvus import DataType, Function, FunctionType, MilvusClient + +from simplemem.core.database import ( + MilvusVectorStoreBackend, + ScoreOrder, + VectorStore, + VectorStoreRecord, +) +from simplemem.core.hybrid_retriever import HybridRetriever +from simplemem.core.models.memory_entry import MemoryEntry + + +class DeterministicEmbedder: + dimension = 3 + + def encode_documents(self, texts): + return np.stack([self._encode(text) for text in texts]) + + def encode_single(self, text, is_query=False): + return self._encode(text) + + @staticmethod + def _encode(text): + text = text.lower() + if "coffee" in text or "espresso" in text: + vector = [1.0, 0.0, 0.0] + elif "apollo" in text or "budget" in text: + vector = [0.0, 1.0, 0.0] + else: + vector = [-1.0, 0.0, 0.0] + return np.array(vector, dtype=np.float32) + + +class DeterministicLLM: + def chat_completion(self, messages, **kwargs): + prompt = messages[-1]["content"] + if "extract key information" in prompt: + return json.dumps( + { + "keywords": ["Apollo"], + "persons": ["Carol"], + "time_expression": None, + "location": None, + "entities": [], + } + ) + if "information requirements analysis" in prompt: + return json.dumps( + { + "reasoning": "Use one semantic query.", + "queries": ["coffee status"], + } + ) + if "determine what specific information is required" in prompt: + return json.dumps( + { + "question_type": "factual", + "key_entities": ["coffee", "Apollo", "Carol"], + "required_info": [ + { + "info_type": "facts", + "description": "Retrieve all three facts", + "priority": "high", + } + ], + "relationships": [], + "minimal_queries_needed": 1, + } + ) + raise AssertionError(f"Unexpected LLM prompt: {prompt[:120]}") + + @staticmethod + def extract_json(response): + return json.loads(response) + + +def _record( + entry_id, + vector, + text, + keywords, + timestamp, + location, + persons, + entities, + topic, +): + return VectorStoreRecord( + entry_id=entry_id, + vector=vector, + metadata={ + "lossless_restatement": text, + "keywords": keywords, + "timestamp": timestamp, + "location": location, + "persons": persons, + "entities": entities, + "topic": topic, + }, + ) + + +@pytest.fixture +def records(): + return [ + _record( + "coffee", + [1.0, 0.0, 0.0], + "Alice drinks espresso at the neighborhood cafe.", + ["coffee", "espresso"], + "2026-01-10T09:00:00", + "New York", + ["Alice"], + ["Neighborhood Cafe"], + "coffee", + ), + _record( + "apollo", + [0.0, 1.0, 0.0], + "The Project Apollo budget budget was approved.", + ["Apollo", "budget"], + "2026-02-10T09:00:00", + "Paris", + ["Bob"], + ["Project Apollo"], + "finance", + ), + _record( + "carol", + [-1.0, 0.0, 0.0], + "Carol planned a trip to Paris.", + ["travel", "Paris"], + "2026-03-10T09:00:00", + "Paris", + ["Carol"], + ["Rail Europe"], + "travel", + ), + _record( + "apollo-brief", + [0.0, 0.0, 1.0], + "Apollo launch status was discussed.", + ["Apollo", "launch"], + "2026-04-10T09:00:00", + "Houston", + ["Dana"], + ["Project Apollo"], + "space", + ), + ] + + +@pytest.fixture +def milvus_backend(tmp_path, records): + backend = MilvusVectorStoreBackend( + collection_name="memory_entries", + vector_dimension=3, + uri=str(tmp_path / "milvus.db"), + ) + backend.insert(records) + try: + yield backend + finally: + backend.close() + + +def test_milvus_lite_preserves_semantic_and_keyword_score_order(milvus_backend): + semantic_results = milvus_backend.semantic_search([1.0, 0.0, 0.0], top_k=4) + keyword_results = milvus_backend.keyword_search(["Apollo", "budget"], top_k=4) + + assert milvus_backend.semantic_score_order == ScoreOrder.ASCENDING + assert milvus_backend.keyword_score_order == ScoreOrder.DESCENDING + assert [result.entry_id for result in semantic_results] == [ + "coffee", + "apollo", + "apollo-brief", + "carol", + ] + assert [result.score for result in semantic_results] == sorted( + result.score for result in semantic_results + ) + assert keyword_results[0].entry_id == "apollo" + assert {result.entry_id for result in keyword_results} == { + "apollo", + "apollo-brief", + } + assert [result.score for result in keyword_results] == sorted( + (result.score for result in keyword_results), reverse=True + ) + assert all(result.score >= 0 for result in keyword_results) + + +def test_milvus_lite_applies_safe_semantic_filters(milvus_backend): + scalar_results = milvus_backend.semantic_search( + [1.0, 0.0, 0.0], + top_k=4, + filters={"topic": ["finance", "travel"]}, + ) + array_results = milvus_backend.semantic_search( + [1.0, 0.0, 0.0], + top_k=4, + filters={"persons": ["Alice", "Carol"]}, + ) + escaped_results = milvus_backend.semantic_search( + [1.0, 0.0, 0.0], + top_k=4, + filters={"topic": 'finance" or true'}, + ) + + assert [result.entry_id for result in scalar_results] == ["apollo", "carol"] + assert [result.entry_id for result in array_results] == ["coffee", "carol"] + assert escaped_results == [] + + with pytest.raises(ValueError, match="Invalid semantic filter field"): + milvus_backend.semantic_search( + [1.0, 0.0, 0.0], + top_k=4, + filters={"topic or true": "finance"}, + ) + with pytest.raises(TypeError, match="scalar filter.*supports only strings"): + milvus_backend.semantic_search( + [1.0, 0.0, 0.0], + top_k=4, + filters={"topic": {"nested": "value"}}, + ) + + +def test_milvus_lite_supports_structured_search(milvus_backend): + assert { + result.entry_id + for result in milvus_backend.structured_search(persons=["Alice", "Carol"]) + } == {"coffee", "carol"} + assert { + result.entry_id for result in milvus_backend.structured_search(location="Paris") + } == {"apollo", "carol"} + assert { + result.entry_id + for result in milvus_backend.structured_search(entities=["Project Apollo"]) + } == {"apollo", "apollo-brief"} + assert [ + result.entry_id + for result in milvus_backend.structured_search( + timestamp_range=("2026-02-01", "2026-02-28"), + ) + ] == ["apollo"] + assert len(milvus_backend.structured_search(location="Paris", top_k=1)) == 1 + + with pytest.raises(ValueError, match="wildcard"): + milvus_backend.structured_search(location="Paris%") + + +def test_milvus_lite_preserves_metadata_and_lifecycle_operations(milvus_backend): + assert milvus_backend.count() == 4 + results = {result.entry_id: result for result in milvus_backend.get_all()} + assert set(results) == {"coffee", "apollo", "carol", "apollo-brief"} + assert results["apollo"].metadata == { + "lossless_restatement": "The Project Apollo budget budget was approved.", + "keywords": ["Apollo", "budget"], + "timestamp": "2026-02-10T09:00:00", + "location": "Paris", + "persons": ["Bob"], + "entities": ["Project Apollo"], + "topic": "finance", + } + + milvus_backend.optimize() + milvus_backend.clear() + + assert milvus_backend.count() == 0 + assert milvus_backend.get_all() == [] + assert milvus_backend.semantic_search([1.0, 0.0, 0.0], top_k=3) == [] + + +def test_milvus_lite_validates_reused_collection_dimension(tmp_path, records): + uri = str(tmp_path / "reused.db") + backend = MilvusVectorStoreBackend("memory_entries", 3, uri=uri) + backend.insert(records[:1]) + backend.close() + + reopened = MilvusVectorStoreBackend("memory_entries", 3, uri=uri) + assert reopened.count() == 1 + reopened.close() + + with pytest.raises(ValueError, match="vector dimension 3; expected 4"): + MilvusVectorStoreBackend("memory_entries", 4, uri=uri) + + +def test_reused_collection_forwards_configured_consistency_level( + tmp_path, records, monkeypatch +): + uri = str(tmp_path / "consistency.db") + initial = MilvusVectorStoreBackend( + "memory_entries", + 3, + uri=uri, + consistency_level="Bounded", + ) + initial.insert(records) + initial.close() + + class RecordingMilvusClient(MilvusClient): + def __init__(self, *args, **kwargs): + self.read_calls = [] + super().__init__(*args, **kwargs) + + def search(self, *args, **kwargs): + self.read_calls.append(("search", kwargs.get("consistency_level"))) + return super().search(*args, **kwargs) + + def query(self, *args, **kwargs): + self.read_calls.append(("query", kwargs.get("consistency_level"))) + return super().query(*args, **kwargs) + + def query_iterator(self, *args, **kwargs): + self.read_calls.append(("query_iterator", kwargs.get("consistency_level"))) + return super().query_iterator(*args, **kwargs) + + monkeypatch.setattr( + MilvusVectorStoreBackend, + "_load_pymilvus", + staticmethod(lambda: (RecordingMilvusClient, DataType, Function, FunctionType)), + ) + backend = MilvusVectorStoreBackend( + "memory_entries", + 3, + uri=uri, + consistency_level="Session", + ) + try: + assert backend.count() == len(records) + backend.semantic_search([1.0, 0.0, 0.0], top_k=1) + backend.keyword_search(["Apollo"], top_k=1) + backend.structured_search(persons=["Alice"], top_k=1) + backend.get_all() + + methods = [method for method, _ in backend.client.read_calls] + assert methods.count("search") == 2 + assert methods.count("query") >= 4 + assert methods.count("query_iterator") == 1 + assert all( + consistency_level == "Session" + for _, consistency_level in backend.client.read_calls + ) + finally: + backend.close() + + +def test_milvus_lite_3_0_cosine_workaround_is_narrow(): + assert MilvusVectorStoreBackend._is_lite_3_0_cosine_distance("./milvus.db", "3.0.0") + assert not MilvusVectorStoreBackend._is_lite_3_0_cosine_distance( + "./milvus.db", "3.1.0" + ) + assert not MilvusVectorStoreBackend._is_lite_3_0_cosine_distance( + "https://example.api.zillizcloud.com", "3.0.0" + ) + assert not MilvusVectorStoreBackend._is_lite_3_0_cosine_distance( + "localhost:19530", "3.0.0" + ) + + +def test_hybrid_retrieval_uses_milvus_for_all_paths(tmp_path): + entries = [ + MemoryEntry( + entry_id="coffee", + lossless_restatement="Alice drinks espresso at the neighborhood cafe.", + keywords=["coffee", "espresso"], + persons=["Alice"], + topic="coffee", + ), + MemoryEntry( + entry_id="apollo", + lossless_restatement="The Project Apollo budget was approved.", + keywords=["Apollo", "budget"], + persons=["Bob"], + topic="finance", + ), + MemoryEntry( + entry_id="carol", + lossless_restatement="Carol planned a trip to Paris.", + keywords=["travel", "Paris"], + persons=["Carol"], + location="Paris", + topic="travel", + ), + ] + store = VectorStore( + db_path=str(tmp_path / "unused-lancedb"), + table_name="memory_entries", + embedding_model=DeterministicEmbedder(), + backend_factory=lambda dimension: MilvusVectorStoreBackend( + collection_name="memory_entries", + vector_dimension=dimension, + uri=str(tmp_path / "hybrid.db"), + ), + ) + store.add_entries(entries) + retriever = HybridRetriever( + llm_client=DeterministicLLM(), + vector_store=store, + semantic_top_k=1, + keyword_top_k=1, + structured_top_k=1, + enable_planning=True, + enable_reflection=False, + enable_parallel_retrieval=False, + ) + + try: + results = retriever.retrieve("coffee Apollo Carol") + assert [entry.entry_id for entry in results] == [ + "coffee", + "apollo", + "carol", + ] + assert not (tmp_path / "unused-lancedb").exists() + finally: + store.backend.close()