Add Milvus vector store backend#876
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (16)
🚧 Files skipped from review as they are similar to previous changes (7)
WalkthroughAdds Milvus Lite and remote Milvus support as a vector-store backend, including configuration fields, collection operations, startup dimension validation, factory integration, documentation, dependency updates, and unit/integration tests. ChangesMilvus vector store support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Application
participant MilvusVectorStore
participant MilvusClient
participant EmbeddingValidator
Application->>MilvusVectorStore: upsert_many(namespace, vectors)
MilvusVectorStore->>MilvusClient: create or validate collection
MilvusVectorStore->>MilvusClient: upsert records
Application->>MilvusVectorStore: query(namespace, embedding, filters)
MilvusVectorStore->>MilvusClient: search with COSINE parameters
MilvusClient-->>MilvusVectorStore: ordered hits and metadata
MilvusVectorStore-->>Application: VectorQueryResult list
EmbeddingValidator->>MilvusVectorStore: probe_namespace_dim(namespace)
MilvusVectorStore->>MilvusClient: describe collection
MilvusClient-->>EmbeddingValidator: declared vector dimension
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/vector_store/milvus.py`:
- Around line 320-398: The query logic in Milvus uses raw search scores as if
they were cosine distance, which makes the max_distance filter and
VectorQueryResult.score inconsistent. Update the async query path in query() to
convert the Milvus score returned from self.client.search into cosine distance
before comparing against max_distance and before assigning score, or adjust the
contract consistently throughout this method. Use the existing query(),
search_kwargs, and VectorQueryResult construction points to apply the conversion
in one place.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 77bb082b-68a7-48d7-aa40-5d891ac05abe
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (12)
.env.templateREADME.mdconfig.toml.exampledocs/v3/contributing/configuration.mdxpyproject.tomlsrc/config.pysrc/startup/embedding_validator.pysrc/vector_store/__init__.pysrc/vector_store/milvus.pytests/conftest.pytests/llm/test_model_config.pytests/vector_store/test_milvus.py
|
Hi @VVoruganti @Rajat-Ahuja1997, could you take a look when you have a chance? This PR adds a Milvus vector store backend following the existing VectorStore interface. I also pushed a follow-up for the score/distance handling and added focused coverage; the current review check is green. |
cfae284 to
fecb346
Compare
VVoruganti
left a comment
VVoruganti
left a comment
There was a problem hiding this comment.
Thanks for this — the backend itself is solid work. I verified parity rather than taking the tests at face value, and it holds up:
- Score convention is correct. I probed
MilvusClient.searchwithmetric_type="COSINE"directly: hits come back as{id, distance, entity}anddistanceis true cosine distance (0.0 / 1.0 / 2.0 for identical / orthogonal / opposite). That matches lancedb's_distanceand turbopuffer's$dist, so ordering andmax_distancethresholding behave identically to the existing backends. - Filters cover every caller shape. Callers only ever pass
leveleq,session_nameeq or{"in": [...]}, andpeer_nameeq (crud/document.py:242,crud/message.py:582,utils/search.py:104,dreamer/surprisal.py:264). All supported, and this backend is stricter than lancedb on malformed operators (clearValueErrorvs. emitting broken SQL) — an improvement. include_attributesTrue/False/["message_id"]all behave like the other backends, the last via the dynamic field.- Clean on ruff, ruff format, and basedpyright (0 errors); 21/21 tests in
tests/vector_store/pass.
Two things I'd want fixed before merge, plus some cleanup. Details inline.
Blocking
MILVUS_CONSISTENCY_LEVELdefault is unsafe on a real Milvus server — see inline onsrc/config.py. This is the one genuine capability gap versus the other backends.- The
"score"branch in_hit_cosine_distanceis unreachable, and the test that covers it asserts behavior no real backend produces — see inline onsrc/vector_store/milvus.pyandtests/vector_store/test_milvus.py.
Cleanup
- Two changes unrelated to Milvus (
tests/conftest.py,tests/llm/test_model_config.py) — inline. - Docs and changelog are half-done. There are zero Milvus mentions in
CHANGELOG.md,docs/changelog/introduction.mdx,CLAUDE.md(the project tree around lines 246-247 and architectural decision #8 at line 289 still say "turbopuffer or lancedb"),docs/v3/contributing/changing-embeddings.mdx:60,scripts/configure_embeddings.py:19-20, anddocs/v3/contributing/configuration.mdx:285and:295. The repo has a clear convention here — "External Vector Store support for turbopuffer and lancedb" got its own CHANGELOG entry, so a third backend should too. The.env.template/config.toml.example/configuration.mdxenv-var updates you did make are all correct; it's just these prose surfaces that were missed.
Everything else in the non-Milvus part of the diff is exactly the mechanical widening you'd expect, and declaring pymilvus as a non-optional core dependency is consistent with how turbopuffer and lancedb are already declared.
The remaining inline notes are non-blocking design observations, not merge conditions.
| MILVUS_URI: str = "./milvus.db" | ||
| MILVUS_TOKEN: str | None = None | ||
| MILVUS_DB_NAME: str | None = None | ||
| MILVUS_CONSISTENCY_LEVEL: ( |
There was a problem hiding this comment.
Blocking: defaulting this to None means a Milvus server or Zilliz Cloud collection is created with Milvus's own default consistency level, Bounded — a just-upserted vector can be invisible to search for several seconds.
That's a real behavioral divergence from the other three backends: pgvector, lancedb, and turbopuffer are all read-your-writes. Honcho's reconciler writes embeddings and search reads them back, so this shows up as silently missing recall rather than an error.
It's double-masked right now:
- the default
MILVUS_URIis Milvus Lite, which is effectively strong, so local dev never sees it; test_milvus_lite_round_trippinsMILVUS_CONSISTENCY_LEVELto"Strong", so the tested config isn't the default config.
There's also a second-order problem: the value is only passed to create_collection, so it's baked in at collection-creation time. An operator who hits stale reads in production can't fix it by setting the env var — existing collections keep whatever default they were born with, and there's no migration path short of dropping them.
Suggest defaulting to "Session" (read-your-writes per client, which is the semantics the other backends provide) rather than None.
| return float(hit["distance"]) | ||
| if "score" in hit: | ||
| return 1.0 - float(hit["score"]) | ||
| return 0.0 |
There was a problem hiding this comment.
Blocking (small): the "score" branch is unreachable. MilvusClient.search returns hits shaped {id, distance, entity} — there is no score key. I confirmed against pymilvus directly:
[{'id': 'same', 'distance': 0.0, ...},
{'id': 'orth', 'distance': 1.0, ...},
{'id': 'opp', 'distance': 2.0, ...}]
So distance is already true cosine distance and the first branch is the only one that ever runs — which is correct and matches lancedb/turbopuffer. The 1.0 - score line is dead code guarding against a response shape this SDK doesn't produce.
Worth deleting along with _ScoreOnlySearchClient: as written, the two tests assert contradictory scoring conventions for an identical-vector hit (0.0 in the round-trip test, 0.1 from the fake) and both pass, so a future reader has no way to tell which one describes reality. The round-trip test against real Milvus Lite already covers scoring properly.
The return 0.0 fallback is fine to keep — lancedb and turbopuffer both default a missing distance to 0.0 the same way.
| "entity": {"message_id": "msg_2"}, | ||
| }, | ||
| ] | ||
| ] |
There was a problem hiding this comment.
This fake emits a score key that MilvusClient.search never returns (real hits are {id, distance, entity}, with distance already being cosine distance). So test_query_converts_milvus_score_to_cosine_distance certifies a code path that can't execute against pymilvus.
Suggest removing this class and that test together with the dead branch in _hit_cosine_distance — test_milvus_lite_round_trip already asserts the real scoring behavior (all_results[0].score == 0.0 for the identical vector) against actual Milvus Lite, which is the stronger check.
| "tests/unified/", | ||
| "tests/live_llm/", | ||
| # Vector-store provider tests exercise SDK behavior directly and do not need DB setup. | ||
| "tests/vector_store/", |
There was a problem hiding this comment.
This change isn't needed for the Milvus tests, and it affects pre-existing tests.
I verified by reverting just these two lines and re-running: tests/vector_store/test_milvus.py passes 6/6 without it. Since tests/vector_store/ also holds test_lancedb.py, test_turbopuffer.py, and test_namespace_dim_probes.py, adding the directory to the blocklist silently changes fixture behavior for those existing tests as a side effect of a PR about Milvus.
Suggest dropping it from this PR. If skipping DB setup for the whole tests/vector_store/ directory is worth doing on its own merits, it'd be easier to review as a separate change.
| assert any( | ||
| "VECTOR_STORE_DIMENSIONS is deprecated" in m for m in messages | ||
| ), f"expected deprecation warning, got {messages!r}" | ||
| assert any("VECTOR_STORE_DIMENSIONS is deprecated" in m for m in messages), ( |
There was a problem hiding this comment.
This reformat is unrelated to Milvus — it's ruff format drift that's already latent on main. Running ruff format (0.15.12) against main's copy of this file produces exactly this hunk, so it isn't something your change introduced; it just got swept in.
Harmless, but it's noise in a backend PR and it'll reappear for the next contributor. Worth pulling out into its own formatting commit.
(The Literal widening further down in this same file is correctly part of this PR — no issue there.)
| for key in STANDARD_METADATA_FIELDS: | ||
| row[key] = self._metadata_varchar_value(metadata.get(key)) | ||
|
|
||
| return row |
There was a problem hiding this comment.
Non-blocking design note: every metadata value gets stored twice — once inside the JSON METADATA_FIELD, and again as a top-level field (explicit column for the six STANDARD_METADATA_FIELDS, dynamic field for everything else).
Since the collection is created with enable_dynamic_field=True, the dynamic field already captures arbitrary keys, which makes the JSON metadata column redundant — roughly 2x metadata storage. It also creates a precedence rule in _entity_metadata that's invisible until the two copies disagree: JSON is merged first, then entity fields overwrite it, so the str()-coerced copy always wins for those six keys.
Related, two smaller things in the same area:
- The
str()coercion in_metadata_varchar_valuesilently changes types round-trip for those six fields. Harmless today (level,session_name, andpeer_nameare all strings), but it's a latent trap for any future non-string value. message_id— the only field actually projected on the hot message-search path (crud/message.pyandutils/search.pyboth passinclude_attributes=["message_id"]) — isn't among the declared fields, so it lands in the dynamic field, while six mostly-unused VARCHAR columns are declared explicitly. It works, but it's backwards from an indexing standpoint. lancedb takes the other approach in_metadata_fields_for_namespace, declaringmessage_idformsgnamespaces.
Also minor: MAX_VARCHAR_LENGTH = 65535 on the primary key is generous for Honcho's ~21-char nanoid IDs.
| if settings.VECTOR_STORE.MILVUS_DB_NAME: | ||
| client_kwargs["db_name"] = settings.VECTOR_STORE.MILVUS_DB_NAME | ||
|
|
||
| self.client = MilvusClient(**client_kwargs) |
There was a problem hiding this comment.
Non-blocking: constructing the client eagerly in __init__ means that with a Milvus server / Zilliz URI, a transient outage at first use surfaces as a raw MilvusException propagating out of get_external_vector_store() — not the VectorStoreError that callers of this module are written to expect.
For comparison, LanceDBVectorStore defers connecting until _get_db(), and turbopuffer's httpx client connects lazily. Since @cache doesn't memoize exceptions this does self-heal on the next call, so it's an error-shape wart rather than a durable failure. Fine to leave; just noting it since Milvus Lite (the default) never exercises this path.
VVoruganti
left a comment
VVoruganti
left a comment
There was a problem hiding this comment.
Converting my previous review to a formal request for changes — the inline comments there still hold and I haven't duplicated them here.
To restate the assessment: the backend is solid and has real parity with the existing stores. I verified the score convention against pymilvus directly (distance is already true cosine distance, matching lancedb's _distance and turbopuffer's $dist), confirmed every caller filter shape is supported, and it's clean on ruff / ruff format / basedpyright with 21/21 vector-store tests passing. Happy to approve once the two blocking items are addressed.
Blocking:
-
MILVUS_CONSISTENCY_LEVELdefaults toNone(inline) → real Milvus/Zilliz collections get Milvus'sBoundeddefault, so a just-upserted vector can be invisible to search for seconds. pgvector, lancedb, and turbopuffer are all read-your-writes, so this is a genuine capability gap rather than a preference. It's invisible in CI and local dev because the defaultMILVUS_URIis Milvus Lite (effectively strong), andtest_milvus_lite_round_trippins"Strong"— so it would surface only in production, as missing recall rather than an error. Compounding it, the value is only passed tocreate_collection, so it's fixed at collection-creation time and can't be corrected later via config. Suggest defaulting to"Session". -
Dead
"score"branch in_hit_cosine_distance, plus a test that certifies it.MilvusClient.searchnever returns ascorekey, so that branch can't execute;_ScoreOnlySearchClientasserts a response shape the SDK doesn't produce. The two tests currently assert contradictory scoring for an identical-vector hit (0.0real vs.0.1from the fake) and both pass, which leaves no way to tell which describes reality. Deleting the branch and the fake is enough —test_milvus_lite_round_tripalready covers scoring against real Milvus Lite.
Non-blocking, and I don't want these to hold up the merge: the docs/changelog gaps (item 4 in the previous review) are worth picking up since the repo has an established convention for new backends, and the remaining inline notes are design observations only.
Signed-off-by: Cheney Zhang <chen.zhang@zilliz.com>
fecb346 to
21a3076
Compare
2 participants
Summary
Testing
Notes
uv run pytest tests/startup/test_embedding_validator.py -q; it is blocked locally because Postgres at localhost:5432 is not running.Summary by CodeRabbit
New Features
Documentation
Tests