Add milvus backend unit test - #40
Conversation
|
Warning Rate limit exceeded
⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📝 WalkthroughWalkthroughRefactors backend APIs and Milvus integration: adds backend details/typed attributes, reduces embedding dim 768→384, normalizes schema defaults/quoting, defaults metadata to empty dict in multiple places, updates AGENTS.md, and adds comprehensive Milvus unit tests. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant Backend as "MilvusEntityBackend"
participant Embed as "EmbeddingModel\n(SentenceTransformer)"
participant Milvus as "MilvusClient"
participant DB as "SQLiteManager"
Client->>Backend: create/search/update/delete namespace/entity
Backend->>Milvus: list/create/drop collection / get_collection_stats
Backend->>Embed: encode content -> embedding
Backend->>Milvus: insert/upsert/search/delete with embedding & filters
Milvus-->>Backend: stats / search results / ack
Backend->>DB: persist/delete namespace metadata
DB-->>Backend: confirm persistence
Backend-->>Client: return result or error
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
kaizen/backend/milvus.py (1)
80-84: Handle emptyentitiesinput explicitly.entities[0]will raiseIndexErroron empty input; returning a clear error is safer.Proposed fix
- entity_type = entities[0].type + if not entities: + raise KaizenException("At least one entity is required.") + entity_type = entities[0].type
🤖 Fix all issues with AI agents
In `@kaizen/backend/milvus.py`:
- Around line 60-63: The code assumes SQLiteManager.get_namespace(namespace_id)
returns a namespace object but it can be None; update the block in the method
that uses SQLiteManager and get_namespace to guard against a missing row: after
calling namespace = db_manager.get_namespace(namespace_id) check if namespace is
None and handle it (e.g., raise a descriptive exception or return
None/appropriate response), and only then set namespace.num_entities =
self.milvus.get_collection_stats(namespace_id)["row_count"] and return
namespace; reference the SQLiteManager, get_namespace, namespace.num_entities,
and self.milvus.get_collection_stats symbols when applying the change.
- Around line 91-95: The code normalizes metadata to {} only for temporary
RecordedEntity instances but passes update.metadata and entity.metadata (which
can be None) directly into Milvus insert/upsert; update the code paths that call
the Milvus insert/upsert routines (where entity and update objects are used) to
ensure metadata is normalized to an empty dict whenever None — e.g., before
creating RecordedEntity, before using entity.metadata in insert, and before
using update.metadata in upsert, replace None with {} (or call a small helper
normalize_metadata(metadata) that returns {} for None) so Milvus always receives
valid JSON metadata.
In `@tests/unit/test_milvus_backend.py`:
- Around line 38-59: Several test stubs have unused parameters causing Ruff
ARG001; update the parameter names to start with an underscore for each unused
param to satisfy linting: rename parameters in always_has_collection
(collection_name → _collection_name), never_has_collection (_collection_name),
noop_create_collection (collection_name, dimension, auto_id, schema →
_collection_name, _dimension, _auto_id, _schema), arbitrary_collection_stats
(collection_name → _collection_name), arbitrary_embedding (text → _text), and
the test stub methods search_entities, insert, resolve_conflicts, and
query—prefix any unused parameters in those method signatures with an underscore
(e.g., params → _params) so Ruff no longer flags ARG001.
- Around line 1-21: The test imports MilvusEntityBackend at module import time
which triggers class-level initialization of milvus and embedding_model before
your fixture patches run; fix by deferring the import so the patches apply
first—inside the fixture (or top of the test) patch
"kaizen.backend.milvus.MilvusClient" and
"kaizen.backend.milvus.SentenceTransformer" and only then import or from-import
MilvusEntityBackend (the class name MilvusEntityBackend and its class attributes
milvus and embedding_model should be referenced after the patch), ensuring no
real client/model are constructed during test collection.
🧹 Nitpick comments (1)
kaizen/backend/milvus.py (1)
163-175: Escape/validate filter values before interpolating. Raw string interpolation into the Milvus DSL can break queries when values contain quotes or operators; a small escape/validation helper would make this more robust.
| def always_has_collection(collection_name: str): | ||
| return True | ||
|
|
||
|
|
||
| def never_has_collection(collection_name: str): | ||
| return False | ||
|
|
||
|
|
||
| def noop_create_collection(collection_name: str, dimension, auto_id, schema): | ||
| pass | ||
|
|
||
|
|
||
| def arbitrary_namespace(namespace_id: str) -> Namespace: | ||
| return Namespace(id=namespace_id, created_at=datetime.datetime.now(datetime.UTC)) | ||
|
|
||
|
|
||
| def arbitrary_collection_stats(collection_name: str): | ||
| return {"row_count": 42} | ||
|
|
||
|
|
||
| def arbitrary_embedding(text: str): | ||
| return [0.1] * 384 |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, find the file
find . -name "test_milvus_backend.py" -type f 2>/dev/null | head -20Repository: AgentToolkit/kaizen
Length of output: 98
🏁 Script executed:
#!/bin/bash
# Read the file to check the code at the specified lines
cat -n ./tests/unit/test_milvus_backend.py | head -80Repository: AgentToolkit/kaizen
Length of output: 3074
🏁 Script executed:
#!/bin/bash
# Check lines 159-167 and 198-206
sed -n '159,167p' ./tests/unit/test_milvus_backend.py
echo "---"
sed -n '198,206p' ./tests/unit/test_milvus_backend.pyRepository: AgentToolkit/kaizen
Length of output: 702
🏁 Script executed:
#!/bin/bash
# Check pyproject.toml for Ruff configuration
cat pyproject.toml | grep -A 50 "\[tool.ruff"Repository: AgentToolkit/kaizen
Length of output: 109
🏁 Script executed:
#!/bin/bash
# Get more context around line 159 to see the class/function structure
sed -n '150,175p' ./tests/unit/test_milvus_backend.pyRepository: AgentToolkit/kaizen
Length of output: 1341
🏁 Script executed:
#!/bin/bash
# Get more context around line 198 to see that section
sed -n '190,210p' ./tests/unit/test_milvus_backend.pyRepository: AgentToolkit/kaizen
Length of output: 814
🏁 Script executed:
#!/bin/bash
# Check full pyproject.toml to see all Ruff configuration
cat pyproject.tomlRepository: AgentToolkit/kaizen
Length of output: 999
🏁 Script executed:
#!/bin/bash
# Verify if arbitrary_namespace uses its parameter
sed -n '50,51p' ./tests/unit/test_milvus_backend.pyRepository: AgentToolkit/kaizen
Length of output: 205
Ruff ARG001 flags these unused stub parameters. Prefix them with _ to keep linting green. As per coding guidelines, use Ruff for linting and formatting (configured in pyproject.toml).
The following functions have unused parameters that need fixing:
- Lines 38-39, 42-43, 54-55, 58-59:
always_has_collection,never_has_collection,arbitrary_collection_stats,arbitrary_embedding - Lines 46-47:
noop_create_collection - Lines 159-167:
search_entities,insert,resolve_conflicts(stub methods in test) - Line 198:
querymethod
Apply the underscore prefix convention (e.g., collection_name → _collection_name) to all unused parameters.
🧰 Tools
🪛 Ruff (0.14.14)
38-38: Unused function argument: collection_name
(ARG001)
42-42: Unused function argument: collection_name
(ARG001)
54-54: Unused function argument: collection_name
(ARG001)
58-58: Unused function argument: text
(ARG001)
🤖 Prompt for AI Agents
In `@tests/unit/test_milvus_backend.py` around lines 38 - 59, Several test stubs
have unused parameters causing Ruff ARG001; update the parameter names to start
with an underscore for each unused param to satisfy linting: rename parameters
in always_has_collection (collection_name → _collection_name),
never_has_collection (_collection_name), noop_create_collection
(collection_name, dimension, auto_id, schema → _collection_name, _dimension,
_auto_id, _schema), arbitrary_collection_stats (collection_name →
_collection_name), arbitrary_embedding (text → _text), and the test stub methods
search_entities, insert, resolve_conflicts, and query—prefix any unused
parameters in those method signatures with an underscore (e.g., params →
_params) so Ruff no longer flags ARG001.
|
This seems to be a critical issue #40 (comment) |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@kaizen/backend/milvus.py`:
- Around line 162-180: In search_entities, the vector-search branch currently
passes an empty filter string when filters is empty; update the call to
self.milvus.query (the branch that sets anns_field, data, limit, search_params)
to use the same fallback expression used earlier: compute the filter as " AND
".join([f"{k} == '{v}'" for k, v in filters.items()]) if len(filters) > 0 else
"id > 0" and pass that into the filter parameter so both non-vector and vector
queries use the consistent fallback.
|
What about this one? #40 (comment) |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
kaizen/backend/milvus.py (1)
85-89: Handle empty entity batches to avoid IndexError.
entities[0]will raise whenentitiesis empty. Add an early guard.🛠️ Proposed fix
def update_entities(self, namespace_id: str, entities: list[Entity], enable_conflict_resolution: bool = True) -> list[EntityUpdate]: self.validate_namespace(namespace_id) + if not entities: + return [] entity_type = entities[0].type
🤖 Fix all issues with AI agents
In `@kaizen/backend/milvus.py`:
- Around line 52-58: The create_namespace method currently calls
create_collection(..., auto_id=False) which conflicts with the entity_schema's
id field (auto_id=True) and causes inserts that omit id (in the insert logic) to
fail; update the create_collection call in create_namespace to use auto_id=True
so the collection matches entity_schema and allows the database to auto-generate
ids (refer to create_namespace, entity_schema, and create_collection in your
change).
- Around line 168-180: The Milvus filter construction is vulnerable to syntax
errors when values contain backslashes or single quotes; add a small helper
(e.g., _escape_milvus_value) that takes a string v and returns v with
backslashes and single quotes escaped (replace "\" -> "\\" and "'" -> "\'") and
use it when building the filter expressions passed to self.milvus.query in both
branches (the "query is None" block and the else block where
anns_field/data/limit are used), replacing f"{k} == '{v}'" with f"{k} ==
'{_escape_milvus_value(v)}'"; ensure filters is handled the same way in both
places.
| def create_namespace(self, namespace_id: str | None = None) -> Namespace: | ||
| """Create a new namespace for entities to exist in.""" | ||
| namespace_id = namespace_id or 'ns_' + str(uuid.uuid4()).replace('-', '_') | ||
| namespace_id = namespace_id or "ns_" + str(uuid.uuid4()).replace("-", "_") | ||
|
|
||
| if not self.milvus.has_collection(namespace_id): | ||
| self.milvus.create_collection(collection_name=namespace_id, dimension=768, auto_id=False, schema=entity_schema) | ||
| self.milvus.create_collection(collection_name=namespace_id, dimension=384, auto_id=False, schema=entity_schema) | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n kaizen/backend/milvus.pyRepository: AgentToolkit/kaizen
Length of output: 11676
Change auto_id=False to auto_id=True to match schema.
The entity_schema defines the id field with auto_id=True (line 213), but create_collection() at line 57 explicitly sets auto_id=False. This conflicts with the insert operations at lines 114–123 and 148–157, which omit the id field entirely, expecting the database to auto-generate it. With auto_id=False, these inserts will fail. Align the parameter with the schema definition.
Suggested fix
- self.milvus.create_collection(collection_name=namespace_id, dimension=384, auto_id=False, schema=entity_schema)
+ self.milvus.create_collection(collection_name=namespace_id, dimension=384, auto_id=True, schema=entity_schema)📝 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.
| def create_namespace(self, namespace_id: str | None = None) -> Namespace: | |
| """Create a new namespace for entities to exist in.""" | |
| namespace_id = namespace_id or 'ns_' + str(uuid.uuid4()).replace('-', '_') | |
| namespace_id = namespace_id or "ns_" + str(uuid.uuid4()).replace("-", "_") | |
| if not self.milvus.has_collection(namespace_id): | |
| self.milvus.create_collection(collection_name=namespace_id, dimension=768, auto_id=False, schema=entity_schema) | |
| self.milvus.create_collection(collection_name=namespace_id, dimension=384, auto_id=False, schema=entity_schema) | |
| def create_namespace(self, namespace_id: str | None = None) -> Namespace: | |
| """Create a new namespace for entities to exist in.""" | |
| namespace_id = namespace_id or "ns_" + str(uuid.uuid4()).replace("-", "_") | |
| if not self.milvus.has_collection(namespace_id): | |
| self.milvus.create_collection(collection_name=namespace_id, dimension=384, auto_id=True, schema=entity_schema) |
🤖 Prompt for AI Agents
In `@kaizen/backend/milvus.py` around lines 52 - 58, The create_namespace method
currently calls create_collection(..., auto_id=False) which conflicts with the
entity_schema's id field (auto_id=True) and causes inserts that omit id (in the
insert logic) to fail; update the create_collection call in create_namespace to
use auto_id=True so the collection matches entity_schema and allows the database
to auto-generate ids (refer to create_namespace, entity_schema, and
create_collection in your change).
There was a problem hiding this comment.
wrong, the schema defines it
There was a problem hiding this comment.
@illeatmyhat, you're absolutely right—I apologize for the confusion. When schema=entity_schema is passed to create_collection(), the schema's field definitions (including auto_id=True in the FieldSchema for the id field) take precedence. The auto_id=False parameter in the create_collection() call is effectively ignored in this case. The code is working as intended, and there's no conflict. Thank you for the correction!
1772ce8 to
01e36b8
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@kaizen/backend/base.py`:
- Around line 44-47: The method search_entities in BaseEntityBackend is missing
the `@abstractmethod` decorator; add `@abstractmethod` immediately above the
search_entities definition so subclasses are required to implement it (and if
abstractmethod isn't already imported from abc, add that import). Ensure the
signature of search_entities remains unchanged and matches other abstract
methods in the class.
In `@kaizen/backend/filesystem.py`:
- Line 68: The raise site in filesystem.py uses a long formatted message in the
raise of NamespaceAlreadyExistsException which triggers Ruff TRY003; either move
the message construction into the exception class (e.g., add an __init__ or
classmethod on NamespaceAlreadyExistsException that builds the message from
namespace_id and then raise NamespaceAlreadyExistsException(namespace_id)) or
suppress the linter at the raise site with a local noqa (TRY003) comment; update
the NamespaceAlreadyExistsException implementation to accept namespace_id and
produce the formatted message so callers (like the raise in filesystem.py)
simply raise the exception with the identifier.
In `@kaizen/schema/core.py`:
- Around line 22-24: The metadata field was changed to disallow None which
breaks callers that pass parsed_metadata=None (e.g., kaizen/frontend/cli/cli.py
and checks in kaizen/sync/phoenix_sync.py); fix by restoring compatibility:
either revert metadata to Optional[dict] with default=None or add a Pydantic
field validator on the model that owns the metadata Field (the metadata Field in
kaizen/schema/core.py) to coerce None to {} (use a pre=True, always=True
validator for "metadata" to return {} when input is None). Ensure the change
preserves the Field description and default behavior for non-None values.
In `@tests/unit/test_milvus_backend.py`:
- Around line 62-72: The test expectation is outdated:
MilvusEntityBackend.ready() now returns a bool, not a dict. Update the
test_ready assertion to expect a boolean (e.g., assert result is True or assert
result == True) while keeping the existing monkeypatch of
milvus.list_collections and the test name; ensure you call
milvus_backend.ready() and assert the returned boolean value from that method.
🧹 Nitpick comments (1)
kaizen/schema/core.py (1)
14-16: Consider addingstrict=Truetozip()to catch mismatched row/field counts, but note this is not currently enforced by Ruff in your project.The
zip(fields, row)call will silently truncate if row length drifts from cursor description length. While Ruff rule B905 (zip-without-explicit-strict) would flag this, B905 is not enabled in your current Ruff configuration (pyproject.tomllacks a[tool.ruff.lint]section with flake8-bugbear selected). The suggestion remains good defensive coding practice for robustness.🛠️ Suggested fix
- return Namespace(**{k: v for k, v in zip(fields, row)}) + return Namespace(**{k: v for k, v in zip(fields, row, strict=True)})
| def search_entities( | ||
| self, | ||
| namespace_id: str, | ||
| query: str | None = None, | ||
| filters: dict | None = None, | ||
| limit: int = 10 | ||
| self, namespace_id: str, query: str | None = None, filters: dict | None = None, limit: int = 10 | ||
| ) -> list[RecordedEntity]: | ||
| pass |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, examine the base.py file to understand the class structure
cat -n kaizen/backend/base.py | head -60Repository: AgentToolkit/kaizen
Length of output: 1741
🏁 Script executed:
# Search for the class definition and imports to understand if it's already abstract
rg -n "^from abc|^import abc|class BaseEntityBackend|@abstractmethod" kaizen/backend/base.py -B2 -A2Repository: AgentToolkit/kaizen
Length of output: 1046
🏁 Script executed:
# Find all subclasses of BaseEntityBackend
rg -n "class.*\(.*BaseEntityBackend" kaizen/backendRepository: AgentToolkit/kaizen
Length of output: 218
🏁 Script executed:
# Check if search_entities is implemented in subclasses
rg -n "def search_entities" kaizen/backend -B3 -A5Repository: AgentToolkit/kaizen
Length of output: 1704
🏁 Script executed:
# Check pyproject.toml for Ruff configuration
cat pyproject.toml | grep -A20 "\[tool.ruff"Repository: AgentToolkit/kaizen
Length of output: 109
Add @abstractmethod decorator to search_entities to maintain consistency with other abstract methods in the class.
All other methods in BaseEntityBackend are marked with @abstractmethod, but search_entities is missing this decorator. This inconsistency weakens the interface contract and allows subclasses to skip implementation without raising an error during instantiation.
🛠️ Suggested fix
+ `@abstractmethod`
def search_entities(
self, namespace_id: str, query: str | None = None, filters: dict | None = None, limit: int = 10
) -> list[RecordedEntity]:
pass📝 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.
| def search_entities( | |
| self, | |
| namespace_id: str, | |
| query: str | None = None, | |
| filters: dict | None = None, | |
| limit: int = 10 | |
| self, namespace_id: str, query: str | None = None, filters: dict | None = None, limit: int = 10 | |
| ) -> list[RecordedEntity]: | |
| pass | |
| `@abstractmethod` | |
| def search_entities( | |
| self, namespace_id: str, query: str | None = None, filters: dict | None = None, limit: int = 10 | |
| ) -> list[RecordedEntity]: | |
| pass |
🧰 Tools
🪛 Ruff (0.14.14)
44-47: BaseEntityBackend.search_entities is an empty method in an abstract base class, but has no abstract decorator
(B027)
🤖 Prompt for AI Agents
In `@kaizen/backend/base.py` around lines 44 - 47, The method search_entities in
BaseEntityBackend is missing the `@abstractmethod` decorator; add `@abstractmethod`
immediately above the search_entities definition so subclasses are required to
implement it (and if abstractmethod isn't already imported from abc, add that
import). Ensure the signature of search_entities remains unchanged and matches
other abstract methods in the class.
| raise NamespaceAlreadyExistsException( | ||
| f'Namespace "{namespace_id}" already exists.' | ||
| ) | ||
| raise NamespaceAlreadyExistsException(f'Namespace "{namespace_id}" already exists.') |
There was a problem hiding this comment.
Ruff TRY003: avoid long message at raise site.
If TRY003 is enabled, move the message into the exception class or suppress locally.
🛠️ Minimal suppression
- raise NamespaceAlreadyExistsException(f'Namespace "{namespace_id}" already exists.')
+ raise NamespaceAlreadyExistsException(f'Namespace "{namespace_id}" already exists.') # noqa: TRY003🧰 Tools
🪛 Ruff (0.14.14)
68-68: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In `@kaizen/backend/filesystem.py` at line 68, The raise site in filesystem.py
uses a long formatted message in the raise of NamespaceAlreadyExistsException
which triggers Ruff TRY003; either move the message construction into the
exception class (e.g., add an __init__ or classmethod on
NamespaceAlreadyExistsException that builds the message from namespace_id and
then raise NamespaceAlreadyExistsException(namespace_id)) or suppress the linter
at the raise site with a local noqa (TRY003) comment; update the
NamespaceAlreadyExistsException implementation to accept namespace_id and
produce the formatted message so callers (like the raise in filesystem.py)
simply raise the exception with the identifier.
| content: str | list | dict = Field(description="Searchable text or structured data.") | ||
| metadata: dict = Field(default_factory=dict, description="Arbitrary metadata which is related to the entity.") | ||
| type: str = Field(description="The type of the entity.") |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's look at the current state of the file
cat -n kaizen/schema/core.py | head -50Repository: AgentToolkit/kaizen
Length of output: 1521
🏁 Script executed:
# Check git history for this file to see what changed
git log --oneline -10 kaizen/schema/core.pyRepository: AgentToolkit/kaizen
Length of output: 81
🏁 Script executed:
# Get the previous version of the metadata field definition
git show HEAD~1:kaizen/schema/core.py 2>/dev/null | cat -n | head -50Repository: AgentToolkit/kaizen
Length of output: 45
🏁 Script executed:
# Check full git log with more detail
git log --all --oneline kaizen/schema/core.pyRepository: AgentToolkit/kaizen
Length of output: 185
🏁 Script executed:
# Check if this is a new file
git log --follow --oneline kaizen/schema/core.py | head -20Repository: AgentToolkit/kaizen
Length of output: 81
🏁 Script executed:
# Search for usages of Entity and metadata in the codebase
rg "Entity\(" --type py -B 2 -A 2 | head -100Repository: AgentToolkit/kaizen
Length of output: 6226
🏁 Script executed:
# Look for test files that might use Entity with metadata
fd "test" --type f -e py | head -20Repository: AgentToolkit/kaizen
Length of output: 257
🏁 Script executed:
# Check the git diff for the metadata field in recent commits
git show 6d486cd:kaizen/schema/core.py 2>/dev/null | cat -n | head -50Repository: AgentToolkit/kaizen
Length of output: 1481
🏁 Script executed:
# Also check the initial commit
git show 46acae5:kaizen/schema/core.py 2>/dev/null | cat -n | head -50Repository: AgentToolkit/kaizen
Length of output: 1865
🏁 Script executed:
# Let's see what changed in the "fix filter" commit
git show 01e36b8 --statRepository: AgentToolkit/kaizen
Length of output: 5176
🏁 Script executed:
# Check if there's a diff between the current and previous version
git diff 6d486cd HEAD -- kaizen/schema/core.pyRepository: AgentToolkit/kaizen
Length of output: 2192
🏁 Script executed:
# Search for any code that explicitly passes metadata=None
rg "metadata\s*=\s*None" --type pyRepository: AgentToolkit/kaizen
Length of output: 249
🏁 Script executed:
# Search for any code that sets metadata to None
rg "\.metadata\s*=\s*None" --type pyRepository: AgentToolkit/kaizen
Length of output: 141
🏁 Script executed:
# Check if there are any loads/parsing scenarios where None might be passed
rg "dict\(" kaizen/schema/core.py kaizen/backend/ --type py -B 2 -A 2Repository: AgentToolkit/kaizen
Length of output: 45
🏁 Script executed:
# Test if the current code actually rejects None with Pydantic v2
python3 <<'EOF'
from pydantic import BaseModel, Field
# Simulate current definition
class EntityCurrent(BaseModel):
content: str | list | dict = Field(description="Searchable text or structured data.")
metadata: dict = Field(default_factory=dict, description="Arbitrary metadata which is related to the entity.")
type: str = Field(description="The type of the entity.")
# Simulate previous definition
class EntityPrevious(BaseModel):
content: str | list | dict = Field(description="Searchable text or structured data.")
metadata: dict | None = Field(default=None, description="Arbitrary metadata which is related to the entity.")
type: str = Field(description="The type of the entity.")
# Test current
try:
e = EntityCurrent(type="test", content="test", metadata=None)
print(f"Current accepts None: {e.metadata}")
except Exception as ex:
print(f"Current rejects None: {type(ex).__name__}: {ex}")
# Test previous
try:
e = EntityPrevious(type="test", content="test", metadata=None)
print(f"Previous accepts None: {e.metadata}")
except Exception as ex:
print(f"Previous rejects None: {type(ex).__name__}: {ex}")
# Test omitting metadata in current
try:
e = EntityCurrent(type="test", content="test")
print(f"Current with omitted metadata: {e.metadata}")
except Exception as ex:
print(f"Current fails when omitted: {type(ex).__name__}: {ex}")
EOFRepository: AgentToolkit/kaizen
Length of output: 381
🏁 Script executed:
# Check the test that uses metadata=None
cat -n tests/unit/test_cli.py | grep -A 5 -B 5 "metadata=None"Repository: AgentToolkit/kaizen
Length of output: 457
🏁 Script executed:
# Check the test that assigns metadata = None
cat -n tests/unit/test_phoenix_sync.py | grep -A 10 -B 10 "metadata = None"Repository: AgentToolkit/kaizen
Length of output: 1016
🏁 Script executed:
# Run the tests to see if they fail
python3 -m pytest tests/unit/test_cli.py -v --tb=short 2>&1 | head -100Repository: AgentToolkit/kaizen
Length of output: 1285
🏁 Script executed:
# Check what happens in the cli.py file with parsed_metadata
cat -n kaizen/frontend/cli/cli.py | grep -A 10 -B 10 "parsed_metadata = None"Repository: AgentToolkit/kaizen
Length of output: 963
🏁 Script executed:
# Check if cli.py actually passes parsed_metadata to Entity creation
cat -n kaizen/frontend/cli/cli.py | sed -n '174,210p'Repository: AgentToolkit/kaizen
Length of output: 1473
🏁 Script executed:
# Look for where Entity is created with metadata in cli.py
rg "Entity\(" kaizen/frontend/cli/cli.py -A 3 -B 1Repository: AgentToolkit/kaizen
Length of output: 169
🏁 Script executed:
# Check if the test at line 503 actually runs (check test function)
cat -n tests/unit/test_cli.py | sed -n '490,510p'Repository: AgentToolkit/kaizen
Length of output: 958
🏁 Script executed:
# Let's trace through the cli.py flow - when parsed_metadata is None, it gets passed to Entity
# This should fail validation. Let's verify the exact code path
cat -n kaizen/frontend/cli/cli.py | sed -n '201,206p'Repository: AgentToolkit/kaizen
Length of output: 217
🏁 Script executed:
# Check if there are any other places where metadata might be None
rg "metadata\s*=" kaizen/ --type py | grep -v "metadata: " | head -20Repository: AgentToolkit/kaizen
Length of output: 1010
🏁 Script executed:
# Let's check if the backend or other code handles None metadata
rg "\.metadata" kaizen/backend/ kaizen/sync/ --type py -B 1 -A 1 | head -50Repository: AgentToolkit/kaizen
Length of output: 2376
Breaking change confirmed: metadata no longer accepts None.
The change from metadata: dict | None = Field(default=None, ...) to metadata: dict = Field(default_factory=dict, ...) will cause runtime failures. Production code in kaizen/frontend/cli/cli.py:204 passes metadata=parsed_metadata where parsed_metadata is initialized as None (line 184) and remains None when no metadata argument is provided. This will raise a Pydantic ValidationError. Additionally, kaizen/sync/phoenix_sync.py contains checks like if e.metadata and e.metadata.get("span_id"), indicating the codebase expects metadata can be None.
Implement one of:
- Ensure all callers provide a dict (including
{}) instead ofNone - Use the field validator to coerce
Noneto{}:
Compatibility shim
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
class Entity(BaseModel):
metadata: dict = Field(default_factory=dict, description="Arbitrary metadata which is related to the entity.")
+
+ `@field_validator`("metadata", mode="before")
+ `@classmethod`
+ def _coerce_metadata(cls, v):
+ return v or {}🤖 Prompt for AI Agents
In `@kaizen/schema/core.py` around lines 22 - 24, The metadata field was changed
to disallow None which breaks callers that pass parsed_metadata=None (e.g.,
kaizen/frontend/cli/cli.py and checks in kaizen/sync/phoenix_sync.py); fix by
restoring compatibility: either revert metadata to Optional[dict] with
default=None or add a Pydantic field validator on the model that owns the
metadata Field (the metadata Field in kaizen/schema/core.py) to coerce None to
{} (use a pre=True, always=True validator for "metadata" to return {} when input
is None). Ensure the change preserves the Field description and default behavior
for non-None values.
|
@illeatmyhat can you check why these are giving errors? |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@tests/e2e/test_e2e_pipeline.py`:
- Line 61: The two subprocess.run invocations that execute the agent script and
perform trace verification (the call assigning to result using
subprocess.run([...], env=env, capture_output=True, text=True) and the second
similar call later) lack timeouts and can hang CI; add timeout=300 for the agent
execution call and timeout=60 for the trace verification call, and wrap each run
in try/except subprocess.TimeoutExpired to fail the test fast while capturing
and including the process stdout/stderr in the test failure message (use
variables result.stdout/result.stderr or the exception's output/ stderr) so
diagnostics are printed when a timeout occurs.
- Around line 132-140: The loop using process.stdout.readline() can block
indefinitely and defeat the timeout; replace it with non-blocking I/O using
selectors: create a selectors.DefaultSelector(), register process.stdout (use
process.stdout.fileno()) and in your loop call
selector.select(timeout=remaining_time) where remaining_time = timeout -
(time.time() - sync_start); if select returns empty, print the timeout message
and break, otherwise call process.stdout.readline() to consume the available
data and continue; also still check process.poll() after reads and
unregister/close the selector when done. Ensure you reference and update the
existing symbols process, sync_start, timeout, process.stdout.readline(), and
process.poll() in the changed code.
In `@tests/unit/test_milvus_backend.py`:
- Around line 237-248: The test test_delete_entity_by_id misses mocking
milvus.query which delete_entity_by_id calls to confirm existence; add a
monkeypatch for milvus_backend.milvus.query (e.g., return a truthy list like
[{"id": 12345}] when called with collection_name="test_namespace" and filter
matching id) so the existence check passes and the delete path (milvus.delete)
is exercised; reference the delete_entity_by_id method, milvus.query, and
milvus.delete when adding the mock.
- Around line 1-21: The test fixture is fine but the delete test misses mocking
milvus.query and some stub functions have unused parameters; update the test to
monkeypatch milvus_backend.milvus.query with a Mock (e.g., returning an empty
list or a list with the entity) before calling delete_entity_by_id, and rename
stub functions used as patches to prefix their parameters with an underscore
(e.g., always_has_collection(_collection_name: str),
never_has_collection(_collection_name: str), and similar stubs) to resolve the
Ruff ARG001 unused-argument warning and ensure delete_entity_by_id sees the
expected query result.
🧹 Nitpick comments (5)
kaizen/schema/core.py (1)
15-16: Addstrict=Truetozip()for safer row parsing.Per Ruff B905, the
zip()call should includestrict=Trueto ensurefieldsandrowhave matching lengths. A mismatch would indicate a data integrity issue that should fail fast rather than silently truncate.Proposed fix
def row_factory(cursor: Cursor, row: Row) -> "Namespace": fields = [column[0] for column in cursor.description] - return Namespace(**{k: v for k, v in zip(fields, row)}) + return Namespace(**{k: v for k, v in zip(fields, row, strict=True)})kaizen/frontend/mcp/mcp_server.py (2)
147-153: Alignlogging.exceptionand f-strings with Ruff (TRY401/RUF010).Ruff flags redundant exception info in
logging.exceptionandstr(e)inside f-strings. Consider simplifying the log message and using{e!s}for the response. As per coding guidelines "Use Ruff for linting and formatting (configured in pyproject.toml)".💡 Suggested diff
- except json.JSONDecodeError as e: - logger.exception(f"Invalid JSON in metadata parameter: {str(e)}") - return json.dumps( - {"error": "Invalid metadata JSON", "message": f"Failed to parse metadata: {str(e)}", "invalid_metadata": metadata} - ) + except json.JSONDecodeError as e: + logger.exception("Invalid JSON in metadata parameter") + return json.dumps( + {"error": "Invalid metadata JSON", "message": f"Failed to parse metadata: {e!s}", "invalid_metadata": metadata} + )
193-195: Fix redundant exception formatting indelete_entitylogging.Ruff warns about including exception text in
logging.exceptionandstr(e)in f-strings. Prefer a structured log message and{e!s}. As per coding guidelines "Use Ruff for linting and formatting (configured in pyproject.toml)".💡 Suggested diff
- except KaizenException as e: - logger.exception(f"Error deleting entity {entity_id}: {str(e)}") - return json.dumps({"success": False, "error": str(e)}) + except KaizenException as e: + logger.exception("Error deleting entity %s", entity_id) + return json.dumps({"success": False, "error": f"{e!s}"})tests/unit/test_milvus_backend.py (2)
158-165: Prefix unused parameters in nested stubs with underscore.These nested stub functions have unused parameters flagged by Ruff ARG001. Apply the underscore prefix convention:
♻️ Proposed fix
- def search_entities(self, namespace_id, query, filters=None, limit=10): + def search_entities(_self, _namespace_id, _query, _filters=None, _limit=10): return [] - def insert(collection_name, data): + def insert(_collection_name, _data): return {"ids": [12345]} - def resolve_conflicts(old_entities, new_entities): + def resolve_conflicts(_old_entities, _new_entities): return [entity_update]
197-206: Prefix unused parameters inquerystub with underscore.The
querystub has multiple unused parameters flagged by Ruff ARG001:♻️ Proposed fix
- def query(collection_name, filter="", output_fields=None, timeout=None, ids=None, partition_names=None, **kwargs): + def query(_collection_name, _filter="", _output_fields=None, _timeout=None, _ids=None, _partition_names=None, **_kwargs): return [ { "id": 123, "type": "fact", "content": "Test content", "created_at": int(datetime.datetime.now(datetime.UTC).timestamp()), "metadata": {}, } ]
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@kaizen/backend/milvus.py`:
- Around line 35-36: The helper _escape_filter currently calls .replace on its
input assuming it's a str which will raise for ints/bools; update _escape_filter
to coerce the input to a string (e.g., value = str(value) or similar) before
performing the .replace calls so non-string filter values are safely handled and
returned as an escaped string; keep the function name _escape_filter and ensure
the return remains a str.
- Around line 93-96: The method update_entities currently dereferences
entities[0] which will raise on an empty list; add an early guard at the top of
update_entities to handle empty input (either return [] immediately or raise a
clear ValueError like "entities must not be empty") and ensure callers expect
the chosen contract; update_entities should validate entities is non-empty
before accessing entities[0] and proceed only after this check.
- Around line 206-210: Replace the broad Exception catch in the close() cleanup
with the specific MilvusException: import MilvusException from pymilvus and
change the except block that currently wraps self.milvus.close() to except
MilvusException as e, then call logger.warning(f"Error closing Milvus client:
{e}") so only Milvus-related errors are swallowed and other unexpected
exceptions propagate; reference symbols: self.milvus.close(), MilvusException,
logger.warning.
🧹 Nitpick comments (2)
kaizen/backend/milvus.py (2)
99-109: Skip building temporary IDs when conflict resolution is disabled.
entities_with_temporary_idsis only used in the conflict-resolution path; move construction inside theif enable_conflict_resolutionblock to avoid unnecessary work.
56-58: Ruff TRY003: move long exception message into the exception class.Ruff flags long inline messages; consider defining a default message in
NamespaceNotFoundException(or a helper) and raising the exception without constructing a long f-string here.
As per coding guidelines,**/*.py: Use Ruff for linting and formatting (configured in pyproject.toml).
3 participants
Summary by CodeRabbit
Documentation
Refactor
Schema
Behavior
Tests