Add Milvus as a first-class Vektori storage backend#49
Merged
Conversation
add MilvusBackend with facts/sentences/episodes collections and user partition-key isolation wire Milvus into storage resolution (storage_backend='milvus' and :19530 URL heuristic) add Milvus support to config/CLI/docs and Docker Compose local stack (etcd, minio, milvus) add optional dependency extras for Milvus (pymilvus) add unit + integration + e2e tests for Milvus backend behavior harden Milvus operations for reliability (nullable string normalization, async loop fallback, read-after-write flush) include Ollama+Milvus example as sanity-check demo (non-core to backend feature)
📝 WalkthroughWalkthroughAdds Milvus as a new async storage backend: implementation of Changes
Sequence Diagram(s)sequenceDiagram
actor User as Client/User
participant Vektori as Vektori API
participant Embedder as Embedder
participant MilvusBackend as MilvusBackend
participant Milvus as Milvus Server
rect rgba(100, 200, 100, 0.5)
Note over User,Milvus: Initialization
User->>Vektori: create(storage_backend="milvus")
Vektori->>MilvusBackend: instantiate(url, prefix, embedding_dim)
Vektori->>Embedder: create()
Vektori->>MilvusBackend: set_sentence_embedder(embedder)
Vektori->>MilvusBackend: await initialize()
MilvusBackend->>Milvus: connect & ensure_collections()
end
rect rgba(100, 150, 200, 0.5)
Note over User,Milvus: Ingest (add)
User->>Vektori: await add(messages...)
Vektori->>Embedder: embed(messages) -> vectors
Vektori->>MilvusBackend: upsert_sentences(messages, vectors)
MilvusBackend->>Milvus: upsert into sentences collection
Milvus-->>MilvusBackend: ack
MilvusBackend-->>Vektori: inserted_count
Vektori-->>User: add result
end
rect rgba(200, 150, 100, 0.5)
Note over User,Milvus: Search
User->>Vektori: await search(query, depth="l1")
Vektori->>Embedder: embed(query) -> vector
Vektori->>MilvusBackend: search_sentences(vector, user_id)
MilvusBackend->>Milvus: search with filter_expr
Milvus-->>MilvusBackend: matched rows
MilvusBackend->>Vektori: normalized {sentences,facts,episodes}
Vektori-->>User: search result
end
rect rgba(200, 100, 100, 0.5)
Note over User,Milvus: Teardown
User->>Vektori: await close()
Vektori->>MilvusBackend: await close()
MilvusBackend->>Milvus: close client
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
There was a problem hiding this comment.
Pull request overview
This PR adds Milvus as an optional, production-oriented storage backend for Vektori, including backend implementation, configuration/factory/CLI wiring, local Docker runtime support, documentation updates, and Milvus-specific unit + integration/e2e tests.
Changes:
- Implement
MilvusBackendwith three collections (facts/sentences/episodes) and session markers stored as typed records. - Wire Milvus into backend factory selection, config defaults/help text, and README usage guidance.
- Add Docker Compose services for a local Milvus stack (etcd + minio + milvus) and introduce Milvus-focused tests and an Ollama sanity example.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| vektori/storage/milvus.py | New Milvus backend implementation (schema, CRUD/search, session markers, flush strategy). |
| vektori/storage/factory.py | Routes storage_backend="milvus" and adds URL heuristic for port 19530. |
| vektori/config.py | Documents "milvus" as a valid storage_backend. |
| vektori/cli.py | Updates CLI help + adds Milvus URL example and config example. |
| tests/unit/test_new_backends_factory.py | Adds factory routing + constructor default tests for Milvus. |
| tests/unit/test_milvus_backend.py | Adds mocked unit tests for key MilvusBackend behaviors. |
| tests/integration/test_milvus_backend.py | Adds live-backend integration tests (auto-skip if Milvus unavailable). |
| tests/integration/test_milvus_e2e.py | Adds end-to-end pipeline tests against live Milvus with mocked providers. |
| README.md | Adds Milvus install + usage examples and updates backend list. |
| pyproject.toml | Adds milvus extra (pymilvus) and includes it in all. |
| examples/ollama_milvus_session_demo.py | Adds local Ollama + Milvus demo script. |
| docker-compose.yml | Adds etcd/minio/milvus services and volumes for local Milvus runtime. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+288
to
+297
| if filter_expr: | ||
| for key in ("filter", "expr"): | ||
| try: | ||
| rows = await self._call("query", **{**kwargs, key: filter_expr}) | ||
| return self._normalize_query_rows(rows) | ||
| except TypeError: | ||
| continue | ||
| rows = await self._call("query", **kwargs) | ||
| return self._normalize_query_rows(rows) | ||
|
|
There was a problem hiding this comment.
In _query(), if filter_expr is provided but neither filter nor expr keyword is accepted, the fallback call executes query without any filter, which can silently return unfiltered rows (cross-user/record_type leakage and incorrect results). Instead of querying without a filter, fail fast (raise) or add an explicit compatibility path that still applies the filter (e.g., detect the correct parameter name once and cache it).
Comment on lines
+457
to
+463
| for key in ("filter", "expr"): | ||
| try: | ||
| await self._call("delete", collection_name=collection_name, **{key: filter_expr}) | ||
| return | ||
| except TypeError: | ||
| continue | ||
| await self._call("delete", collection_name=collection_name) |
There was a problem hiding this comment.
_delete() falls back to calling delete without passing filter_expr if neither filter nor expr keyword is accepted. That risks deleting the entire collection when the client API signature differs. Safer behavior is to raise an error (or detect/correct the parameter name) rather than performing an unfiltered delete.
Suggested change
| for key in ("filter", "expr"): | |
| try: | |
| await self._call("delete", collection_name=collection_name, **{key: filter_expr}) | |
| return | |
| except TypeError: | |
| continue | |
| await self._call("delete", collection_name=collection_name) | |
| last_error: TypeError | None = None | |
| for key in ("filter", "expr"): | |
| try: | |
| await self._call("delete", collection_name=collection_name, **{key: filter_expr}) | |
| return | |
| except TypeError as e: | |
| last_error = e | |
| continue | |
| raise TypeError( | |
| "Milvus client delete() does not accept a supported filter argument " | |
| f"for filtered delete on collection {collection_name!r}; " | |
| "refusing to perform an unfiltered delete." | |
| ) from last_error |
Comment on lines
+726
to
+733
| if not sentences: | ||
| return 0 | ||
|
|
||
| rows: list[dict[str, Any]] = [] | ||
|
|
||
| for sent, emb in zip(sentences, embeddings): | ||
| if self.embedding_dim and len(emb) != self.embedding_dim: | ||
| raise ValueError( |
There was a problem hiding this comment.
upsert_sentences() iterates with zip(sentences, embeddings), which silently truncates if the lengths differ. Consider validating len(sentences) == len(embeddings) up front and raising a clear error to avoid dropping sentences or embeddings without noticing.
Comment on lines
+1152
to
+1164
|
|
||
| for fid in fact_ids: | ||
| row = await self._get_one( | ||
| self._facts_col, | ||
| f'id == "{_escape(fid)}" and record_type == "{_FACT_RECORD}"', | ||
| ["source_sentence_ids", "id"], | ||
| ) | ||
| if not row: | ||
| continue | ||
| for sid in _csv_to_ids(row.get("source_sentence_ids")): | ||
| if sid not in seen: | ||
| seen.add(sid) | ||
| result.append(sid) |
There was a problem hiding this comment.
get_source_sentences() issues one _get_one() query per fact ID. For larger retrievals this becomes N+1 round trips and can dominate latency. Consider batching (e.g., a single query using an id in [...] / IN-style expression, or a client method that retrieves by primary keys) and then merging/deduping in memory.
Suggested change
| for fid in fact_ids: | |
| row = await self._get_one( | |
| self._facts_col, | |
| f'id == "{_escape(fid)}" and record_type == "{_FACT_RECORD}"', | |
| ["source_sentence_ids", "id"], | |
| ) | |
| if not row: | |
| continue | |
| for sid in _csv_to_ids(row.get("source_sentence_ids")): | |
| if sid not in seen: | |
| seen.add(sid) | |
| result.append(sid) | |
| batch_size = 200 | |
| for start in range(0, len(fact_ids), batch_size): | |
| batch_ids = fact_ids[start : start + batch_size] | |
| id_expr = " or ".join([f'id == "{_escape(fid)}"' for fid in batch_ids]) | |
| rows = await self._query_all( | |
| collection_name=self._facts_col, | |
| filter_expr=f'({id_expr}) and record_type == "{_FACT_RECORD}"', | |
| output_fields=["source_sentence_ids", "id"], | |
| batch_size=batch_size, | |
| ) | |
| for row in rows: | |
| for sid in _csv_to_ids(row.get("source_sentence_ids")): | |
| if sid not in seen: | |
| seen.add(sid) | |
| result.append(sid) |
Comment on lines
+1168
to
+1185
| async def get_sentences_by_ids(self, sentence_ids: list[str]) -> list[dict[str, Any]]: | ||
| if not sentence_ids: | ||
| return [] | ||
|
|
||
| rows: list[dict[str, Any]] = [] | ||
| for sid in sentence_ids: | ||
| row = await self._get_one( | ||
| self._sentences_col, | ||
| ( | ||
| f'id == "{_escape(sid)}" and ' | ||
| f'record_type == "{_SENTENCE_RECORD}" and ' | ||
| "is_active == true" | ||
| ), | ||
| self._sentence_output_fields(), | ||
| ) | ||
| if row: | ||
| rows.append(self._to_sentence(row)) | ||
|
|
There was a problem hiding this comment.
get_sentences_by_ids() queries Milvus one ID at a time via _get_one(), which is an N+1 pattern. If Milvus supports an id in [...] filter or bulk get by primary key, batching this into a single call would significantly reduce latency for L1/L2 retrieval paths.
Suggested change
| async def get_sentences_by_ids(self, sentence_ids: list[str]) -> list[dict[str, Any]]: | |
| if not sentence_ids: | |
| return [] | |
| rows: list[dict[str, Any]] = [] | |
| for sid in sentence_ids: | |
| row = await self._get_one( | |
| self._sentences_col, | |
| ( | |
| f'id == "{_escape(sid)}" and ' | |
| f'record_type == "{_SENTENCE_RECORD}" and ' | |
| "is_active == true" | |
| ), | |
| self._sentence_output_fields(), | |
| ) | |
| if row: | |
| rows.append(self._to_sentence(row)) | |
| async def _get_sentences_rows_by_ids( | |
| self, sentence_ids: list[str], batch_size: int = 100 | |
| ) -> list[dict[str, Any]]: | |
| if not sentence_ids: | |
| return [] | |
| seen_ids: set[str] = set() | |
| ordered_ids: list[str] = [] | |
| for sid in sentence_ids: | |
| if sid and sid not in seen_ids: | |
| seen_ids.add(sid) | |
| ordered_ids.append(sid) | |
| rows: list[dict[str, Any]] = [] | |
| for start in range(0, len(ordered_ids), batch_size): | |
| batch_ids = ordered_ids[start : start + batch_size] | |
| id_filter = " or ".join([f'id == "{_escape(sid)}"' for sid in batch_ids]) | |
| filter_expr = ( | |
| f"({id_filter}) and " | |
| f'record_type == "{_SENTENCE_RECORD}" and ' | |
| "is_active == true" | |
| ) | |
| rows.extend( | |
| await self._query_all( | |
| collection_name=self._sentences_col, | |
| filter_expr=filter_expr, | |
| output_fields=self._sentence_output_fields(), | |
| batch_size=batch_size, | |
| ) | |
| ) | |
| return rows | |
| async def get_sentences_by_ids(self, sentence_ids: list[str]) -> list[dict[str, Any]]: | |
| if not sentence_ids: | |
| return [] | |
| rows = [ | |
| self._to_sentence(row) | |
| for row in await self._get_sentences_rows_by_ids(sentence_ids) | |
| ] |
| ETCD_SNAPSHOT_COUNT: "50000" | ||
| command: | ||
| - /usr/local/bin/etcd | ||
| - -advertise-client-urls=http://127.0.0.1:2379 |
There was a problem hiding this comment.
etcd is configured with -advertise-client-urls=http://127.0.0.1:2379. In a Compose network this typically causes other containers (like milvus) to receive/learn a loopback endpoint that isn’t reachable from outside the etcd container. Prefer advertising the service DNS name (e.g., http://etcd:2379) so Milvus can reliably connect.
Suggested change
| - -advertise-client-urls=http://127.0.0.1:2379 | |
| - -advertise-client-urls=http://etcd:2379 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
docker-compose.yml (1)
50-96: Good addition for local Milvus development stack.The etcd, minio, and milvus services are correctly configured with:
- Proper environment variables for service discovery (
ETCD_ENDPOINTS,MINIO_ADDRESS).- Persistent named volumes for data retention.
- Healthcheck on the Milvus service with generous retries (20) to accommodate startup time.
Consider adding healthchecks to
etcdandminioservices and usingdepends_onwithcondition: service_healthyfor more robust startup ordering in CI environments. This is optional for local development.♻️ Optional: Add healthchecks for more robust startup
etcd: image: quay.io/coreos/etcd:v3.5.5 environment: ETCD_AUTO_COMPACTION_MODE: revision ETCD_AUTO_COMPACTION_RETENTION: "1000" ETCD_QUOTA_BACKEND_BYTES: "4294967296" ETCD_SNAPSHOT_COUNT: "50000" command: - /usr/local/bin/etcd - -advertise-client-urls=http://127.0.0.1:2379 - -listen-client-urls=http://0.0.0.0:2379 - --data-dir=/etcd ports: - "2379:2379" volumes: - etcddata:/etcd + healthcheck: + test: ["CMD-SHELL", "etcdctl endpoint health || exit 1"] + interval: 5s + timeout: 5s + retries: 10 minio: image: minio/minio:RELEASE.2024-02-17T01-15-57Z environment: MINIO_ROOT_USER: minioadmin MINIO_ROOT_PASSWORD: minioadmin command: minio server /minio_data ports: - "9000:9000" volumes: - miniodata:/minio_data + healthcheck: + test: ["CMD-SHELL", "curl -sf http://localhost:9000/minio/health/live || exit 1"] + interval: 5s + timeout: 5s + retries: 10 milvus: image: milvusdb/milvus:v2.5.3 command: ["milvus", "run", "standalone"] environment: ETCD_ENDPOINTS: etcd:2379 MINIO_ADDRESS: minio:9000 ports: - "19530:19530" - "9091:9091" volumes: - milvusdata:/var/lib/milvus depends_on: - - etcd - - minio + etcd: + condition: service_healthy + minio: + condition: service_healthy healthcheck: test: ["CMD-SHELL", "curl -sf http://localhost:9091/healthz || exit 1"] interval: 5s timeout: 5s retries: 20🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docker-compose.yml` around lines 50 - 96, Add healthchecks for the etcd and minio services and update Milvus's startup dependency to wait for those healthchecks; specifically, for the etcd service (identifier: etcd) add a healthcheck that probes its client port (e.g., curl http://127.0.0.1:2379/health or etcdctl endpoint status) with reasonable interval/timeout/retries, and for the minio service (identifier: minio) add a healthcheck that hits the MinIO health endpoint (e.g., curl -sf http://127.0.0.1:9000/minio/health/live or mc admin info) with interval/timeout/retries; then modify the milvus service's depends_on block (identifier: milvus) to use condition: service_healthy for etcd and minio so Milvus waits until both are healthy before starting.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@vektori/storage/milvus.py`:
- Around line 813-819: The stubbed find_sentences_by_similarity must be
implemented: compute embeddings for the input quotes using the same encoder used
elsewhere (e.g., the class's embedding helper), perform vector searches against
the Milvus collection with a filter for session_id if provided, convert returned
hits whose similarity score meets the threshold into their stored sentence/text
metadata (use the collection's metadata field name, e.g., "sentence" or "text"),
deduplicate and preserve order, and return that list; ensure you use the
existing client/collection access patterns in this class (the
find_sentences_by_similarity function, the class's embedding helper and Milvus
collection/search methods) and handle empty inputs and errors gracefully.
- Around line 1319-1342: The upsert path currently overwrites stored session
metadata with {} when metadata is passed as None; fix by computing a
metadata_value that preserves existing metadata when metadata is None (i.e., use
existing.get("metadata") if existing and metadata is None) otherwise use
self._json_value(metadata), then assign that metadata_value to the
row["metadata"] field instead of always calling self._json_value(metadata);
update the upsert logic around the row construction (reference symbols: row,
existing, metadata, self._json_value) to avoid erasing stored metadata on
metadata-less upserts.
- Around line 1210-1217: The dedup key currently omits agent_id causing
cross-agent collisions; update insert_episode() so episode_id generation and the
dedupe lookup include agent_id: include agent_id in the uuid5 input (e.g.
f"{user_id}::{agent_id}::{text}") when assigning episode_id and add an
additional condition to the _get_one query (e.g. `and agent_id ==
"{_escape(agent_id)}"`) against self._episodes_col using the same
_EPISODE_RECORD and _episode_output_fields() symbols so episodes are scoped per
agent.
- Around line 456-463: The _delete method currently tries calling
self._call("delete", collection_name=collection_name, **{key: filter_expr}) for
keys "filter" and "expr" and then falls back to an unfiltered delete; change
this so that you only perform the delete when one of those parameter forms
succeeds—on first successful call, call self._flush_collection(collection_name)
and return—and if both parameter attempts raise TypeError (i.e. the client
doesn't support either), raise a clear exception (e.g. RuntimeError or
TypeError) instead of calling an unfiltered delete via _call; ensure you
reference the _delete method, the _call invocation, the
filter_expr/collection_name params, and the _flush_collection call in the fix.
- Around line 288-296: The current code silently falls back to unfiltered
operations when both "filter" and "expr" are rejected; update the logic in the
Milvus helpers (the block that calls self._call("query", **{**kwargs, key:
filter_expr})) to fail closed: attempt with "filter" then "expr", and if both
attempts raise TypeError, do not re-run the operation without filters—instead
raise a clear exception (e.g., RuntimeError or ValueError) indicating the client
does not support filtering so the caller can handle it; apply the same pattern
to the analogous _search and _delete call sites and reference the same symbols
(filter_expr, kwargs, self._call("query") / self._call("search") /
self._call("delete")).
---
Nitpick comments:
In `@docker-compose.yml`:
- Around line 50-96: Add healthchecks for the etcd and minio services and update
Milvus's startup dependency to wait for those healthchecks; specifically, for
the etcd service (identifier: etcd) add a healthcheck that probes its client
port (e.g., curl http://127.0.0.1:2379/health or etcdctl endpoint status) with
reasonable interval/timeout/retries, and for the minio service (identifier:
minio) add a healthcheck that hits the MinIO health endpoint (e.g., curl -sf
http://127.0.0.1:9000/minio/health/live or mc admin info) with
interval/timeout/retries; then modify the milvus service's depends_on block
(identifier: milvus) to use condition: service_healthy for etcd and minio so
Milvus waits until both are healthy before starting.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: d7d8042d-a0e1-491d-9c0b-b149bca1b438
📒 Files selected for processing (12)
README.mddocker-compose.ymlexamples/ollama_milvus_session_demo.pypyproject.tomltests/integration/test_milvus_backend.pytests/integration/test_milvus_e2e.pytests/unit/test_milvus_backend.pytests/unit/test_new_backends_factory.pyvektori/cli.pyvektori/config.pyvektori/storage/factory.pyvektori/storage/milvus.py
Contributor
|
@snigenigmatic is the api integration also done? can developers provide api keys and use milvus thru cloud version as well? |
Contributor
Author
|
@laxmanclo need to test API based providers, low on credits. |
Contributor
|
@snigenigmatic no worries, i'll test the API on Milvus on my end. |
Contributor
Author
|
@laxmanclo thanks, in the meantime, i will look into the improvements mentioned above |
…lth-gated startup checks
There was a problem hiding this comment.
🧹 Nitpick comments (4)
vektori/storage/milvus.py (4)
1359-1364: Use_EPISODE_RECORDconstant for consistency.Line 1361 uses the hardcoded string
"episode"instead of the_EPISODE_RECORDconstant used elsewhere in the codebase.♻️ Use constant for record type
rows = await self._query_all( collection_name=self._episodes_col, - filter_expr='record_type == "episode" and is_active == true', + filter_expr=f'record_type == "{_EPISODE_RECORD}" and is_active == true', output_fields=self._episode_output_fields(), batch_size=200, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/milvus.py` around lines 1359 - 1364, Replace the hardcoded filter string 'record_type == "episode"' with the module constant _EPISODE_RECORD in the call to _query_all so the filter_expr uses the constant (e.g., build the expression using _EPISODE_RECORD) for consistency with other code; update the call in the block that invokes self._query_all with collection_name=self._episodes_col and output_fields=self._episode_output_fields() to construct filter_expr using _EPISODE_RECORD instead of the literal "episode".
251-264: String-based loop mismatch detection is fragile but pragmatic.The check for
"attached to a different loop"(line 254) handles pytest-asyncio edge cases, but could break if gRPC changes error message format. Consider catchingRuntimeErrormore broadly or documenting this as a known limitation.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/milvus.py` around lines 251 - 264, The current string-based detection for a gRPC loop-mismatch is fragile; update the exception handling so it catches RuntimeError (instead of a broad Exception) and treats RuntimeError as the signal to fall back to the sync client (keep the logger warning, import MilvusClient, set self._client and call via asyncio.to_thread), while re-raising any non-RuntimeError exceptions; locate the block that references self._client, the string check "attached to a different loop", and MilvusClient to implement this change.
38-39: Consider preserving timezone info in ISO timestamps.
_utcnow_iso()strips timezone info viareplace(tzinfo=None). If downstream code or clients expect timezone-aware timestamps, this could cause subtle issues. Consider using.isoformat()directly or appendingZsuffix for explicit UTC.♻️ Optional: Preserve UTC indicator
def _utcnow_iso() -> str: - return datetime.now(timezone.utc).replace(tzinfo=None).isoformat() + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/milvus.py` around lines 38 - 39, The helper _utcnow_iso currently strips timezone info (via replace(tzinfo=None)), which can break consumers expecting UTC-aware timestamps; update _utcnow_iso to return an ISO 8601 UTC-aware string by either calling datetime.now(timezone.utc).isoformat() (preserving the +00:00) or by generating a naive UTC time and appending a "Z" suffix so the returned value explicitly indicates UTC; adjust any callers if they relied on naive strings.
473-483: Fallback upsert path swallows delete errors which may mask issues.When the primary
upsertfails and the code falls back to delete-then-insert, exceptions from_delete(including the newRuntimeErrorfor unsupported filters) are silently swallowed (lines 478-481). This could lead to duplicate key errors on the subsequent insert if the old records weren't actually deleted.Consider logging the delete failure or allowing it to propagate in the fallback path.
♻️ Log delete failures in fallback path
if ids: ors = " or ".join([f'id == "{_escape(i)}"' for i in ids]) try: await self._delete(collection_name, ors) - except Exception: - pass + except Exception as e: + logger.warning("Fallback upsert: delete failed, may cause duplicates: %s", e) await self._call("insert", collection_name=collection_name, data=rows)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/milvus.py` around lines 473 - 483, The fallback delete-then-insert path currently swallows all exceptions from await self._delete(...), which can mask failures and cause duplicate-key insert errors; change the bare except to capture the exception (except Exception as e) and either log the failure with context (collection_name, ors or ids) using the instance logger (e.g. self._logger or configured logger) and re-raise to stop the insert, or handle RuntimeError from _delete (unsupported filter) separately by logging a warning and continuing; ensure the changes are applied around the _delete call in the upsert fallback so that failures are not silently ignored before calling _call("insert", ...) and _flush_collection.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@vektori/storage/milvus.py`:
- Around line 1359-1364: Replace the hardcoded filter string 'record_type ==
"episode"' with the module constant _EPISODE_RECORD in the call to _query_all so
the filter_expr uses the constant (e.g., build the expression using
_EPISODE_RECORD) for consistency with other code; update the call in the block
that invokes self._query_all with collection_name=self._episodes_col and
output_fields=self._episode_output_fields() to construct filter_expr using
_EPISODE_RECORD instead of the literal "episode".
- Around line 251-264: The current string-based detection for a gRPC
loop-mismatch is fragile; update the exception handling so it catches
RuntimeError (instead of a broad Exception) and treats RuntimeError as the
signal to fall back to the sync client (keep the logger warning, import
MilvusClient, set self._client and call via asyncio.to_thread), while re-raising
any non-RuntimeError exceptions; locate the block that references self._client,
the string check "attached to a different loop", and MilvusClient to implement
this change.
- Around line 38-39: The helper _utcnow_iso currently strips timezone info (via
replace(tzinfo=None)), which can break consumers expecting UTC-aware timestamps;
update _utcnow_iso to return an ISO 8601 UTC-aware string by either calling
datetime.now(timezone.utc).isoformat() (preserving the +00:00) or by generating
a naive UTC time and appending a "Z" suffix so the returned value explicitly
indicates UTC; adjust any callers if they relied on naive strings.
- Around line 473-483: The fallback delete-then-insert path currently swallows
all exceptions from await self._delete(...), which can mask failures and cause
duplicate-key insert errors; change the bare except to capture the exception
(except Exception as e) and either log the failure with context
(collection_name, ors or ids) using the instance logger (e.g. self._logger or
configured logger) and re-raise to stop the insert, or handle RuntimeError from
_delete (unsupported filter) separately by logging a warning and continuing;
ensure the changes are applied around the _delete call in the upsert fallback so
that failures are not silently ignored before calling _call("insert", ...) and
_flush_collection.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 66c37818-863f-40e7-8ce9-63dd50a9de13
📒 Files selected for processing (4)
docker-compose.ymltests/unit/test_milvus_backend.pyvektori/client.pyvektori/storage/milvus.py
…nstant - make _utcnow_iso return UTC-aware ISO timestamps - catch RuntimeError for async loop mismatch and fall back to sync MilvusClient - stop swallowing _delete failures in upsert fallback (warn on unsupported delete, re-raise others) - replace hardcoded episode filter with _EPISODE_RECORD
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@vektori/storage/milvus.py`:
- Around line 1420-1435: The session marker and sentence lookups must be
tenant-scoped: change _session_marker_id to include user_id (e.g., hash
f"session::{user_id}::{session_id}") and update all callers (e.g.,
upsert_session and get_session) so they pass user_id when generating the marker;
also modify the sentence fetch/filter queries against _sentences_col (the
_get_one and any sentence list queries around upsert_session and the block at
1491-1497) to add a user_id equality clause (e.g., and user_id == "<user_id>")
so marker rows and session sentences are isolated per user/tenant. Ensure every
place that builds or queries by the marker or session rows uses the new
user-scoped marker and filter.
- Around line 1248-1258: The current loop over by_fact reads a row with
_get_one, appends CSV ids and calls _upsert_rows, which causes last-writer-wins
races; change this to a compare-and-swap retry: read the row via _get_one,
compute the merged ids using _csv_to_ids/_ids_to_csv, then call _upsert_rows or
an update API with a conditional filter that the row's previous
source_sentence_ids equals the value you read (i.e. include the original
source_sentence_ids in the update predicate); on failure (no rows updated)
re-read and retry until success (or limit retries). Apply the same pattern to
the analogous block around lines 1355-1365 so updates on _facts_col are atomic
and do not drop concurrent links.
- Around line 1208-1217: The query in expand_session_context() incorrectly
restricts results to the same turn by using the equality filter turn_number ==
{turn_number}, which prevents fetching neighboring sentences across turn
boundaries; update the filter passed to _query_all (for collection
self._sentences_col, record type _SENTENCE_RECORD) to remove the turn_number
equality constraint and instead only filter by session_id (using
_escape(session_id)), the sentence_index range (sentence_index >=
{sentence_index - window} and sentence_index <= {sentence_index + window}), and
is_active == true so surrounding sentences from adjacent turns are included.
- Around line 466-495: The fallback delete+insert is currently triggered by a
broad except Exception around self._call("upsert", ...) which masks real write
errors; change this to only use the fallback when the error unmistakably
indicates that upsert is unsupported (e.g., catch a specific NotImplementedError
/ client-specific UnsupportedOperationError or inspect the exception message for
an "unsupported" / "not implemented" indicator), re-raise any other exceptions
so transient RPC/validation failures are not treated as capability gaps, and
keep the existing logic that calls _delete(...) and then _call("insert", ...)
only inside that narrowed except path; reference the _call("upsert",
collection_name=collection_name, data=rows) invocation and the fallback block
that calls _delete and _call("insert") when making the change.
- Around line 788-793: The code currently zips sentences and embeddings which
silently truncates when their lengths differ; before the loop that does for
sent, emb in zip(sentences, embeddings) (and before the early-return for empty
sentences), validate that both inputs are non-empty and have equal lengths
(len(sentences) == len(embeddings)); if not, raise a ValueError (or log and
return an error) with a clear message including the lengths so the caller fails
fast and the mismatch is obvious.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
| "mentions": 1, | ||
| "is_active": True, | ||
| "superseded_by": superseded_by_target or "", | ||
| "event_time": event_time.isoformat() if event_time else "", |
There was a problem hiding this comment.
Normalize event_time before doing string-based date filters.
These predicates compare raw ISO strings. That is only chronologically correct if every stored/query value uses the same offset, so mixed offsets like -05:00 and +00:00 will return wrong before_date / after_date results.
Also applies to: 1051-1054
- Scope session markers by user_id to prevent cross-tenant collisions
- Add user_id filter to session sentence queries in get_session
- Use compare-and-swap retry loop in insert_fact_sources and
insert_episode_fact to prevent last-writer-wins races
- Remove turn_number equality constraint from expand_session_context
so the window spans adjacent turns
- Narrow _upsert_rows fallback to only trigger on unsupported-upsert
errors instead of catching all exceptions
- Validate sentences/embeddings length match before zip in
upsert_sentences
- Fix _call to catch pymilvus-wrapped RuntimeError on event loop
mismatch
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (2)
vektori/storage/milvus.py (2)
1035-1035:⚠️ Potential issue | 🟠 MajorNormalize
event_timebefore filtering on it.
event_timeis persisted as rawisoformat()output and later compared lexicographically. Mixed offsets — plus the empty-string default for missing timestamps — makebefore_date/after_dateresults incorrect. Persist a canonical UTC form (or epoch) and exclude empty values when date filters are active.Also applies to: 1068-1071
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/milvus.py` at line 1035, Persist event_time in a canonical UTC form and avoid empty-string timestamps when date filters are used: replace the current "event_time": event_time.isoformat() if event_time else "" logic with a normalized UTC representation (e.g., event_time.astimezone(datetime.timezone.utc).isoformat() or event_time.timestamp()) so stored values are comparable lexicographically/chronologically, and when date filters (before_date/after_date) are active omit the field or store a null/absent value instead of "" so filtering logic in the code that uses "event_time" (the variable and the "event_time" dict key found at the shown locations and the similar block around lines 1068-1071) behaves correctly.
1277-1289:⚠️ Potential issue | 🟠 MajorThis is still read-check-write, not compare-and-swap.
The extra
_get_one(... == previous_csv)narrows the race window, but two writers can still read the same old value, both pass the check, and the last_upsert_rowswins. If these links must survive concurrent updates, this needs serialization or a true conditional update instead of rewriting CSV blobs.Also applies to: 1406-1417
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/milvus.py` around lines 1277 - 1289, The current logic around _get_one + _upsert_rows is still a read-check-write race; replace it with an atomic conditional update so the write only succeeds if source_sentence_ids equals the previous CSV on the server. Instead of calling _get_one(...) then _upsert_rows(...), use Milvus/Milvus-client's server-side conditional update (update with an expression/filter like `id == "<fid>" and source_sentence_ids == "<previous_csv>"`) or your storage-layer helper that performs an atomic update, failing the update when the expression doesn't match; update the code paths that use this pattern (the block referencing _get_one, _facts_col, source_sentence_ids, previous_csv and the subsequent _upsert_rows call, and the duplicate block around lines 1406-1417) to use that conditional update and handle the false/failed update case by retrying or aborting as appropriate.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@vektori/storage/milvus.py`:
- Around line 1224-1232: The current _query_all call on self._sentences_col uses
only session_id and a sentence_index range which is per-turn and can pull
unrelated sentences; fix by making the window turn-aware: first fetch the seed
sentence (querying for record_type == _SENTENCE_RECORD, session_id and the exact
sentence_index) to get its turn_number (or global ordering key), then either (A)
constrain the follow-up _query_all to include turn_number == seed_turn and the
sentence_index range within that turn, or (B) fetch the ordered sentence stream
for the session (order by turn_number, sentence_index) and slice client-side
around the seed position to build the window; update the code paths that call
_query_all with these new bounds (references: _query_all, self._sentences_col,
_SENTENCE_RECORD, session_id, sentence_index, window).
- Around line 822-829: The current dedupe query in the call to self._get_one
(using self._sentences_col and record type _SENTENCE_RECORD) only matches on
content_hash and can cross tenants/users; change the WHERE clause to also scope
by the tenant/user identifier stored with sentences (e.g., add a clause like and
user_id == "{_escape(user_id)}" or the appropriate tenant field) so the lookup
becomes content_hash + user/tenant-scoped. Update the query construction around
the existing f-string used for content_hash and ensure you use the same escaping
helper (_escape) and the correct variable name for the current user's id.
- Around line 249-265: The RuntimeError except block is too broad and causes any
RuntimeError to trigger the sync fallback; update the exception handling around
the await fn(**kwargs) call so you only fall back when the exception text
indicates the loop mismatch. Concretely, change the bare "except RuntimeError"
to capture the exception (e.g., "except RuntimeError as exc") and re-raise
unless "attached to a different loop" (or the same loop-mismatch substring)
appears in str(exc); keep the existing Milvus-wrapped Exception check but ensure
both branches only fall through to creating MilvusClient and calling fn via
asyncio.to_thread when the message confirms a loop mismatch, otherwise re-raise
the original exception to avoid replaying real failures for functions like the
mutating methods on self._client.
---
Duplicate comments:
In `@vektori/storage/milvus.py`:
- Line 1035: Persist event_time in a canonical UTC form and avoid empty-string
timestamps when date filters are used: replace the current "event_time":
event_time.isoformat() if event_time else "" logic with a normalized UTC
representation (e.g., event_time.astimezone(datetime.timezone.utc).isoformat()
or event_time.timestamp()) so stored values are comparable
lexicographically/chronologically, and when date filters
(before_date/after_date) are active omit the field or store a null/absent value
instead of "" so filtering logic in the code that uses "event_time" (the
variable and the "event_time" dict key found at the shown locations and the
similar block around lines 1068-1071) behaves correctly.
- Around line 1277-1289: The current logic around _get_one + _upsert_rows is
still a read-check-write race; replace it with an atomic conditional update so
the write only succeeds if source_sentence_ids equals the previous CSV on the
server. Instead of calling _get_one(...) then _upsert_rows(...), use
Milvus/Milvus-client's server-side conditional update (update with an
expression/filter like `id == "<fid>" and source_sentence_ids ==
"<previous_csv>"`) or your storage-layer helper that performs an atomic update,
failing the update when the expression doesn't match; update the code paths that
use this pattern (the block referencing _get_one, _facts_col,
source_sentence_ids, previous_csv and the subsequent _upsert_rows call, and the
duplicate block around lines 1406-1417) to use that conditional update and
handle the false/failed update case by retrying or aborting as appropriate.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
Comment on lines
+249
to
+265
| try: | ||
| return await fn(**kwargs) | ||
| except RuntimeError: | ||
| pass # loop mismatch — fall through to sync fallback | ||
| except Exception as exc: | ||
| # pymilvus decorators wrap RuntimeError in MilvusException | ||
| # before our handler sees it; detect the wrapped loop error. | ||
| if "attached to a different loop" not in str(exc): | ||
| raise | ||
| # pytest-asyncio (or other multi-loop runtimes) may reuse a backend | ||
| # instance across event loops, which breaks grpc.aio internals. | ||
| logger.warning("AsyncMilvusClient loop mismatch detected; falling back to sync client") | ||
| from pymilvus import MilvusClient | ||
|
|
||
| self._client = MilvusClient(uri=self.url) | ||
| fn = getattr(self._client, method) | ||
| return await asyncio.to_thread(fn, **kwargs) |
There was a problem hiding this comment.
Only fall back to the sync client on an actual loop-mismatch error.
The bare except RuntimeError will treat every runtime failure as a loop issue and replay the Milvus call on a new sync client. For mutating operations, that can resend a write after a real partial failure.
Proposed fix
if inspect.iscoroutinefunction(fn):
try:
return await fn(**kwargs)
- except RuntimeError:
- pass # loop mismatch — fall through to sync fallback
+ except RuntimeError as exc:
+ if "attached to a different loop" not in str(exc):
+ raise
except Exception as exc:
# pymilvus decorators wrap RuntimeError in MilvusException
# before our handler sees it; detect the wrapped loop error.
if "attached to a different loop" not in str(exc):
raise🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/milvus.py` around lines 249 - 265, The RuntimeError except
block is too broad and causes any RuntimeError to trigger the sync fallback;
update the exception handling around the await fn(**kwargs) call so you only
fall back when the exception text indicates the loop mismatch. Concretely,
change the bare "except RuntimeError" to capture the exception (e.g., "except
RuntimeError as exc") and re-raise unless "attached to a different loop" (or the
same loop-mismatch substring) appears in str(exc); keep the existing
Milvus-wrapped Exception check but ensure both branches only fall through to
creating MilvusClient and calling fn via asyncio.to_thread when the message
confirms a loop mismatch, otherwise re-raise the original exception to avoid
replaying real failures for functions like the mutating methods on self._client.
Comment on lines
+822
to
+829
| existing = await self._get_one( | ||
| self._sentences_col, | ||
| ( | ||
| f'record_type == "{_SENTENCE_RECORD}" and ' | ||
| f'content_hash == "{_escape(content_hash)}"' | ||
| ), | ||
| self._sentence_output_fields(), | ||
| ) |
There was a problem hiding this comment.
Scope sentence deduplication by tenant.
This lookup only keys on content_hash, but that hash is derived from session/turn/text, not user_id. If two users reuse the same session id and sentence positions, they can collapse onto the same row and overwrite each other's sentence state.
Proposed fix
existing = await self._get_one(
self._sentences_col,
(
f'record_type == "{_SENTENCE_RECORD}" and '
- f'content_hash == "{_escape(content_hash)}"'
+ f'content_hash == "{_escape(content_hash)}" and '
+ f'user_id == "{_escape(user_id)}"'
),
self._sentence_output_fields(),
)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/milvus.py` around lines 822 - 829, The current dedupe query
in the call to self._get_one (using self._sentences_col and record type
_SENTENCE_RECORD) only matches on content_hash and can cross tenants/users;
change the WHERE clause to also scope by the tenant/user identifier stored with
sentences (e.g., add a clause like and user_id == "{_escape(user_id)}" or the
appropriate tenant field) so the lookup becomes content_hash +
user/tenant-scoped. Update the query construction around the existing f-string
used for content_hash and ensure you use the same escaping helper (_escape) and
the correct variable name for the current user's id.
Comment on lines
+1224
to
+1232
| rows = await self._query_all( | ||
| collection_name=self._sentences_col, | ||
| filter_expr=( | ||
| f'record_type == "{_SENTENCE_RECORD}" and ' | ||
| f'session_id == "{_escape(session_id)}" and ' | ||
| f"sentence_index >= {sentence_index - window} and " | ||
| f"sentence_index <= {sentence_index + window} and " | ||
| "is_active == true" | ||
| ), |
There was a problem hiding this comment.
Don't window a whole session on sentence_index alone.
After dropping the turn_number equality, this range now matches every turn whose local sentence_index falls in [seed - window, seed + window]. Because sentence_index is per-turn, a seed near a boundary can pull unrelated sentences from the entire session. Build the window from the session's ordered sentence stream, or add turn-aware bounds.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/milvus.py` around lines 1224 - 1232, The current _query_all
call on self._sentences_col uses only session_id and a sentence_index range
which is per-turn and can pull unrelated sentences; fix by making the window
turn-aware: first fetch the seed sentence (querying for record_type ==
_SENTENCE_RECORD, session_id and the exact sentence_index) to get its
turn_number (or global ordering key), then either (A) constrain the follow-up
_query_all to include turn_number == seed_turn and the sentence_index range
within that turn, or (B) fetch the ordered sentence stream for the session
(order by turn_number, sentence_index) and slice client-side around the seed
position to build the window; update the code paths that call _query_all with
these new bounds (references: _query_all, self._sentences_col, _SENTENCE_RECORD,
session_id, sentence_index, window).
This was referenced Apr 15, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR introduces Milvus support as a production-ready storage backend in Vektori, with full backend implementation, factory/config/CLI wiring, documentation updates, Docker local runtime support, and dedicated test coverage.
The Ollama example run is included only as a sanity validation path.
Problem / Motivation
Vektori needs a high-scale vector database option beyond existing backends. Milvus provides scalable ANN search and operational patterns suitable for larger deployments.
Scope
Out of Scope
###Key Design Decisions
Validation
Risks & Mitigations
Backward Compatibility
Summary by CodeRabbit
New Features
Documentation
Chores
Tests