Pr 49 milvus#52
Merged
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)
…lth-gated startup checks
…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
- 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
- Added `milvus_token` to `VektoriConfig` and `Vektori` client initialization - Passed token to both `MilvusClient` and `AsyncMilvusClient` - Added CLI support for `MILVUS_TOKEN` / `MILVUS_API_KEY` - Updated README with Cloud connection examples - Added unit tests for factory token parsing
📝 WalkthroughWalkthroughThis PR adds comprehensive Milvus vector database backend support to Vektori, including a full backend implementation with collection schema management, query/filtering logic, configuration wiring, Docker infrastructure, and extensive unit/integration/e2e test coverage alongside documentation and usage examples. Changes
Sequence DiagramssequenceDiagram
actor User
participant Vektori as Vektori Client
participant Embedder
participant MilvusBackend
participant Milvus as Milvus Server
User->>Vektori: add(messages, session_id, user_id)
activate Vektori
Vektori->>Embedder: embed(text) for each message
Embedder-->>Vektori: embedding vectors
Vektori->>MilvusBackend: upsert_sentences(sentences, embeddings, user_id)
activate MilvusBackend
MilvusBackend->>Milvus: upsert into sentences collection
Milvus-->>MilvusBackend: ack
MilvusBackend->>Milvus: flush (make visible)
Milvus-->>MilvusBackend: flushed
MilvusBackend-->>Vektori: count
deactivate MilvusBackend
Vektori-->>User: result (ok, sentence_ids)
deactivate Vektori
sequenceDiagram
actor User
participant Vektori as Vektori Client
participant Embedder
participant MilvusBackend
participant Milvus as Milvus Server
User->>Vektori: search(query, depth=l1, user_id)
activate Vektori
Vektori->>Embedder: embed(query)
Embedder-->>Vektori: query_embedding
Vektori->>MilvusBackend: search_sentences(embedding, user_id, limit=10)
activate MilvusBackend
MilvusBackend->>Milvus: vector search (cosine) in sentences collection
Milvus-->>MilvusBackend: results with distance scores
MilvusBackend->>MilvusBackend: map distance to result dict
MilvusBackend-->>Vektori: sentence_results
deactivate MilvusBackend
alt depth == l1
Vektori->>MilvusBackend: search_facts(embedding, user_id)
activate MilvusBackend
MilvusBackend->>Milvus: vector search + filter (active, temporal constraints)
Milvus-->>MilvusBackend: fact results
MilvusBackend-->>Vektori: fact_results
deactivate MilvusBackend
end
Vektori-->>User: {sentences, facts, episodes}
deactivate Vektori
Estimated code review effort🎯 4 (Complex) | ⏱️ ~55 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 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 |
Contributor
Author
|
@snigenigmatic Let me know what you think! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (1)
README.md (1)
262-264: Consider documenting both Milvus env var aliases in the cloud example.CLI accepts both
MILVUS_TOKENandMILVUS_API_KEY; showing both here would reduce setup confusion.📝 Suggested doc tweak
-# Milvus Cloud -MILVUS_TOKEN=your-api-key VEKTORI_STORAGE_BACKEND=milvus VEKTORI_DATABASE_URL=https://your-cluster-endpoint vektori add "I prefer dark mode" --user-id u1 +# Milvus Cloud (either env var works) +MILVUS_TOKEN=your-api-key VEKTORI_STORAGE_BACKEND=milvus VEKTORI_DATABASE_URL=https://your-cluster-endpoint vektori add "I prefer dark mode" --user-id u1 +# or +MILVUS_API_KEY=your-api-key VEKTORI_STORAGE_BACKEND=milvus VEKTORI_DATABASE_URL=https://your-cluster-endpoint vektori add "I prefer dark mode" --user-id u1🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@README.md` around lines 262 - 264, Update the Milvus Cloud example to mention both accepted environment variable names by showing both MILVUS_TOKEN and MILVUS_API_KEY (e.g., list both alternatives or include both in the example) and clarify that the CLI accepts either; update the example command that uses vektori add "I prefer dark mode" --user-id u1 to include both env var aliases so readers know either MILVUS_TOKEN or MILVUS_API_KEY will work.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@README.md`:
- Around line 53-54: Update the README pip install line so the “all backends”
extras actually include Postgres: change the extras list referenced by the
package string 'vektori[neo4j,qdrant,milvus]' to include Postgres (e.g. add
'postgres') so the install command reflects all backends; update the string
where 'vektori[neo4j,qdrant,milvus]' appears.
In `@vektori/storage/milvus.py`:
- Around line 922-926: The session-local sentence filters currently only
constrain session_id and is_active which allows cross-tenant reads when
session_ids collide; update the filter building in the Milvus sentence read
paths (the filter_expr construction used by the code around the shown diff and
the similar blocks at the other locations you flagged) to include tenant scoping
(e.g., user_id and/or agent_id) by passing the tenant context into the helpers
or resolving it from the seed row (use expand_session_context() to obtain tenant
fields when needed) and then add conditions like user_id == "{...}" and/or
agent_id == "{...}" to the f-string filters so all reads (including
find_sentence_containing() and L2 expansion callers) are restricted to the same
tenant.
- Around line 1357-1396: The override of insert_episode drops the metadata
parameter and always writes an empty dict; restore the metadata parameter on
insert_episode to match the abstract StorageBackend.insert_episode signature,
accept something like metadata: dict | None, and persist the caller-supplied
metadata by setting row["metadata"] = self._json_value(metadata or {}) before
inserting; keep existing behavior for missing metadata, and leave episode_id
generation, the _EPISODE_RECORD value, and lookups against
_episodes_col/_episode_output_fields unchanged.
- Around line 1143-1164: The find_fact_by_text implementation returns a full
fact dict but the storage contract (base.py) requires returning a fact id or
None; update find_fact_by_text to query only the id field from self._facts_col
(use self._get_one with a projection that requests only the id column instead of
self._fact_output_fields()), and then return the id string (e.g. row["id"] or
the actual id key used in your rows) or None if no row; keep the same filter
construction (record_type/_escape/agent_id) and do not call self._to_fact().
- Around line 1523-1529: upsert_session() is overwriting an existing marker's
agent_id with an empty string when called with agent_id=None; change the upsert
logic so it preserves the stored agent_id instead of replacing it: when building
the row for insertion/update (the dict containing "id", "text", "embedding",
"user_id", "agent_id", "session_id"), if agent_id is None retrieve the existing
marker by marker_id (or read the current agent_id if available) and use that
value, otherwise use the provided agent_id; alternatively only set "agent_id" in
the row when agent_id is not None to avoid erasing the stored agent_id.
- Around line 824-845: The dedup logic uses generate_content_hash (based on
session_id/turn/text) and _get_one against self._sentences_col to find existing
sentence rows by content_hash only, causing cross-tenant collisions and
incorrect increment of mentions and overwrites of tenant fields; update the
lookup in the block that calls _get_one (and any subsequent upsert path using
sid/mentions) to include tenant scope (sent["user_id"] and/or sent["agent_id"])
in the query predicate so existing is only returned for the same tenant, and
ensure the upsert/mentions increment also uses those tenant fields when choosing
sid and writing back.
---
Nitpick comments:
In `@README.md`:
- Around line 262-264: Update the Milvus Cloud example to mention both accepted
environment variable names by showing both MILVUS_TOKEN and MILVUS_API_KEY
(e.g., list both alternatives or include both in the example) and clarify that
the CLI accepts either; update the example command that uses vektori add "I
prefer dark mode" --user-id u1 to include both env var aliases so readers know
either MILVUS_TOKEN or MILVUS_API_KEY will work.
🪄 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: ce9c26cb-4818-46d2-bdeb-4b628ef0230a
📒 Files selected for processing (13)
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/client.pyvektori/config.pyvektori/storage/factory.pyvektori/storage/milvus.py
Comment on lines
+53
to
54
| pip install 'vektori[neo4j,qdrant,milvus]' # all backends | ||
| ``` |
There was a problem hiding this comment.
Fix “all backends” install command to include Postgres dependencies.
Line 53 currently labels vektori[neo4j,qdrant,milvus] as “all backends”, but it skips Postgres extras.
📌 Suggested doc fix
-pip install 'vektori[neo4j,qdrant,milvus]' # all backends
+pip install 'vektori[all]' # all backends (Postgres + Neo4j + Qdrant + Milvus)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| pip install 'vektori[neo4j,qdrant,milvus]' # all backends | |
| ``` | |
| pip install 'vektori[all]' # all backends (Postgres + Neo4j + Qdrant + Milvus) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@README.md` around lines 53 - 54, Update the README pip install line so the
“all backends” extras actually include Postgres: change the extras list
referenced by the package string 'vektori[neo4j,qdrant,milvus]' to include
Postgres (e.g. add 'postgres') so the install command reflects all backends;
update the string where 'vektori[neo4j,qdrant,milvus]' appears.
Comment on lines
+824
to
+845
| content_hash = generate_content_hash( | ||
| sent["session_id"], | ||
| f"{sent['turn_number']}_{sent['sentence_index']}", | ||
| sent["text"], | ||
| ) | ||
|
|
||
| existing = await self._get_one( | ||
| self._sentences_col, | ||
| ( | ||
| f'record_type == "{_SENTENCE_RECORD}" and ' | ||
| f'content_hash == "{_escape(content_hash)}"' | ||
| ), | ||
| self._sentence_output_fields(), | ||
| ) | ||
|
|
||
| if existing: | ||
| sid = existing.get("id") | ||
| mentions = int(existing.get("mentions") or 1) + 1 | ||
| else: | ||
| sid = sent["id"] | ||
| mentions = 1 | ||
|
|
There was a problem hiding this comment.
Scope sentence deduplication by tenant.
content_hash here is derived from session/turn/text only, and the lookup reuses any matching row without checking user_id or agent_id. Two tenants reusing the same session id and utterance will merge into one record, increment mentions, and overwrite the original tenant fields on upsert.
Suggested fix
existing = await self._get_one(
self._sentences_col,
(
f'record_type == "{_SENTENCE_RECORD}" and '
+ f'user_id == "{_escape(user_id)}" and '
+ f'agent_id == "{_escape(agent_id or "")}" and '
f'content_hash == "{_escape(content_hash)}"'
),
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 824 - 845, The dedup logic uses
generate_content_hash (based on session_id/turn/text) and _get_one against
self._sentences_col to find existing sentence rows by content_hash only, causing
cross-tenant collisions and incorrect increment of mentions and overwrites of
tenant fields; update the lookup in the block that calls _get_one (and any
subsequent upsert path using sid/mentions) to include tenant scope
(sent["user_id"] and/or sent["agent_id"]) in the query predicate so existing is
only returned for the same tenant, and ensure the upsert/mentions increment also
uses those tenant fields when choosing sid and writing back.
Comment on lines
+922
to
+926
| filter_expr = ( | ||
| f'record_type == "{_SENTENCE_RECORD}" and ' | ||
| f'session_id == "{_escape(session_id)}" and ' | ||
| "is_active == true" | ||
| ) |
There was a problem hiding this comment.
Add tenant scoping to session-local sentence reads.
These filters only constrain session_id (and find_sentence_containing() also skips is_active). If two users reuse the same session id, quote matching and L2 expansion can pull another tenant's sentences into the current result set. expand_session_context() can at least add the seed row's user_id/agent_id; the other helpers need the same tenant context passed in or resolved before querying.
Also applies to: 967-971, 997-999, 1234-1240
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/milvus.py` around lines 922 - 926, The session-local sentence
filters currently only constrain session_id and is_active which allows
cross-tenant reads when session_ids collide; update the filter building in the
Milvus sentence read paths (the filter_expr construction used by the code around
the shown diff and the similar blocks at the other locations you flagged) to
include tenant scoping (e.g., user_id and/or agent_id) by passing the tenant
context into the helpers or resolving it from the seed row (use
expand_session_context() to obtain tenant fields when needed) and then add
conditions like user_id == "{...}" and/or agent_id == "{...}" to the f-string
filters so all reads (including find_sentence_containing() and L2 expansion
callers) are restricted to the same tenant.
Comment on lines
+1143
to
+1164
| async def find_fact_by_text( | ||
| self, | ||
| user_id: str, | ||
| text: str, | ||
| agent_id: str | None = None, | ||
| ) -> dict[str, Any] | None: | ||
| parts = [ | ||
| f'record_type == "{_FACT_RECORD}"', | ||
| f'user_id == "{_escape(user_id)}"', | ||
| "is_active == true", | ||
| f'text == "{_escape(text)}"', | ||
| ] | ||
| if agent_id is not None: | ||
| parts.append(f'agent_id == "{_escape(agent_id)}"') | ||
|
|
||
| row = await self._get_one( | ||
| self._facts_col, | ||
| " and ".join(parts), | ||
| self._fact_output_fields(), | ||
| ) | ||
| return self._to_fact(row) if row else None | ||
|
|
There was a problem hiding this comment.
Match find_fact_by_text() to the storage contract.
vektori/storage/base.py declares this method as returning a fact id or None, but this implementation returns a full fact dict. Any caller that forwards the result into id-based APIs will now get the wrong type back.
Suggested fix
async def find_fact_by_text(
self,
user_id: str,
text: str,
agent_id: str | None = None,
- ) -> dict[str, Any] | None:
+ ) -> str | None:
@@
- return self._to_fact(row) if row else None
+ return str(row["id"]) if row else None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async def find_fact_by_text( | |
| self, | |
| user_id: str, | |
| text: str, | |
| agent_id: str | None = None, | |
| ) -> dict[str, Any] | None: | |
| parts = [ | |
| f'record_type == "{_FACT_RECORD}"', | |
| f'user_id == "{_escape(user_id)}"', | |
| "is_active == true", | |
| f'text == "{_escape(text)}"', | |
| ] | |
| if agent_id is not None: | |
| parts.append(f'agent_id == "{_escape(agent_id)}"') | |
| row = await self._get_one( | |
| self._facts_col, | |
| " and ".join(parts), | |
| self._fact_output_fields(), | |
| ) | |
| return self._to_fact(row) if row else None | |
| async def find_fact_by_text( | |
| self, | |
| user_id: str, | |
| text: str, | |
| agent_id: str | None = None, | |
| ) -> str | None: | |
| parts = [ | |
| f'record_type == "{_FACT_RECORD}"', | |
| f'user_id == "{_escape(user_id)}"', | |
| "is_active == true", | |
| f'text == "{_escape(text)}"', | |
| ] | |
| if agent_id is not None: | |
| parts.append(f'agent_id == "{_escape(agent_id)}"') | |
| row = await self._get_one( | |
| self._facts_col, | |
| " and ".join(parts), | |
| self._fact_output_fields(), | |
| ) | |
| return str(row["id"]) if row else None |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/milvus.py` around lines 1143 - 1164, The find_fact_by_text
implementation returns a full fact dict but the storage contract (base.py)
requires returning a fact id or None; update find_fact_by_text to query only the
id field from self._facts_col (use self._get_one with a projection that requests
only the id column instead of self._fact_output_fields()), and then return the
id string (e.g. row["id"] or the actual id key used in your rows) or None if no
row; keep the same filter construction (record_type/_escape/agent_id) and do not
call self._to_fact().
Comment on lines
+1357
to
+1396
| async def insert_episode( | ||
| self, | ||
| text: str, | ||
| embedding: list[float], | ||
| user_id: str, | ||
| agent_id: str | None = None, | ||
| session_id: str | None = None, | ||
| ) -> str: | ||
| if self.embedding_dim and len(embedding) != self.embedding_dim: | ||
| raise ValueError( | ||
| f"Embedding has {len(embedding)} dimensions, expected {self.embedding_dim}" | ||
| ) | ||
|
|
||
| agent_scope = agent_id or "" | ||
| episode_id = str(uuid.uuid5(uuid.NAMESPACE_OID, f"{user_id}::{agent_scope}::{text}")) | ||
| existing = await self._get_one( | ||
| self._episodes_col, | ||
| ( | ||
| f'id == "{_escape(episode_id)}" and ' | ||
| f'record_type == "{_EPISODE_RECORD}" and ' | ||
| f'agent_id == "{_escape(agent_scope)}"' | ||
| ), | ||
| self._episode_output_fields(), | ||
| ) | ||
| if existing: | ||
| return episode_id | ||
|
|
||
| row = { | ||
| "id": episode_id, | ||
| "text": text, | ||
| "embedding": embedding, | ||
| "user_id": user_id, | ||
| "agent_id": agent_id or "", | ||
| "session_id": session_id or "", | ||
| "is_active": True, | ||
| "created_at": _utcnow_iso(), | ||
| "fact_ids": "", | ||
| "record_type": _EPISODE_RECORD, | ||
| "metadata": self._json_value({}), | ||
| } |
There was a problem hiding this comment.
Restore the metadata parameter on insert_episode().
The abstract StorageBackend.insert_episode() signature includes metadata, but this override drops it and always persists {}. That breaks the backend contract and loses caller-supplied episode metadata.
Suggested fix
async def insert_episode(
self,
text: str,
embedding: list[float],
user_id: str,
agent_id: str | None = None,
session_id: str | None = None,
+ metadata: dict[str, Any] | None = None,
) -> str:
@@
- "metadata": self._json_value({}),
+ "metadata": self._json_value(metadata),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async def insert_episode( | |
| self, | |
| text: str, | |
| embedding: list[float], | |
| user_id: str, | |
| agent_id: str | None = None, | |
| session_id: str | None = None, | |
| ) -> str: | |
| if self.embedding_dim and len(embedding) != self.embedding_dim: | |
| raise ValueError( | |
| f"Embedding has {len(embedding)} dimensions, expected {self.embedding_dim}" | |
| ) | |
| agent_scope = agent_id or "" | |
| episode_id = str(uuid.uuid5(uuid.NAMESPACE_OID, f"{user_id}::{agent_scope}::{text}")) | |
| existing = await self._get_one( | |
| self._episodes_col, | |
| ( | |
| f'id == "{_escape(episode_id)}" and ' | |
| f'record_type == "{_EPISODE_RECORD}" and ' | |
| f'agent_id == "{_escape(agent_scope)}"' | |
| ), | |
| self._episode_output_fields(), | |
| ) | |
| if existing: | |
| return episode_id | |
| row = { | |
| "id": episode_id, | |
| "text": text, | |
| "embedding": embedding, | |
| "user_id": user_id, | |
| "agent_id": agent_id or "", | |
| "session_id": session_id or "", | |
| "is_active": True, | |
| "created_at": _utcnow_iso(), | |
| "fact_ids": "", | |
| "record_type": _EPISODE_RECORD, | |
| "metadata": self._json_value({}), | |
| } | |
| async def insert_episode( | |
| self, | |
| text: str, | |
| embedding: list[float], | |
| user_id: str, | |
| agent_id: str | None = None, | |
| session_id: str | None = None, | |
| metadata: dict[str, Any] | None = None, | |
| ) -> str: | |
| if self.embedding_dim and len(embedding) != self.embedding_dim: | |
| raise ValueError( | |
| f"Embedding has {len(embedding)} dimensions, expected {self.embedding_dim}" | |
| ) | |
| agent_scope = agent_id or "" | |
| episode_id = str(uuid.uuid5(uuid.NAMESPACE_OID, f"{user_id}::{agent_scope}::{text}")) | |
| existing = await self._get_one( | |
| self._episodes_col, | |
| ( | |
| f'id == "{_escape(episode_id)}" and ' | |
| f'record_type == "{_EPISODE_RECORD}" and ' | |
| f'agent_id == "{_escape(agent_scope)}"' | |
| ), | |
| self._episode_output_fields(), | |
| ) | |
| if existing: | |
| return episode_id | |
| row = { | |
| "id": episode_id, | |
| "text": text, | |
| "embedding": embedding, | |
| "user_id": user_id, | |
| "agent_id": agent_id or "", | |
| "session_id": session_id or "", | |
| "is_active": True, | |
| "created_at": _utcnow_iso(), | |
| "fact_ids": "", | |
| "record_type": _EPISODE_RECORD, | |
| "metadata": self._json_value(metadata), | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/milvus.py` around lines 1357 - 1396, The override of
insert_episode drops the metadata parameter and always writes an empty dict;
restore the metadata parameter on insert_episode to match the abstract
StorageBackend.insert_episode signature, accept something like metadata: dict |
None, and persist the caller-supplied metadata by setting row["metadata"] =
self._json_value(metadata or {}) before inserting; keep existing behavior for
missing metadata, and leave episode_id generation, the _EPISODE_RECORD value,
and lookups against _episodes_col/_episode_output_fields unchanged.
Comment on lines
+1523
to
+1529
| row = { | ||
| "id": marker_id, | ||
| "text": "", | ||
| "embedding": [0.0] * self.embedding_dim, | ||
| "user_id": user_id, | ||
| "agent_id": agent_id or "", | ||
| "session_id": session_id, |
There was a problem hiding this comment.
Don't erase the stored agent_id on session updates.
When upsert_session() is called for an existing session with agent_id=None, this rewrites the marker with agent_id="". Subsequent agent-scoped session queries stop matching that session.
Suggested fix
row = {
"id": marker_id,
"text": "",
"embedding": [0.0] * self.embedding_dim,
"user_id": user_id,
- "agent_id": agent_id or "",
+ "agent_id": (
+ agent_id
+ if agent_id is not None
+ else (existing.get("agent_id", "") if existing else "")
+ ),
"session_id": session_id,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/milvus.py` around lines 1523 - 1529, upsert_session() is
overwriting an existing marker's agent_id with an empty string when called with
agent_id=None; change the upsert logic so it preserves the stored agent_id
instead of replacing it: when building the row for insertion/update (the dict
containing "id", "text", "embedding", "user_id", "agent_id", "session_id"), if
agent_id is None retrieve the existing marker by marker_id (or read the current
agent_id if available) and use that value, otherwise use the provided agent_id;
alternatively only set "agent_id" in the row when agent_id is not None to avoid
erasing the stored agent_id.
This was referenced Apr 13, 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
2 participants
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.
I noticed that the cloud usage via API keys wasn't supported yet, so I went ahead and added support for Milvus / Zilliz Cloud authentication directly to this branch.
What I added:
Added milvus_token to VektoriConfig which gets passed directly down to both the async and sync MilvusClients.
Enabled CLI support via MILVUS_TOKEN or MILVUS_API_KEY environment variables.
Included unit tests to verify the factory passes the token correctly.
Updated the README.md with examples of how to connect to Zilliz Cloud using the new token configuration.
Summary by CodeRabbit
New Features
Documentation
Configuration
vektori[milvus]