feat: add Milvus vector index provider#404
feat: add Milvus vector index provider#404zc277584121 wants to merge 1 commit intoNevaMind-AI:mainfrom
Conversation
Wire up an external VectorIndex abstraction alongside the existing metadata
store, and ship a Milvus / Milvus Lite / Zilliz Cloud implementation as the
first provider. Users can now opt into Milvus with:
database_config = {
"metadata_store": {"provider": "inmemory"},
"vector_index": {"provider": "milvus"}, # ./milvus.db by default
}
The inmemory metadata backend mirrors create/update/delete into the vector
index and routes vector_search_items through it, preserving scope filters
(user_id, agent_id, ...) via Milvus dynamic fields. Salience ranking still
runs locally because it relies on per-item reinforcement/recency factors.
Scope is additive and non-breaking: default configs keep using the
bruteforce/pgvector paths. pymilvus and milvus-lite are gated behind the
new `[milvus]` optional-dependencies extra.
Adds ADR 0003, docs/milvus.md, and a runnable example under
examples/milvus_vector_index.py.
There was a problem hiding this comment.
Pull request overview
Adds an optional external VectorIndex abstraction and introduces a Milvus-backed provider (Milvus / Milvus Lite / Zilliz Cloud), wiring it into the inmemory metadata backend for write mirroring and similarity search delegation.
Changes:
- Introduce
VectorIndexprotocol +MilvusVectorIndeximplementation and a small factory (build_vector_index). - Wire the in-memory backend to mirror item mutations into the configured vector index and route non-salience searches through it.
- Add Milvus optional dependencies, docs, example, changelog entry, and Milvus Lite–backed tests.
Reviewed changes
Copilot reviewed 16 out of 17 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| uv.lock | Locks new optional Milvus dependency graph (pymilvus/milvus-lite/etc.). |
| tests/test_milvus_vector_index.py | Adds integration tests using Milvus Lite for upsert/search/filter/delete and in-memory repo routing. |
| src/memu/database/vector_index/milvus.py | Implements Milvus-backed VectorIndex with dynamic-field filtering and AUTOINDEX+COSINE search. |
| src/memu/database/vector_index/interfaces.py | Defines the backend-agnostic VectorIndex protocol. |
| src/memu/database/vector_index/init.py | Adds build_vector_index() constructor for the new provider. |
| src/memu/database/inmemory/repositories/memory_item_repo.py | Mirrors create/update/delete/clear to the external index; delegates search (except salience). |
| src/memu/database/inmemory/repo.py | Passes VectorIndex into the in-memory repos and closes it on store close. |
| src/memu/database/inmemory/init.py | Threads vector_index through the in-memory builder. |
| src/memu/database/factory.py | Constructs vector index and passes it into the in-memory backend. |
| src/memu/app/settings.py | Extends VectorIndexConfig to support milvus (uri/token/collection/dim) and defaults uri. |
| pyproject.toml | Adds [milvus] optional extra, deptry ignores, and mypy override for pymilvus. |
| examples/milvus_vector_index.py | Adds runnable end-to-end Milvus Lite example (currently with an incorrect sample file path). |
| docs/milvus.md | Adds Milvus configuration and usage guide. |
| docs/architecture.md | Documents external vector index concept and Milvus provider. |
| docs/adr/0003-external-vector-index.md | Adds ADR for external vector index architecture decision. |
| README.md | Adds quickstart snippet for Milvus extra + example. |
| CHANGELOG.md | Adds Unreleased feature note for Milvus external vector index. |
Comments suppressed due to low confidence (1)
src/memu/database/factory.py:53
build_database()constructs the external vector index unconditionally, but only theinmemorybackend receives/owns it. Forsqlite/postgresthis will instantiate (and potentially connect to) Milvus and then leak it unused, andvector_index.provider="milvus"will silently have no effect. Consider only building/passing the index for backends that support it today, or raising a clear error whenmilvusis configured with an unsupported metadata provider (or ensure the created index is always closed).
vector_index = build_vector_index(config.vector_index)
provider = config.metadata_store.provider
if provider == "inmemory":
return build_inmemory_database(
config=config,
user_model=user_model,
vector_index=vector_index,
)
elif provider == "postgres":
# Lazy import to avoid requiring pgvector when not using postgres
from memu.database.postgres import build_postgres_database
return build_postgres_database(config=config, user_model=user_model)
elif provider == "sqlite":
# Lazy import to avoid loading SQLite dependencies when not needed
from memu.database.sqlite import build_sqlite_database
return build_sqlite_database(config=config, user_model=user_model)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def _build_filter_expr(where: Mapping[str, Any] | None) -> str: | ||
| if not where: | ||
| return "" | ||
| parts: list[str] = [] | ||
| for key, value in where.items(): | ||
| literal = _format_scalar(value) | ||
| if literal is None: | ||
| logger.debug("Skipping unsupported filter key %s=%r for Milvus", key, value) | ||
| continue | ||
| parts.append(f"{key} == {literal}") | ||
| return " and ".join(parts) |
There was a problem hiding this comment.
_build_filter_expr() interpolates where keys directly into the Milvus boolean expression. Because where ultimately comes from user-provided filters, this allows expression injection / invalid field names (e.g. keys containing spaces, operators, or parentheses). Validate/whitelist field names (e.g. ^[A-Za-z_][A-Za-z0-9_]*$) and fail closed (return no hits or raise) when filters can’t be safely represented, rather than silently broadening the query.
| def _ensure_collection(self, dim: int) -> None: | ||
| with self._lock: | ||
| if self._ready: | ||
| return | ||
| if self._client.has_collection(collection_name=self._collection_name): | ||
| self._ready = True | ||
| self._dim = dim | ||
| return | ||
|
|
||
| from pymilvus import DataType | ||
|
|
||
| schema = self._client.create_schema( | ||
| auto_id=False, | ||
| enable_dynamic_field=True, | ||
| ) | ||
| schema.add_field( | ||
| field_name=_ID_FIELD, | ||
| datatype=DataType.VARCHAR, | ||
| is_primary=True, | ||
| max_length=_DEFAULT_ID_MAX_LENGTH, | ||
| ) | ||
| schema.add_field( | ||
| field_name=_VECTOR_FIELD, | ||
| datatype=DataType.FLOAT_VECTOR, | ||
| dim=dim, | ||
| ) | ||
|
|
||
| index_params = self._client.prepare_index_params() | ||
| index_params.add_index( | ||
| field_name=_VECTOR_FIELD, | ||
| index_type="AUTOINDEX", | ||
| metric_type="COSINE", | ||
| ) | ||
|
|
||
| self._client.create_collection( | ||
| collection_name=self._collection_name, | ||
| schema=schema, | ||
| index_params=index_params, | ||
| ) | ||
| self._dim = dim | ||
| self._ready = True | ||
|
|
||
| def upsert( | ||
| self, | ||
| item_id: str, | ||
| vector: list[float], | ||
| scope: Mapping[str, Any] | None = None, | ||
| ) -> None: | ||
| if not vector: | ||
| logger.debug("Skipping Milvus upsert for %s: empty vector", item_id) | ||
| return | ||
| self._ensure_collection(len(vector)) | ||
|
|
There was a problem hiding this comment.
Dimension handling is inconsistent: _ensure_collection() returns early when _ready is already true, so subsequent upsert() calls with a different vector length won’t be detected until Milvus errors (or worse, if a pre-existing collection has a different dim). Consider storing/validating self._dim and raising a clear error when len(vector) (or provided dim) doesn’t match the established collection dimension.
| def delete(self, item_id: str) -> None: | ||
| if not self._ready: | ||
| return | ||
| self._client.delete( | ||
| collection_name=self._collection_name, | ||
| filter=f'{_ID_FIELD} == "{item_id}"', | ||
| ) | ||
|
|
||
| def delete_many(self, item_ids: Iterable[str]) -> None: | ||
| ids = list(item_ids) | ||
| if not ids or not self._ready: | ||
| return | ||
| quoted = ", ".join(f'"{i}"' for i in ids) | ||
| self._client.delete( | ||
| collection_name=self._collection_name, | ||
| filter=f"{_ID_FIELD} in [{quoted}]", | ||
| ) |
There was a problem hiding this comment.
delete() / delete_many() build Milvus filter expressions by inserting raw item_id values into quoted strings. If an ID ever contains quotes/backslashes, the expression can break or be exploited similarly to injection. Prefer escaping IDs (reuse the string escaping used in _format_scalar) or use an API variant that deletes by primary keys without constructing an expression string.
| ## Supported Combinations | ||
|
|
||
| The Milvus vector index can be paired with any metadata store. memU keeps the metadata record (summary, categories, scope, reinforcement stats, ...) in the metadata store and mirrors embeddings into Milvus on create / update / delete. | ||
|
|
||
| - **Available now**: `inmemory` metadata store + Milvus vector index (feature-complete). | ||
| - **Planned**: `sqlite` and `postgres` metadata stores wired to Milvus. See `docs/adr/0003-external-vector-index.md` for the rollout plan. |
There was a problem hiding this comment.
This section says the Milvus vector index “can be paired with any metadata store”, but in code only the inmemory backend is currently wired to pass a VectorIndex instance; sqlite/postgres ignore provider="milvus" today. Consider adjusting the wording to match the actual supported combinations (or documenting the current limitation more explicitly).
| file_path = os.path.abspath("example/example_conversation.json") | ||
| if not Path(file_path).exists(): | ||
| msg = ( | ||
| f"Example conversation not found at {file_path}. " | ||
| "Run from the repository root so the 'example/' folder is visible." |
There was a problem hiding this comment.
The example looks for example/example_conversation.json, but this repository doesn’t have an example/ directory; the sample conversation JSON lives under tests/example/ and there are also sample conversations under examples/resources/. As written, the script will exit immediately; update the path/message to point at an existing file location.
| file_path = os.path.abspath("example/example_conversation.json") | |
| if not Path(file_path).exists(): | |
| msg = ( | |
| f"Example conversation not found at {file_path}. " | |
| "Run from the repository root so the 'example/' folder is visible." | |
| file_path = os.path.abspath("tests/example/example_conversation.json") | |
| if not Path(file_path).exists(): | |
| msg = ( | |
| f"Example conversation not found at {file_path}. " | |
| "Run from the repository root so 'tests/example/example_conversation.json' is available." |
| `database_config.vector_index` can point at an external index that lives outside the metadata store. Providers: | ||
|
|
||
| - `bruteforce` / `pgvector`: default, handled inside the metadata backend (see above). | ||
| - `milvus`: delegates similarity search to Milvus / Milvus Lite / Zilliz Cloud via the `VectorIndex` protocol in `memu/database/vector_index/`. Metadata backends mirror embedding writes through this interface on create/update/delete; searches route through it except for `ranking="salience"`, which still runs locally. See `docs/milvus.md` and `docs/adr/0003-external-vector-index.md`. |
There was a problem hiding this comment.
The milvus bullet implies all metadata backends mirror writes and route searches through VectorIndex, but in code only the inmemory backend is currently wired to receive a VectorIndex instance; sqlite/postgres ignore it today. Consider tightening this wording (e.g., “currently only inmemory mirrors…”) to avoid suggesting functionality that isn’t implemented yet.
| - `milvus`: delegates similarity search to Milvus / Milvus Lite / Zilliz Cloud via the `VectorIndex` protocol in `memu/database/vector_index/`. Metadata backends mirror embedding writes through this interface on create/update/delete; searches route through it except for `ranking="salience"`, which still runs locally. See `docs/milvus.md` and `docs/adr/0003-external-vector-index.md`. | |
| - `milvus`: delegates similarity search to Milvus / Milvus Lite / Zilliz Cloud via the `VectorIndex` protocol in `memu/database/vector_index/`. Currently, only the `inmemory` metadata backend mirrors embedding writes through this interface on create/update/delete and routes searches through it; `ranking="salience"` still runs locally. See `docs/milvus.md` and `docs/adr/0003-external-vector-index.md`. |
| def build_vector_index(config: VectorIndexConfig | None) -> VectorIndex | None: | ||
| """Construct an external VectorIndex instance if configured. | ||
|
|
||
| Returns None for providers that do not require a separate index | ||
| (``bruteforce``, ``pgvector``, ``none``) because those are handled | ||
| inside the metadata backend itself. | ||
| """ | ||
| if config is None: | ||
| return None | ||
| provider = config.provider | ||
| if provider == "milvus": | ||
| from memu.database.vector_index.milvus import MilvusVectorIndex | ||
|
|
||
| return MilvusVectorIndex( | ||
| uri=config.uri or "./milvus.db", | ||
| token=config.token, | ||
| collection_name=config.collection_name, | ||
| dim=config.dim, | ||
| ) | ||
| return None |
There was a problem hiding this comment.
build_vector_index() returns None for provider="none", but the in-memory repository will then fall back to local cosine search, so vector search is not actually disabled. If none is meant to disable search, consider returning a no-op VectorIndex whose search() returns [] (and whose writes are no-ops), or plumb the provider into repositories so they can short-circuit search when configured.
📝 Pull Request Summary
Wire up an optional
VectorIndexabstraction next to the existing metadata store, and ship a Milvus / Milvus Lite / Zilliz Cloud implementation as the first provider. The inmemory metadata backend now mirrors writes into the external index and routes similarity search through it.The change is fully additive: default configs keep using
bruteforce/pgvector, andpymilvus/milvus-litelive behind a new[milvus]extra.✅ What does this PR do?
memu/database/vector_index/with aVectorIndexprotocol (upsert/delete/delete_many/search/close) and aMilvusVectorIndeximplementation built onMilvusClient(usesAUTOINDEX+COSINE, dynamic fields for scope filters, and a./milvus.dbMilvus Lite default for zero-setup local).VectorIndexConfig.providernow accepts"milvus"withuri,token,collection_name,dimfields;DatabaseConfig.model_post_initdefaultsurito./milvus.dbwhen left blank.build_databaseconstructs the index once and passes it into the inmemory backend;InMemoryStore.close()closes the index.InMemoryMemoryItemRepositorymirrorscreate_item/create_item_reinforce/update_item/delete_item/clear_itemsto the vector index and delegatesvector_search_itemsto it (exceptranking="salience", which still needs per-item reinforcement/recency factors and runs locally).[project.optional-dependencies].milvus:pymilvus,milvus-lite, and a conditionalsetuptools<81pin for Python 3.13 (works around milvus-lite <=2.5.1 still importing the removedpkg_resourcesAPI).tests/test_milvus_vector_index.pycovering upsert/search, scope filtering, delete/delete_many, end-to-end through the inmemory repo, andclear_itemspropagation. Real Milvus Lite, not mocks.docs/milvus.mduser guide,docs/adr/0003-external-vector-index.mdarchitecture decision,docs/architecture.mdsection,README.mdquickstart snippet,examples/milvus_vector_index.pyrunnable end-to-end example.🤔 Why is this change needed?
settings.pyalready modelledmetadata_storeandvector_indexas orthogonal concerns, but every shipped backend kept embeddings inside its own tables. Two footprints were unserved:Milvus Lite also preserves the zero-setup local story the
inmemorybackend already promises — no Docker, no separate service, just a single file.🔍 Type of Change
✅ PR Quality Checklist
feat:)docs/milvus.md, ADR 0003,architecture.md,README.md, example, CHANGELOGUnreleased📌 Optional
VectorIndexinto thesqliteandpostgresmetadata backends (intentionally out of scope for this PR to keep it reviewable).Local verification