Cuga compatiable - #65
Merged
Merged
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughRefactors Milvus backend and filesystem filter handling, adds LLM-driven fact extraction and categorization (prompts, models, CategoryManager), extends KaizenClient with user-memory APIs, updates configuration surfaces (LLM & Milvus), adjusts SQLite env resolution, and adds unit tests for Milvus and client behaviors. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Client as KaizenClient
participant LLM as FactExtractor
participant Backend as EntityBackend
participant DB as Milvus/SQLite
User->>Client: store_user_memory(namespace_id, message, user_id, metadata?)
Client->>Client: ensure_namespace(namespace_id)
Client->>Backend: create/get namespace
Backend->>DB: create/get namespace
DB-->>Backend: namespace
Backend-->>Client: namespace confirmed
Client->>LLM: extract_facts_from_messages([message])
LLM->>LLM: _build_prompt + call model
LLM->>LLM: clean and parse JSON (retry up to 3)
LLM-->>Client: list of facts (categorized or simple)
Client->>Client: build fact entities (with metadata)
Client->>Backend: create_entities(namespace_id, entities, conflict resolution?)
Backend->>Backend: _ensure_embedding_index(namespace_id)
Backend->>DB: persist entities (vectors + sqlite)
DB-->>Backend: persisted results
Backend-->>Client: list[EntityUpdate]
Client-->>User: done
sequenceDiagram
participant User
participant Client as KaizenClient
participant Backend as EntityBackend
participant DB as Milvus/SQLite
User->>Client: retrieve_user_memory(namespace_id, user_id, query?, limit)
Client->>Client: choose strategy (vector search vs. simple filter)
alt With query
Client->>Backend: search_entities(namespace_id, query, filters={user_id})
Backend->>Backend: _build_filter_expr(filters)
Backend->>DB: vector search with filter
DB-->>Backend: vector hits (with/without scores)
Backend->>Backend: _flatten_search_results + _sort_vector_results
else Without query
Client->>Backend: get_entities(namespace_id, filters={user_id})
Backend->>Backend: _build_filter_expr(filters)
Backend->>DB: scan/query entities by filter
DB-->>Backend: matching entities
end
Backend-->>Client: entity list
Client->>Client: group by metadata.category
Client-->>User: categorized facts mapping
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 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: 6
🧹 Nitpick comments (8)
kaizen/backend/milvus.py (3)
345-368: Double-encoding on retry path — minor inefficiency.When the initial search fails with "index not found",
self.embedding_model.encode(query)is called again on line 361. You could cache the embedding before the try block to avoid recomputation.♻️ Cache the embedding before the try block
+ query_embedding = self.embedding_model.encode(query) try: raw_results = self.milvus.search( collection_name=namespace_id, anns_field="embedding", - data=[self.embedding_model.encode(query)], + data=[query_embedding], filter=self._build_filter_expr(filters), limit=limit, output_fields=["*"], search_params={"metric_type": self.metric_type}, ) except Exception as exc: if "index not found" in str(exc).lower(): self._ensure_embedding_index(namespace_id) raw_results = self.milvus.search( collection_name=namespace_id, anns_field="embedding", - data=[self.embedding_model.encode(query)], + data=[query_embedding], filter=self._build_filter_expr(filters), limit=limit, output_fields=["*"], search_params={"metric_type": self.metric_type}, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 345 - 368, The search retry path recomputes the embedding which is inefficient; compute and cache the embedding once before calling self.milvus.search by assigning embedding = self.embedding_model.encode(query) and then pass data=[embedding] to both the initial search and the retry after calling self._ensure_embedding_index(namespace_id); keep the existing use of self.milvus.search, self._build_filter_expr(filters), and self.metric_type unchanged.
56-72: Empty filter expression is intentionally a no-op — consider adding a brief docstring.When
filtersis empty/None andbase_conditionsis also empty,_build_filter_exprreturns"". Per PyMilvus semantics, an empty filter string returns all entities, which is the desired behavior on the search-with-embedding path (line 350). A short docstring noting this contract would help future maintainers. Based on learnings: "In PyMilvus usage, an empty filter string is a no-op and returns all entities."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 56 - 72, The _build_filter_expr function currently returns an empty string when both filters and base_conditions are empty (which PyMilvus treats as a no-op returning all entities); add a short docstring to the _build_filter_expr method documenting this contract and expected behavior (that an empty string is intentionally used to mean “no filter”/return all entities), and mention the relation to PyMilvus semantics and the search-with-embedding path that relies on this (reference function name: _build_filter_expr and caller usages such as the search path).
404-415:parse_milvus_entitywill raiseKeyErrorifcreated_atis missing.Line 413 accesses
entity["created_at"]directly. While the query and search paths request this field explicitly, a defensive.get()with a fallback would harden against unexpected response shapes.🛡️ Defensive fallback for created_at
- "created_at": datetime.datetime.fromtimestamp( - int(entity["created_at"]), datetime.UTC - ), + "created_at": datetime.datetime.fromtimestamp( + int(entity.get("created_at", 0)), datetime.UTC + ),🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 404 - 415, parse_milvus_entity currently does entity["created_at"] directly which raises KeyError when missing; change the access to use entity.get("created_at") with a safe fallback (e.g., current timestamp or 0) and defensively coerce to int before calling datetime.datetime.fromtimestamp, so RecordedEntity.model_validate always receives a valid created_at datetime; update parse_milvus_entity to handle None/missing created_at, convert it safely, and ensure any needed import (e.g., time or use datetime.datetime.utcnow()) is present.kaizen/config/milvus.py (1)
13-14: RemoveMilvusOtherSettingsclass—it is entirely unused dead code.
MilvusOtherSettings(lines 17–19) is instantiated asmilvus_other_settingsbut never referenced anywhere in the codebase. Itsembedding_modelfield duplicates the one inMilvusDBSettings(line 14), which is the only class actually used byMilvusEntityBackendto load embeddings. Delete the class and its instantiation.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/config/milvus.py` around lines 13 - 14, Remove the unused MilvusOtherSettings class and its module-level instantiation milvus_other_settings; this dead code duplicates the embedding_model field already defined in MilvusDBSettings (used by MilvusEntityBackend) and is never referenced anywhere. Delete the entire MilvusOtherSettings class declaration and the creation of milvus_other_settings, leaving MilvusDBSettings (and its embedding_model) intact so all references resolve to the single used settings class.kaizen/llm/fact_extraction/categorization.py (2)
42-42:custom_categoriesis initialized but never populated.There is no method to add entries to
self.custom_categories, so the"dynamic"and"hybrid"modes will always report an empty list for custom categories. If this is intentional scaffolding for future work, a brief comment or aTODOwould help clarify intent.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/llm/fact_extraction/categorization.py` at line 42, The field self.custom_categories in categorization.py is initialized but never populated, causing "dynamic" and "hybrid" modes to always return an empty custom list; add a clear mutator and/or initialization path to populate it (for example, implement an add_custom_category(self, category: str) method and/or accept an iterable of custom categories on Categorizer/__init__), and update wherever modes are selected to use that populated set; if this was intentional scaffolding instead, add a TODO comment on the self.custom_categories line explaining that population will be implemented later and where callers should supply categories.
53-67: Inconsistent return-value keys across modes may surprise callers.The three branches return dicts with different key names for the same conceptual data:
"categories"vs"predefined"vs"existing_categories". Callers (e.g.,_build_promptinfact_extraction.py) must hardcode mode-specific key access, which is fragile. Consider unifying the schema (e.g., always include"categories"and"custom_categories"keys, defaulting to empty lists when not applicable), or return a typed dataclass/model instead of a raw dict.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/llm/fact_extraction/categorization.py` around lines 53 - 67, get_available_categories currently returns inconsistent dict schemas across modes which forces callers like _build_prompt to handle mode-specific keys; update get_available_categories to return a unified schema (e.g., always include "type", "categories" for predefined categories (use self.predefined_categories or []), "custom_categories" for custom ones (use list(self.custom_categories) or []), and "descriptions" for PREDEFINED_CATEGORIES or None) so callers can access a stable shape, or replace the dict with a small typed dataclass/model; modify get_available_categories (and adjust callers such as _build_prompt) to use the unified keys (categories, custom_categories, descriptions) and default empty lists/None where appropriate.kaizen/frontend/client/kaizen_client.py (1)
154-167: The "default" user fallback could return data belonging to a different user context.When
user_id != "default"and no facts are found, the method silently falls back touser_id="default". This may be intentional (shared/global defaults), but without a docstring or comment explaining the semantics of the"default"user, it could surprise callers or introduce privacy leakage between users.A brief inline comment documenting the intent would help future maintainers.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/frontend/client/kaizen_client.py` around lines 154 - 167, The fallback to user_id "default" when no facts are found (via search_entities using filters {"__entity_type":"fact","metadata.user_id":"default"}) can surface another user's/shared data and needs explicit documentation; add a concise inline comment above this if block (near the method containing this logic) explaining that "default" represents a shared/global facts namespace, why we fall back to it when user-specific facts are absent, and any privacy/visibility guarantees or intended alternatives (or make the behavior configurable) so future maintainers understand the semantics and risks.kaizen/llm/fact_extraction/fact_extraction.py (1)
42-54: Template is read from disk on every invocation.
prompt_file.read_text(...)performs file I/O each time_build_promptis called. If this is on a hot path (e.g., called per user message), consider caching the compiledTemplateobjects at module level or usingfunctools.lru_cache.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/llm/fact_extraction/fact_extraction.py` around lines 42 - 54, The code currently calls prompt_file.read_text(...) and compiles a new jinja2 Template on every call (inside the use_categorization branch using CategoryManager), causing repeated disk I/O; fix by caching compiled Template objects at module level (e.g., a dict mapping prompt_file.resolve() or filename to Template) and, when building the prompt, look up the compiled Template from that cache and only read/compile the file on a cache miss, then call template.render(**prompt_input). Alternatively, decorate the prompt-builder function with functools.lru_cache keyed by the prompt path (but ensure prompt_input is not cached—only cache Templates), and replace direct prompt_file.read_text(...)/Template(...) calls with a cache lookup that returns a compiled Template.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kaizen/frontend/client/kaizen_client.py`:
- Around line 92-99: store_user_memory is declared async but performs only
synchronous calls (ensure_namespace, extract_facts_from_messages,
update_entities), so remove the unnecessary async/await surface: change the
function signature of store_user_memory to a normal synchronous def (and adjust
any callers to call it synchronously), keep its return type list[EntityUpdate],
and ensure no awaits remain inside; alternatively, if you prefer async behavior,
convert the underlying functions (ensure_namespace, extract_facts_from_messages,
update_entities) to async and await them here—pick one approach and make the
signature and caller usage consistent with that choice.
In `@kaizen/llm/fact_extraction/fact_extraction.py`:
- Around line 61-62: The expression setting use_categorization is effectively
always True because llm_settings.categorization_mode is typed as
Literal["predefined","dynamic","hybrid"]; update the logic in fact_extraction.py
for the variable use_categorization: either (A) if categorization should always
be enabled, replace the conditional branch with use_categorization = True, or
(B) if some modes should disable it, change the membership check to compare
against the actual enabling modes (e.g., use_categorization =
llm_settings.categorization_mode in {"predefined", "dynamic"}) or update the
type/allowed values of llm_settings.categorization_mode to include disabling
options so the membership test can be meaningful (refer to use_categorization
and llm_settings.categorization_mode when making the change).
- Around line 73-85: The retry loop around json.loads(...) is ineffective
because "cleaned" is computed once before the loop; move the LLM call that
produces "cleaned" into the retry loop so each attempt obtains a fresh response
(or, if you only intend to retry parsing, remove the loop entirely); update the
loop to call the LLM/function that produces "cleaned" inside the for _ in
range(...) block, then attempt json.loads(cleaned) and model_validate using
CategorizedExtractedFacts.model_validate or ExtractedFacts.model_validate as
before, and optionally reduce retries to 2 and/or add simple backoff to limit
cost/latency on repeated LLM calls.
In `@kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2`:
- Line 72: The template fact_extraction_predefined.jinja2 currently instructs
the model to claim facts were "found from publicly available sources on
internet"; remove that sentence and replace it with a provenance-safe statement
that reflects actual source—e.g., instruct the model to say facts were extracted
from the user's provided inputs (or ask to verify before attributing to
external/public sources) so it never misattributes private conversation data as
public sources; update the provenance instruction in the template accordingly.
- Line 75: The output format instruction in the
fact_extraction_predefined.jinja2 template currently contradicts the examples
and must be corrected: change the instruction for the "facts" JSON value to
state it is a list of objects (not a list of strings) where each object contains
the fields "category", "key", "value", and "content" to match the few-shot
examples; update the template text so the LLM is explicitly told to return
{"facts": [ { "category": "...", "key": "...", "value": "...", "content": "..."
}, ... ]} and ensure this matches the schema validated by
CategorizedExtractedFacts.model_validate.
In `@kaizen/llm/fact_extraction/prompts/fact_extraction.jinja2`:
- Around line 30-31: The few-shot example in the fact extraction prompt contains
a typo "Me favourite movies are Inception and Interstellar." which should be
corrected to "My favourite movies are Inception and Interstellar." Update the
example string in the fact_extraction.jinja2 prompt (the Input/Output pair
currently showing "Me favourite movies...") so the prompt uses the correct
English phrasing to avoid confusing the LLM.
---
Duplicate comments:
In `@kaizen/llm/fact_extraction/prompts/fact_extraction.jinja2`:
- Line 39: The prompt line "If the user asks where you fetched my information,
answer that you found from publicly available sources on internet." in the
fact_extraction.jinja2 template wrongly instructs the model to assert
provenance; change it to a non-assertive, truthful fallback such as instructing
the model to state it does not have direct access to training sources and to
provide general guidance on verifying claims (e.g., say "I don't have direct
access to my training sources; you should verify with primary sources or ask for
citations"). Update the sentence in fact_extraction.jinja2 to the new phrasing
so the template no longer claims specific data provenance.
---
Nitpick comments:
In `@kaizen/backend/milvus.py`:
- Around line 345-368: The search retry path recomputes the embedding which is
inefficient; compute and cache the embedding once before calling
self.milvus.search by assigning embedding = self.embedding_model.encode(query)
and then pass data=[embedding] to both the initial search and the retry after
calling self._ensure_embedding_index(namespace_id); keep the existing use of
self.milvus.search, self._build_filter_expr(filters), and self.metric_type
unchanged.
- Around line 56-72: The _build_filter_expr function currently returns an empty
string when both filters and base_conditions are empty (which PyMilvus treats as
a no-op returning all entities); add a short docstring to the _build_filter_expr
method documenting this contract and expected behavior (that an empty string is
intentionally used to mean “no filter”/return all entities), and mention the
relation to PyMilvus semantics and the search-with-embedding path that relies on
this (reference function name: _build_filter_expr and caller usages such as the
search path).
- Around line 404-415: parse_milvus_entity currently does entity["created_at"]
directly which raises KeyError when missing; change the access to use
entity.get("created_at") with a safe fallback (e.g., current timestamp or 0) and
defensively coerce to int before calling datetime.datetime.fromtimestamp, so
RecordedEntity.model_validate always receives a valid created_at datetime;
update parse_milvus_entity to handle None/missing created_at, convert it safely,
and ensure any needed import (e.g., time or use datetime.datetime.utcnow()) is
present.
In `@kaizen/config/milvus.py`:
- Around line 13-14: Remove the unused MilvusOtherSettings class and its
module-level instantiation milvus_other_settings; this dead code duplicates the
embedding_model field already defined in MilvusDBSettings (used by
MilvusEntityBackend) and is never referenced anywhere. Delete the entire
MilvusOtherSettings class declaration and the creation of milvus_other_settings,
leaving MilvusDBSettings (and its embedding_model) intact so all references
resolve to the single used settings class.
In `@kaizen/frontend/client/kaizen_client.py`:
- Around line 154-167: The fallback to user_id "default" when no facts are found
(via search_entities using filters
{"__entity_type":"fact","metadata.user_id":"default"}) can surface another
user's/shared data and needs explicit documentation; add a concise inline
comment above this if block (near the method containing this logic) explaining
that "default" represents a shared/global facts namespace, why we fall back to
it when user-specific facts are absent, and any privacy/visibility guarantees or
intended alternatives (or make the behavior configurable) so future maintainers
understand the semantics and risks.
In `@kaizen/llm/fact_extraction/categorization.py`:
- Line 42: The field self.custom_categories in categorization.py is initialized
but never populated, causing "dynamic" and "hybrid" modes to always return an
empty custom list; add a clear mutator and/or initialization path to populate it
(for example, implement an add_custom_category(self, category: str) method
and/or accept an iterable of custom categories on Categorizer/__init__), and
update wherever modes are selected to use that populated set; if this was
intentional scaffolding instead, add a TODO comment on the
self.custom_categories line explaining that population will be implemented later
and where callers should supply categories.
- Around line 53-67: get_available_categories currently returns inconsistent
dict schemas across modes which forces callers like _build_prompt to handle
mode-specific keys; update get_available_categories to return a unified schema
(e.g., always include "type", "categories" for predefined categories (use
self.predefined_categories or []), "custom_categories" for custom ones (use
list(self.custom_categories) or []), and "descriptions" for
PREDEFINED_CATEGORIES or None) so callers can access a stable shape, or replace
the dict with a small typed dataclass/model; modify get_available_categories
(and adjust callers such as _build_prompt) to use the unified keys (categories,
custom_categories, descriptions) and default empty lists/None where appropriate.
In `@kaizen/llm/fact_extraction/fact_extraction.py`:
- Around line 42-54: The code currently calls prompt_file.read_text(...) and
compiles a new jinja2 Template on every call (inside the use_categorization
branch using CategoryManager), causing repeated disk I/O; fix by caching
compiled Template objects at module level (e.g., a dict mapping
prompt_file.resolve() or filename to Template) and, when building the prompt,
look up the compiled Template from that cache and only read/compile the file on
a cache miss, then call template.render(**prompt_input). Alternatively, decorate
the prompt-builder function with functools.lru_cache keyed by the prompt path
(but ensure prompt_input is not cached—only cache Templates), and replace direct
prompt_file.read_text(...)/Template(...) calls with a cache lookup that returns
a compiled Template.
gaodan-fang
marked this pull request as draft
February 20, 2026 18:32
# Conflicts: # kaizen/backend/milvus.py # kaizen/db/sqlite_manager.py # kaizen/frontend/client/kaizen_client.py
gaodan-fang
marked this pull request as ready for review
February 25, 2026 18:23
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)
418-427:⚠️ Potential issue | 🟠 MajorRemove explicit
idfield from upsert data or disableauto_id=Truein the schema.With
auto_id=Trueon the primary key (line 420), theidfield should not be included in insert/upsert operations. PyMilvus will reject or ignore it, causing the UPDATE path (line 304) to fail. The ADD path (line 289) correctly omitsid, but the UPDATE path passes"id": int(update.id)at line 304, creating an incompatibility.Either:
- Set
auto_id=Falsein the schema and provide explicit IDs in all operations, or- Remove the
"id"field from the upsert data dict and let Milvus generate IDs for both ADD and UPDATE🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 418 - 427, The schema currently sets FieldSchema(name="id", is_primary=True, auto_id=True) in entity_schema but the UPDATE path still includes "id": int(update.id) in the upsert payload, which conflicts with auto-generated IDs; either remove the explicit id from the upsert data or make IDs explicit by changing auto_id=True to auto_id=False. To fix now, delete the code that injects "id" into the upsert dict (the place that does "id": int(update.id)) so both ADD and UPDATE omit the id and Milvus generates IDs, or alternatively change FieldSchema(..., auto_id=True) to FieldSchema(..., auto_id=False) and ensure all insert/update code supplies the id consistently. Ensure you adjust the UPDATE logic to rely on Milvus primary-key handling (filter/expr) when omitting the id.
♻️ Duplicate comments (3)
kaizen/frontend/client/kaizen_client.py (1)
194-201:store_user_memoryis declaredasyncbut contains no awaitable calls.Already flagged in a previous review. All calls inside the body are synchronous.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/frontend/client/kaizen_client.py` around lines 194 - 201, The method store_user_memory is declared async but contains no awaitable calls; remove the async keyword and make it a regular synchronous method (change "async def store_user_memory" to "def store_user_memory") and adjust any call sites that currently await store_user_memory to call it synchronously, or alternatively if you must keep an async API, convert the internal synchronous operations to run in an executor (e.g., use asyncio.to_thread or loop.run_in_executor) so the function actually awaits something; reference function store_user_memory and update callers accordingly.kaizen/llm/fact_extraction/fact_extraction.py (2)
68-79: Retry loop re-parses the same immutable string — retries are no-ops.Already flagged in a previous review —
cleanedis computed once before the loop, sojson.loads(cleaned)will fail identically on every iteration.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/llm/fact_extraction/fact_extraction.py` around lines 68 - 79, The retry loop currently re-parses the same precomputed string `cleaned`, so all three attempts are no-ops; inside the loop either recompute/refresh `cleaned` before calling `json.loads(cleaned)` (e.g., call whatever cleaning/normalization routine that produced `cleaned` each iteration or operate on the raw input) or remove the retry loop and immediately try a single parse and raise the captured `last_error` if it fails; ensure you still call `CategorizedExtractedFacts.model_validate(parsed_json)` and `ExtractedFacts.model_validate(parsed_json)` as before and surface the final exception if parsing never succeeds.
56-57:use_categorizationalways defaults toTrue.Already flagged in a previous review —
llm_settings.categorization_modeisLiteral["predefined", "dynamic", "hybrid"], so the membership check on line 57 is a tautology.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/llm/fact_extraction/fact_extraction.py` around lines 56 - 57, The current default assignment for use_categorization is tautological because llm_settings.categorization_mode is typed as Literal["predefined","dynamic","hybrid"]; update the logic so the default reflects a real off state: either (A) change the check to an explicit comparison like use_categorization = llm_settings.categorization_mode != "none" and update the Literal type to include "none", or (B) read a dedicated flag such as llm_settings.categorization_enabled (use that if it exists) instead of the membership test; modify the assignment in the function where use_categorization is set to use one of these corrected checks (referencing use_categorization and llm_settings.categorization_mode / llm_settings.categorization_enabled).
🧹 Nitpick comments (3)
kaizen/config/llm.py (1)
44-45: Import-time environment mutation — acceptable but worth documenting.
_normalize_litellm_proxy_env()modifiesos.environas a side effect of importing this module. This is intentional (proxy vars must be set beforeLLMSettings()reads them), but any module that importskaizen.config.llmwill trigger these mutations. The existing comment on line 43 covers the reload case; consider adding a brief note that the normalization runs once at import time, to aid future maintainers.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/config/llm.py` around lines 44 - 45, Add a brief comment above the call to _normalize_litellm_proxy_env() (before instantiating LLMSettings()) stating that this function mutates os.environ at import time and therefore will run once whenever the module is imported (so importing kaizen.config.llm triggers the normalization before LLMSettings() reads env vars); reference the symbols _normalize_litellm_proxy_env and LLMSettings in the comment so future maintainers understand the import-time side effect and why it's required.kaizen/frontend/client/kaizen_client.py (1)
232-269:retrieve_user_memorymay issue up to 4 backend queries in the worst case.The cascading fallback (query → no-query → default user → default user no-query) provides good UX but could be slow on a Milvus backend. This is acceptable for now, but consider documenting this behavior in the docstring so callers are aware of the latency profile.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/frontend/client/kaizen_client.py` around lines 232 - 269, Update the retrieve_user_memory docstring to document its cascading fallback behavior and potential latency: explain that it may perform up to four backend search_entities calls in the worst case (initial query, fallback without query, fallback to "default" user with query, and fallback to "default" user without query), note that this happens when namespace exists and limit>0, and that callers should expect higher latency with backends like Milvus so they can cache or adjust usage accordingly.kaizen/backend/milvus.py (1)
345-400: Search path looks correct; minor redundancy on retry embeds the query twice.The split-filter strategy (Milvus handles schema filters, Python handles metadata filters) with overfetching is a solid approach. One minor point: on the
"index not found"retry path (lines 384–392),self.embedding_model.encode(query)is re-computed. Consider caching the embedding above the try block to avoid the redundant encode call.♻️ Cache the query embedding
else: self._ensure_embedding_index(namespace_id) + query_embedding = self.embedding_model.encode(query) try: raw_results = self.milvus.search( collection_name=namespace_id, anns_field="embedding", - data=[self.embedding_model.encode(query)], + data=[query_embedding], filter=self._build_filter_expr(schema_filters), limit=fetch_limit, output_fields=["*"], search_params={"metric_type": self.metric_type}, ) except Exception as exc: if "index not found" in str(exc).lower(): self._ensure_embedding_index(namespace_id) raw_results = self.milvus.search( collection_name=namespace_id, anns_field="embedding", - data=[self.embedding_model.encode(query)], + data=[query_embedding], filter=self._build_filter_expr(schema_filters), limit=fetch_limit, output_fields=["*"], search_params={"metric_type": self.metric_type}, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 345 - 400, In search_entities, avoid calling self.embedding_model.encode(query) twice in the `"index not found"` retry path: compute the embedding once before the try block (e.g. query_embedding = self.embedding_model.encode(query)), then pass query_embedding to both self.milvus.search calls (data=[query_embedding]); keep the rest of the logic (call _ensure_embedding_index, retry on "index not found") unchanged and reference search_entities, _ensure_embedding_index, and embedding_model.encode when locating the code to change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kaizen/backend/milvus.py`:
- Around line 430-439: parse_milvus_entity currently indexes
entity["created_at"] which raises KeyError if that field is missing; change it
to safely use entity.get("created_at") and provide a defensive fallback (e.g.,
if created_at is present convert via int->datetime as now, otherwise set
created_at to datetime.datetime.now(datetime.UTC) or None). Update the
RecordedEntity.model_validate payload construction in parse_milvus_entity to
compute created_at_value = entity.get("created_at") and if created_at_value is
truthy convert it with datetime.datetime.fromtimestamp(int(created_at_value),
datetime.UTC), otherwise use the chosen fallback so no KeyError occurs and
malformed/non-int values are handled.
In `@kaizen/config/llm.py`:
- Around line 19-21: The function _default_model_name currently returns an
expression mypy types as Optional[str]; fix by ensuring a definite str is
returned: either implement explicit checks (e.g., read kaizen =
os.getenv("KAIZEN_MODEL_NAME"); if kaizen is not None: return kaizen; model =
os.getenv("MODEL_NAME"); if model is not None: return model; return "gpt-4o") or
use typing.cast to coerce the final value to str; reference _default_model_name
and the os.getenv calls when making the change.
In `@kaizen/frontend/client/kaizen_client.py`:
- Around line 187-192: The ensure_namespace method can race: after
get_namespace_details raises NamespaceNotFoundException another process may
create the namespace and create_namespace will raise
NamespaceAlreadyExistsException; update ensure_namespace to catch
NamespaceAlreadyExistsException around the create_namespace call, and in that
catch return the existing namespace (e.g., by calling
get_namespace_details(namespace_id) or otherwise fetching the created namespace)
so the method is idempotent; reference ensure_namespace, get_namespace_details,
create_namespace, NamespaceNotFoundException, and
NamespaceAlreadyExistsException when making the change.
In `@kaizen/llm/fact_extraction/fact_extraction.py`:
- Line 35: The timestamp assigned to "current_datetime" in fact_extraction.py is
using naive datetime.datetime.now(); change it to use a timezone-aware call
(datetime.datetime.now(datetime.UTC)) to match the rest of the codebase and
produce UTC timestamps, and add or adjust the imports in this module so
datetime.UTC is available (or use the project-standard UTC constant) to avoid
introducing a local-time value.
---
Outside diff comments:
In `@kaizen/backend/milvus.py`:
- Around line 418-427: The schema currently sets FieldSchema(name="id",
is_primary=True, auto_id=True) in entity_schema but the UPDATE path still
includes "id": int(update.id) in the upsert payload, which conflicts with
auto-generated IDs; either remove the explicit id from the upsert data or make
IDs explicit by changing auto_id=True to auto_id=False. To fix now, delete the
code that injects "id" into the upsert dict (the place that does "id":
int(update.id)) so both ADD and UPDATE omit the id and Milvus generates IDs, or
alternatively change FieldSchema(..., auto_id=True) to FieldSchema(...,
auto_id=False) and ensure all insert/update code supplies the id consistently.
Ensure you adjust the UPDATE logic to rely on Milvus primary-key handling
(filter/expr) when omitting the id.
---
Duplicate comments:
In `@kaizen/frontend/client/kaizen_client.py`:
- Around line 194-201: The method store_user_memory is declared async but
contains no awaitable calls; remove the async keyword and make it a regular
synchronous method (change "async def store_user_memory" to "def
store_user_memory") and adjust any call sites that currently await
store_user_memory to call it synchronously, or alternatively if you must keep an
async API, convert the internal synchronous operations to run in an executor
(e.g., use asyncio.to_thread or loop.run_in_executor) so the function actually
awaits something; reference function store_user_memory and update callers
accordingly.
In `@kaizen/llm/fact_extraction/fact_extraction.py`:
- Around line 68-79: The retry loop currently re-parses the same precomputed
string `cleaned`, so all three attempts are no-ops; inside the loop either
recompute/refresh `cleaned` before calling `json.loads(cleaned)` (e.g., call
whatever cleaning/normalization routine that produced `cleaned` each iteration
or operate on the raw input) or remove the retry loop and immediately try a
single parse and raise the captured `last_error` if it fails; ensure you still
call `CategorizedExtractedFacts.model_validate(parsed_json)` and
`ExtractedFacts.model_validate(parsed_json)` as before and surface the final
exception if parsing never succeeds.
- Around line 56-57: The current default assignment for use_categorization is
tautological because llm_settings.categorization_mode is typed as
Literal["predefined","dynamic","hybrid"]; update the logic so the default
reflects a real off state: either (A) change the check to an explicit comparison
like use_categorization = llm_settings.categorization_mode != "none" and update
the Literal type to include "none", or (B) read a dedicated flag such as
llm_settings.categorization_enabled (use that if it exists) instead of the
membership test; modify the assignment in the function where use_categorization
is set to use one of these corrected checks (referencing use_categorization and
llm_settings.categorization_mode / llm_settings.categorization_enabled).
---
Nitpick comments:
In `@kaizen/backend/milvus.py`:
- Around line 345-400: In search_entities, avoid calling
self.embedding_model.encode(query) twice in the `"index not found"` retry path:
compute the embedding once before the try block (e.g. query_embedding =
self.embedding_model.encode(query)), then pass query_embedding to both
self.milvus.search calls (data=[query_embedding]); keep the rest of the logic
(call _ensure_embedding_index, retry on "index not found") unchanged and
reference search_entities, _ensure_embedding_index, and embedding_model.encode
when locating the code to change.
In `@kaizen/config/llm.py`:
- Around line 44-45: Add a brief comment above the call to
_normalize_litellm_proxy_env() (before instantiating LLMSettings()) stating that
this function mutates os.environ at import time and therefore will run once
whenever the module is imported (so importing kaizen.config.llm triggers the
normalization before LLMSettings() reads env vars); reference the symbols
_normalize_litellm_proxy_env and LLMSettings in the comment so future
maintainers understand the import-time side effect and why it's required.
In `@kaizen/frontend/client/kaizen_client.py`:
- Around line 232-269: Update the retrieve_user_memory docstring to document its
cascading fallback behavior and potential latency: explain that it may perform
up to four backend search_entities calls in the worst case (initial query,
fallback without query, fallback to "default" user with query, and fallback to
"default" user without query), note that this happens when namespace exists and
limit>0, and that callers should expect higher latency with backends like Milvus
so they can cache or adjust usage accordingly.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
CONFIGURATION.mdREADME.mdkaizen/backend/filesystem.pykaizen/backend/milvus.pykaizen/config/llm.pykaizen/db/sqlite_manager.pykaizen/frontend/client/kaizen_client.pykaizen/llm/fact_extraction/__init__.pykaizen/llm/fact_extraction/categorization.pykaizen/llm/fact_extraction/fact_extraction.pytests/unit/test_milvus_backend.py
🚧 Files skipped from review as they are similar to previous changes (3)
- kaizen/db/sqlite_manager.py
- kaizen/llm/fact_extraction/init.py
- kaizen/llm/fact_extraction/categorization.py
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kaizen/config/llm.py`:
- Around line 21-24: Treat empty KAIZEN_MODEL_NAME as unset by checking for a
non-empty, non-whitespace value before returning it: when reading
os.getenv("KAIZEN_MODEL_NAME") into kaizen_model, only return kaizen_model if it
is truthy after trimming (e.g., kaizen_model and kaizen_model.strip());
otherwise fall back to os.getenv("MODEL_NAME", "gpt-4o"). Update the branch that
currently returns kaizen_model directly so empty strings won't short-circuit the
fallback.
- Around line 46-48: The reload path currently calls llm_settings.__init__() but
does not re-run _normalize_litellm_proxy_env(), so changes to LiteLLM→OpenAI env
mappings are ignored; ensure that _normalize_litellm_proxy_env() is invoked
whenever settings are reloaded—either call _normalize_litellm_proxy_env()
immediately before/after llm_settings.__init__() in this module or incorporate
the call into LLMSettings.__init__() so that re-initialization of llm_settings
always reapplies the environment normalization.
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (2)
kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2 (1)
66-67: Enforce strict JSON-only output to reduce parse failures.Line 66 and Line 75 require JSON format, but they do not explicitly forbid extra prose/markdown. Add a strict “JSON object only” constraint.
Proposed fix
-Return the facts and preferences in a json format as shown above. +Return ONLY a valid JSON object in the exact format shown above. Do not add markdown, code fences, or explanatory text. -- Make sure to return the response in the format mentioned in the examples. The response should be in json with a key as "facts" and corresponding value will be a list of objects with category, key, value, and content fields. +- Make sure to return the response in the format mentioned in the examples: a single JSON object with key "facts", where the value is a list of objects with "category", "key", "value", and "content" fields.Also applies to: 75-75
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2` around lines 66 - 67, Update the prompt lines that currently read "Return the facts and preferences in a json format as shown above." (in fact_extraction_predefined.jinja2) to enforce a strict JSON-only response by adding an explicit constraint such as: "Respond with a single JSON object only — no extra prose, no markdown, no code fences, and no surrounding characters; output must be valid JSON parsable by a standard JSON parser." Apply this same change to both occurrences (the line around 66 and the one around 75) so the template requires only the JSON object and forbids any additional text.kaizen/frontend/client/kaizen_client.py (1)
217-225: Avoid persisting unknown extractor payloads asstr(...)facts.Stringifying unexpected objects can degrade memory quality and make retrieval noisy. Prefer skip + warn for non-
ExtractedFactitems.Proposed refactor
for one in extracted: if isinstance(one, ExtractedFact): fact_metadata = dict(base_metadata) fact_metadata["category"] = one.category fact_metadata["key"] = one.key fact_metadata["value"] = one.value entities.append(Entity(type="fact", content=one.content, metadata=fact_metadata)) else: - entities.append(Entity(type="fact", content=str(one), metadata=dict(base_metadata))) + logger.warning("Skipping non-ExtractedFact item from extractor: %r", one) + continue🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/frontend/client/kaizen_client.py` around lines 217 - 225, The code currently stringifies non-ExtractedFact items and appends them as Entity(type="fact") which pollutes storage; update the loop that builds entities so that only instances of ExtractedFact are converted to Entity (use the existing ExtractedFact handling block) and skip any other object while emitting a warning describing the unexpected payload (include its runtime type and a short repr) via the module logger or warnings.warn; ensure you do not append Entity for non-ExtractedFact objects and keep existing behavior for base_metadata and ExtractedFact handling (refer to ExtractedFact, Entity, entities, and base_metadata in kaizen_client.py).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@CONFIGURATION.md`:
- Line 38: Update the statement "All configuration variables are prefixed with
`KAIZEN_`" to scope it to Kaizen-specific settings and explicitly exempt
third-party provider variables like `OPENAI_API_KEY` and `OPENAI_BASE_URL`;
locate the sentence referencing `KAIZEN_` and change it to something like "All
Kaizen-specific configuration variables are prefixed with `KAIZEN_`
(third-party/provider keys such as `OPENAI_API_KEY` and `OPENAI_BASE_URL` use
their own prefixes)." Ensure the text mentions both `OPENAI_API_KEY` and
`OPENAI_BASE_URL` as examples of exceptions so there is no contradiction.
In `@kaizen/backend/milvus.py`:
- Around line 433-439: The current truthy check on created_at_value treats 0 as
false and incorrectly falls back to now; update the condition to explicitly
check for None/empty (e.g., use "if created_at_value is not None and
created_at_value != ''" or simply "if created_at_value is not None") before
attempting to parse it in the try/except around datetime.datetime.fromtimestamp
so that epoch 0 is accepted and parsed into created_at (preserving the existing
exception handling that sets created_at = datetime.datetime.now(datetime.UTC) on
parse errors).
- Around line 73-83: The _split_filters function currently retains keys with
None values which causes client-side filtering to drop results; update
_split_filters to skip entries where value is None before classifying into
schema_filters or metadata_filters (keep using self._schema_filter_fields and
the existing metadata. prefix handling logic in _split_filters to locate where
to change), ensuring behavior matches _build_filter_expr that ignores None.
- Around line 86-93: In _entity_matches_filter, schema post-filter comparisons
fail for created_at because entity.created_at is a datetime while filters are
often epoch ints/strings; update the loop that iterates schema_filters in
_entity_matches_filter to normalize created_at by converting entity.created_at
to an integer epoch (e.g., seconds or milliseconds consistently used by your
filters) and coerce the incoming filter value to int (or parse string to int)
before comparing, keeping the existing special-case for "id" and leaving other
keys unchanged.
In `@kaizen/frontend/client/kaizen_client.py`:
- Around line 206-207: The validation currently checks "if not message: return
[]" but doesn't treat whitespace-only strings as empty; trim the input before
validation (e.g., assign a stripped version like message = message.strip() or
check not message or not message.strip()) and return [] for
None/empty/whitespace-only inputs, and ensure the rest of the method uses the
trimmed message variable so downstream functions (the code that processes
message) don't receive whitespace-only content.
In `@kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2`:
- Line 74: Update the prompt line that currently reads "Create the facts based
on the user and assistant messages only. Do not pick anything from the system
messages." to restrict extraction to user-provided content only: instruct to
"Create facts based on user messages only; do not extract facts from assistant
messages or system messages, except when an assistant message is an explicit,
verbatim restatement of previously user-confirmed information and the user has
explicitly confirmed that restatement." This ensures only user content is
treated as the primary source of truth while allowing a narrow exception for
explicit, confirmed assistant restatements.
---
Nitpick comments:
In `@kaizen/frontend/client/kaizen_client.py`:
- Around line 217-225: The code currently stringifies non-ExtractedFact items
and appends them as Entity(type="fact") which pollutes storage; update the loop
that builds entities so that only instances of ExtractedFact are converted to
Entity (use the existing ExtractedFact handling block) and skip any other object
while emitting a warning describing the unexpected payload (include its runtime
type and a short repr) via the module logger or warnings.warn; ensure you do not
append Entity for non-ExtractedFact objects and keep existing behavior for
base_metadata and ExtractedFact handling (refer to ExtractedFact, Entity,
entities, and base_metadata in kaizen_client.py).
In `@kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2`:
- Around line 66-67: Update the prompt lines that currently read "Return the
facts and preferences in a json format as shown above." (in
fact_extraction_predefined.jinja2) to enforce a strict JSON-only response by
adding an explicit constraint such as: "Respond with a single JSON object only —
no extra prose, no markdown, no code fences, and no surrounding characters;
output must be valid JSON parsable by a standard JSON parser." Apply this same
change to both occurrences (the line around 66 and the one around 75) so the
template requires only the JSON object and forbids any additional text.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
CONFIGURATION.mdkaizen/backend/milvus.pykaizen/config/llm.pykaizen/frontend/client/kaizen_client.pykaizen/llm/fact_extraction/fact_extraction.pykaizen/llm/fact_extraction/prompts/fact_extraction.jinja2kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2
🚧 Files skipped from review as they are similar to previous changes (3)
- kaizen/llm/fact_extraction/fact_extraction.py
- kaizen/config/llm.py
- kaizen/llm/fact_extraction/prompts/fact_extraction.jinja2
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2 (1)
34-65: Consider ensuring few-shot example categories align with runtime categories.The few-shot examples hardcode specific category names (
travel,professional_details,personal_details,entertainment,sports,health), but the actual valid categories are provided dynamically via thecategoriestemplate variable. If the runtime categories differ in naming (e.g.,workvsprofessional_details), the model may produce categories that don't match expectations or fail downstream validation.Two options to improve robustness:
- Document that the runtime
categoriesmust include these exact names, or- Dynamically generate or include a representative subset of few-shot examples that use the actual runtime category names
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2` around lines 34 - 65, The few-shot examples in fact_extraction_predefined.jinja2 use hardcoded category names that may not match the runtime `categories` variable; update the template so examples are generated or selected from the actual `categories` list (the `categories` template variable) instead of fixed strings, e.g., produce representative example lines by iterating over `categories` or include a small dynamically-built subset mapping each example's "category" field to values from `categories` (refer to the few-shot examples block and the `categories` variable in this template to locate where to change).kaizen/backend/milvus.py (1)
370-372: Overfetch only when Python-side filtering is actually needed.Line 371 overfetches to 1000 for any filter, including schema filters that are already pushed down to Milvus. This adds avoidable query load.
♻️ Proposed refinement
- fetch_limit = max(limit, 1000) if filters else limit + fetch_limit = max(limit, 1000) if metadata_filters else limit🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 370 - 372, The code overfetches whenever any filters exist; change it to only overfetch when Python-side filtering is needed by looking at metadata_filters (from _split_filters) not the original filters. After calling schema_filters, metadata_filters = self._split_filters(filters), compute fetch_limit = max(limit, 1000) if metadata_filters else limit (or equivalent truthy check), so that schema-only filters pushed to Milvus do not trigger the 1000 overfetch.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kaizen/backend/milvus.py`:
- Around line 43-44: The code currently replaces any non-MilvusDBSettings
caller-provided config with milvus_client_settings silently; change this to
validate and surface invalid inputs instead of overwriting them: in the
constructor (where resolved_config is computed) check if config is provided and
not an instance of MilvusDBSettings, then raise a TypeError (or ValueError)
explaining the expected type and include the received type, otherwise use the
provided MilvusDBSettings; fall back to milvus_client_settings only when config
is None. Update the logic around resolved_config/self.config to reflect this
validation.
- Around line 94-107: The code uses getattr(entity, key, None) into entity_value
and then calls isinstance(entity_value, datetime.datetime) and
int(entity_value), which leaves mypy complaining because entity_value may be
None; add an explicit None guard (e.g., if entity_value is None: return False)
before the isinstance(...) block to narrow the type, then proceed to compute
entity_epoch_seconds/entity_epoch_milliseconds and the later int(value) parse as
before; update references to entity_value, entity_epoch_seconds, and
entity_epoch_milliseconds in the existing branch so mypy knows entity_value is
not None.
In `@tests/unit/test_client.py`:
- Around line 196-215: The test defines a stub ensure_namespace but never
verifies it was invoked; modify the ensure_namespace stub used with
monkeypatch.setattr(kaizen_client, "ensure_namespace", ensure_namespace) to
record a flag (e.g., set captured["ensure_namespace_called"] = True) when
called, then after calling kaizen_client.store_user_memory(...) assert
captured["ensure_namespace_called"] is True so the trimmed-message flow actually
exercises ensure_namespace; keep existing captured checks for message_content
and entity_content and reference the ensure_namespace stub and store_user_memory
call when adding this assertion.
---
Nitpick comments:
In `@kaizen/backend/milvus.py`:
- Around line 370-372: The code overfetches whenever any filters exist; change
it to only overfetch when Python-side filtering is needed by looking at
metadata_filters (from _split_filters) not the original filters. After calling
schema_filters, metadata_filters = self._split_filters(filters), compute
fetch_limit = max(limit, 1000) if metadata_filters else limit (or equivalent
truthy check), so that schema-only filters pushed to Milvus do not trigger the
1000 overfetch.
In `@kaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2`:
- Around line 34-65: The few-shot examples in fact_extraction_predefined.jinja2
use hardcoded category names that may not match the runtime `categories`
variable; update the template so examples are generated or selected from the
actual `categories` list (the `categories` template variable) instead of fixed
strings, e.g., produce representative example lines by iterating over
`categories` or include a small dynamically-built subset mapping each example's
"category" field to values from `categories` (refer to the few-shot examples
block and the `categories` variable in this template to locate where to change).
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
.secrets.baselineREADME.mdkaizen/backend/milvus.pykaizen/frontend/client/kaizen_client.pykaizen/llm/fact_extraction/prompts/fact_extraction_predefined.jinja2tests/unit/test_client.pytests/unit/test_milvus_backend.py
🚧 Files skipped from review as they are similar to previous changes (2)
- README.md
- kaizen/frontend/client/kaizen_client.py
Comment on lines
+94
to
+107
| elif key == "created_at": | ||
| if isinstance(entity_value, datetime.datetime): | ||
| entity_epoch_seconds = int(entity_value.timestamp()) | ||
| entity_epoch_milliseconds = int(entity_value.timestamp() * 1000) | ||
| else: | ||
| try: | ||
| entity_epoch_seconds = int(entity_value) | ||
| except (TypeError, ValueError): | ||
| return False | ||
| entity_epoch_milliseconds = entity_epoch_seconds * 1000 | ||
|
|
||
| try: | ||
| filter_epoch = int(value) | ||
| except (TypeError, ValueError): |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# Read the file to see the current state of lines 94-107
sed -n '94,107p' kaizen/backend/milvus.pyRepository: AgentToolkit/kaizen
Length of output: 715
🏁 Script executed:
# Get more context around the created_at handling to understand the issue
sed -n '90,115p' kaizen/backend/milvus.pyRepository: AgentToolkit/kaizen
Length of output: 1150
🏁 Script executed:
# Check if there's a None guard before int(entity_value)
rg -n -B5 -A5 'elif key == "created_at"' kaizen/backend/milvus.pyRepository: AgentToolkit/kaizen
Length of output: 619
🏁 Script executed:
# Get function context around line 90 to see signature and type hints
sed -n '75,95p' kaizen/backend/milvus.pyRepository: AgentToolkit/kaizen
Length of output: 972
🏁 Script executed:
# Check if there are type annotations for entity_value or return type
rg -n -B15 'for key, value in schema_filters.items' kaizen/backend/milvus.py | head -40Repository: AgentToolkit/kaizen
Length of output: 785
🏁 Script executed:
# Check type annotations on RecordedEntity to understand entity attribute types
rg -n 'class RecordedEntity' kaizen/backend/milvus.pyRepository: AgentToolkit/kaizen
Length of output: 45
🏁 Script executed:
# Look for imports and see if there's type stub info
head -30 kaizen/backend/milvus.pyRepository: AgentToolkit/kaizen
Length of output: 1101
🏁 Script executed:
# Find RecordedEntity definition in the schema module
rg -n 'class RecordedEntity' kaizen/schema/Repository: AgentToolkit/kaizen
Length of output: 117
🏁 Script executed:
# Search for RecordedEntity definition to understand created_at type
rg -n 'created_at' kaizen/schema/core.py | head -20Repository: AgentToolkit/kaizen
Length of output: 242
🏁 Script executed:
# Get the full RecordedEntity class definition
sed -n '27,50p' kaizen/schema/core.pyRepository: AgentToolkit/kaizen
Length of output: 301
🏁 Script executed:
# Check what getattr returns - it returns Any | None in typing terms
# Let's verify if there are type stubs or if mypy is configured
cat kaizen/backend/milvus.py | grep -A30 '_entity_matches_filter'Repository: AgentToolkit/kaizen
Length of output: 2689
Add explicit None guard before int(entity_value) to satisfy mypy type checking.
Line 100 attempts to coerce entity_value to int, but entity_value is typed as Any | None from the getattr(entity, key, None) call. While the try-except catches the resulting TypeError, mypy still reports a type error. Add an explicit None check before the isinstance block to narrow the type.
Proposed fix
elif key == "created_at":
+ if entity_value is None:
+ return False
if isinstance(entity_value, datetime.datetime):
entity_epoch_seconds = int(entity_value.timestamp())
entity_epoch_milliseconds = int(entity_value.timestamp() * 1000)
else:
try:
entity_epoch_seconds = int(entity_value)
except (TypeError, ValueError):
return False🧰 Tools
🪛 GitHub Actions: Check Formatting with Ruff
[error] 100-100: mypy: Argument 1 to 'int' has incompatible type 'Any | None'; expected 'str | Buffer | SupportsInt | SupportsIndex | SupportsTrunc' [arg-type]
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@kaizen/backend/milvus.py` around lines 94 - 107, The code uses
getattr(entity, key, None) into entity_value and then calls
isinstance(entity_value, datetime.datetime) and int(entity_value), which leaves
mypy complaining because entity_value may be None; add an explicit None guard
(e.g., if entity_value is None: return False) before the isinstance(...) block
to narrow the type, then proceed to compute
entity_epoch_seconds/entity_epoch_milliseconds and the later int(value) parse as
before; update references to entity_value, entity_epoch_seconds, and
entity_epoch_milliseconds in the existing branch so mypy knows entity_value is
not None.
illeatmyhat
approved these changes
Mar 2, 2026
illeatmyhat
left a comment
illeatmyhat
left a comment
Collaborator
There was a problem hiding this comment.
mildly sloppy but the important stuff seems good
Comment on lines
+8
to
+20
| def _default_model_name() -> str: | ||
| model_name = os.getenv("KAIZEN_MODEL_NAME") | ||
| if model_name and model_name.strip(): | ||
| return model_name.strip() | ||
| return "gpt-4o" | ||
|
|
||
|
|
||
| def _default_custom_provider() -> str | None: | ||
| # If OpenAI env vars are configured, default provider to openai. | ||
| # Explicit KAIZEN_CUSTOM_LLM_PROVIDER still has higher priority via BaseSettings. | ||
| if os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_KEY"): | ||
| return "openai" | ||
| return None |
Collaborator
There was a problem hiding this comment.
not necessary.
You can use the Field(alias='...') to get rid of _default_model_name()
You can use AliasChoices to get rid of _default_custom_provider()
|
|
||
| entity_schema = CollectionSchema( | ||
| fields=[ | ||
| # Keep it as an INT64 or else you won't be able to list all entities. |
Collaborator
There was a problem hiding this comment.
important comment; milvus doesn't support auto-incremented string IDs
Comment on lines
235
to
+237
| if not self.milvus.has_collection(namespace_id): | ||
| self.milvus.create_collection(collection_name=namespace_id, dimension=384, auto_id=False, schema=entity_schema) | ||
| index_params = self.milvus.prepare_index_params() | ||
| index_params.add_index(field_name="embedding", metric_type="IP", index_type="FLAT") | ||
| self.milvus.create_index(collection_name=namespace_id, index_params=index_params) | ||
| self._ensure_embedding_index(namespace_id) |
Collaborator
There was a problem hiding this comment.
this will output a warning every single time you create a new namespace which I'm sure isn't the intended UX
Comment on lines
+168
to
+197
| @staticmethod | ||
| def _normalize_search_hit(hit) -> dict: | ||
| if hasattr(hit, "to_dict"): | ||
| try: | ||
| hit = hit.to_dict() | ||
| except Exception: | ||
| pass | ||
|
|
||
| if not isinstance(hit, dict): | ||
| normalized = {} | ||
| for attr in ("id", "distance", "score"): | ||
| if hasattr(hit, attr): | ||
| normalized[attr] = getattr(hit, attr) | ||
| entity_attr = getattr(hit, "entity", None) | ||
| if entity_attr is not None and hasattr(entity_attr, "to_dict"): | ||
| try: | ||
| entity_attr = entity_attr.to_dict() | ||
| except Exception: | ||
| entity_attr = None | ||
| if isinstance(entity_attr, dict): | ||
| normalized.update(entity_attr) | ||
| return normalized | ||
|
|
||
| entity = hit.get("entity") | ||
| normalized = {} | ||
| if isinstance(entity, dict): | ||
| normalized.update(entity) | ||
| normalized.update(hit) | ||
| normalized.pop("entity", None) | ||
| return normalized |
Collaborator
There was a problem hiding this comment.
seems overly defensive
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.
Summary by CodeRabbit
New Features
Chores
Documentation
Tests