Open
Conversation
Please make sure all the checkboxes are checked:
|
Contributor
WalkthroughAdds a new agent-memory subsystem: decorator, runtime, models, LLM input injection, persistence, example, CI job, and unit/integration/e2e tests; exposes Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
🚥 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 docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
cognee/modules/agent_memory/models.py (2)
34-37: Consider documenting themodel_post_initbehavior.The method overrides both
idandbelongs_to_setunconditionally, which means user-provided values are silently discarded. This is likely intentional for deterministic tracing, but a brief docstring would clarify this design choice.📝 Suggested docstring
def model_post_init(self, __context: Any) -> None: + """Recompute deterministic id from text and reset belongs_to_set to agent_traces.""" super().model_post_init(__context) self.id = uuid5(NAMESPACE_OID, f"AgentTrace:{self.text}") self.belongs_to_set = [_agent_traces_nodeset()]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/models.py` around lines 34 - 37, Add a concise docstring to the model_post_init method explaining that it intentionally overrides any user-provided id and belongs_to_set values to produce deterministic tracing identifiers (it sets self.id via uuid5 with NAMESPACE_OID and self.text and sets self.belongs_to_set to [_agent_traces_nodeset()]), and note the side-effect that incoming values are discarded so callers should not rely on passing those fields before post-init; reference the method name model_post_init and the attributes id and belongs_to_set as well as uuid5/NAMESPACE_OID/_agent_traces_nodeset to make the behavior discoverable.
32-32: Mutable class-level default may cause shared state issues.The
metadatadict is defined as a class attribute with a mutable default. If any code modifiesinstance.metadata, it will affect all instances. Consider usingField(default_factory=...)for safety.♻️ Proposed fix
- metadata: dict = {"index_fields": ["text"]} + metadata: dict = Field(default_factory=lambda: {"index_fields": ["text"]})🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/models.py` at line 32, The class-level mutable default for the attribute `metadata` can cause shared-state bugs across instances; change the attribute to use a Field default factory (e.g., Field(default_factory=lambda: {"index_fields": ["text"]})) so each instance gets its own dict; locate the `metadata` attribute in the model class in this file and replace the class-level dict with a Field(default_factory=...) import from pydantic (or the appropriate model library) and ensure tests/usage still expect the same structure.cognee/modules/agent_memory/runtime.py (2)
84-130: Consider validatingmemory_top_kbounds.The validation checks types for most parameters but doesn't validate that
memory_top_kis a positive integer. Negative or zero values could cause unexpected behavior in the search call.🛡️ Proposed validation
if dataset_name is not None and (not isinstance(dataset_name, str) or not dataset_name.strip()): raise CogneeValidationError( "dataset_name must be a non-empty string when provided.", log=False, ) + if not isinstance(memory_top_k, int) or memory_top_k < 1: + raise CogneeValidationError( + "memory_top_k must be a positive integer.", + log=False, + ) return AgentMemoryConfig(🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/runtime.py` around lines 84 - 130, The validate_agent_memory_config function lacks bounds checking for memory_top_k; add a validation that memory_top_k is an integer greater than 0 and raise CogneeValidationError with a clear message (e.g., "memory_top_k must be a positive integer") if not; update validate_agent_memory_config (and its signature handling of memory_top_k) so the check happens before constructing AgentMemoryConfig, referencing memory_top_k, CogneeValidationError, validate_agent_memory_config, and AgentMemoryConfig to locate the change.
289-290: Redundantasyncio.create_taskwrapping.
await asyncio.create_task(coro)is equivalent toawait corobut adds unnecessary overhead. If the intent is fire-and-forget persistence, remove theawait. If synchronous completion is required (current behavior), call the coroutine directly.♻️ Proposed fix for synchronous execution
try: - await asyncio.create_task(_persist_trace_in_task()) + await _persist_trace_in_task() except Exception as error:Or for true fire-and-forget (trace persists in background):
♻️ Alternative for fire-and-forget
try: - await asyncio.create_task(_persist_trace_in_task()) + asyncio.create_task(_persist_trace_in_task()) except Exception as error:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/runtime.py` around lines 289 - 290, The current call wraps the coroutine _persist_trace_in_task() in asyncio.create_task and immediately awaits it, which is redundant; either (A) if you need synchronous completion, replace await asyncio.create_task(_persist_trace_in_task()) with awaiting the coroutine directly by calling await _persist_trace_in_task(), or (B) if you want fire-and-forget persistence, remove the await and just call asyncio.create_task(_persist_trace_in_task()) so the task runs in background—update the invocation accordingly in runtime.py where _persist_trace_in_task is called.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/guides/agent_memory_quickstart.py`:
- Around line 34-39: Update the type hints so primitive types are allowed:
change the `response_model` parameter type in
`LLMGateway.acreate_structured_output` and the interface signature in
`llm_interface.py` from `Type[BaseModel]` to `Union[Type[BaseModel], type]`, and
update adapter method return annotations from `BaseModel` to `Union[BaseModel,
Any]` (or `Union[BaseModel, Any]` where used) so implementations that return
primitives (e.g., `str`, `int`, `bool`) are correctly typed; adjust all adapter
implementations' signatures to match (references:
`LLMGateway.acreate_structured_output`, the interface method in
`llm_interface.py`, and each adapter class/method that declares `response_model`
or returns parsed output).
---
Nitpick comments:
In `@cognee/modules/agent_memory/models.py`:
- Around line 34-37: Add a concise docstring to the model_post_init method
explaining that it intentionally overrides any user-provided id and
belongs_to_set values to produce deterministic tracing identifiers (it sets
self.id via uuid5 with NAMESPACE_OID and self.text and sets self.belongs_to_set
to [_agent_traces_nodeset()]), and note the side-effect that incoming values are
discarded so callers should not rely on passing those fields before post-init;
reference the method name model_post_init and the attributes id and
belongs_to_set as well as uuid5/NAMESPACE_OID/_agent_traces_nodeset to make the
behavior discoverable.
- Line 32: The class-level mutable default for the attribute `metadata` can
cause shared-state bugs across instances; change the attribute to use a Field
default factory (e.g., Field(default_factory=lambda: {"index_fields":
["text"]})) so each instance gets its own dict; locate the `metadata` attribute
in the model class in this file and replace the class-level dict with a
Field(default_factory=...) import from pydantic (or the appropriate model
library) and ensure tests/usage still expect the same structure.
In `@cognee/modules/agent_memory/runtime.py`:
- Around line 84-130: The validate_agent_memory_config function lacks bounds
checking for memory_top_k; add a validation that memory_top_k is an integer
greater than 0 and raise CogneeValidationError with a clear message (e.g.,
"memory_top_k must be a positive integer") if not; update
validate_agent_memory_config (and its signature handling of memory_top_k) so the
check happens before constructing AgentMemoryConfig, referencing memory_top_k,
CogneeValidationError, validate_agent_memory_config, and AgentMemoryConfig to
locate the change.
- Around line 289-290: The current call wraps the coroutine
_persist_trace_in_task() in asyncio.create_task and immediately awaits it, which
is redundant; either (A) if you need synchronous completion, replace await
asyncio.create_task(_persist_trace_in_task()) with awaiting the coroutine
directly by calling await _persist_trace_in_task(), or (B) if you want
fire-and-forget persistence, remove the await and just call
asyncio.create_task(_persist_trace_in_task()) so the task runs in
background—update the invocation accordingly in runtime.py where
_persist_trace_in_task is called.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 1e904c33-9b49-440e-8143-2b8b664a3278
📒 Files selected for processing (7)
cognee/__init__.pycognee/infrastructure/llm/LLMGateway.pycognee/modules/agent_memory/__init__.pycognee/modules/agent_memory/decorator.pycognee/modules/agent_memory/models.pycognee/modules/agent_memory/runtime.pyexamples/guides/agent_memory_quickstart.py
hajdul88
changed the title
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
cognee/tests/test_agent_memory_e2e.py (1)
1-2: Move this end-to-end test under a supportedcognee/tests/subdirectory.Keeping
cognee/tests/test_agent_memory_e2e.pyat the top level breaks the repo's test layout convention and forces a one-off workflow target. Moving it undercognee/tests/integration/(or another approved subdirectory) will keep discovery and CI wiring consistent. As per coding guidelines,cognee/tests/**/*.py: "Organize tests in cognee/tests/ directory with subdirectories: unit/, integration/, cli_tests/, tasks/."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/tests/test_agent_memory_e2e.py` around lines 1 - 2, Move the end-to-end test file from cognee/tests/test_agent_memory_e2e.py into an approved subdirectory under cognee/tests (e.g., cognee/tests/integration/test_agent_memory_e2e.py); update any imports or test discovery references that point to the old module path (search for references to test_agent_memory_e2e or the top-level path cognee.tests.test_agent_memory_e2e) and ensure the test module lives under the cognee/tests/** hierarchy so CI and pytest discovery follow the project's convention.cognee/tests/integration/modules/agent_memory/test_agent_memory_integration.py (1)
146-167: Either remove theLLMGatewaymonkeypatch or route the assertion through it.
shared_memory_agent()returnscontext.memory_contextdirectly, so this monkeypatch never affects the test outcome. As written, the test validates retrieval only; a regression in the downstream memory-injection path would still pass.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/tests/integration/modules/agent_memory/test_agent_memory_integration.py` around lines 146 - 167, The test currently monkeypatches LLMGateway.acreate_structured_output but returns context.memory_context directly in shared_memory_agent(), so the monkeypatch is unused; either remove the monkeypatch block entirely, or change shared_memory_agent() to route the memory through the gateway (call LLMGateway.acreate_structured_output with text_input=context.memory_context and assert on the gateway's return value) so the monkeypatched acreate_structured_output is exercised; update references to get_current_agent_memory_context(), shared_memory_agent, and the LLMGateway.acreate_structured_output call accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cognee/modules/agent_memory/runtime.py`:
- Around line 137-140: The resolve_agent_scope function currently falls back to
get_default_user() when neither config.user nor session_user.get() exists;
change this to fail closed: in resolve_agent_scope (working with
AgentMemoryConfig and returning AgentScope) check config.user or
session_user.get(), and if neither is present, raise an explicit exception
(e.g., PermissionError or a custom UnauthorizedError) instead of calling
get_default_user(); remove or gate the get_default_user() fallback so
access-controlled flows never inherit a different principal.
- Around line 87-134: The current validate_agent_memory_config lets blank memory
queries, invalid memory_top_k, and partial user-like objects slip through and
cause later failures; update validate_agent_memory_config to (1) require
memory_query_fixed and memory_query_from_method when provided to be non-empty
after .strip(), (2) enforce memory_top_k is an int >= 1 (and optionally bound it
to a sensible max), and (3) require the user argument to be a real User model
(use isinstance(user, User)) rather than merely having an id; when normalizing,
strip and store the trimmed query strings and raise CogneeValidationError with
clear messages on violations so resolve_agent_scope and search receive validated
inputs.
- Around line 180-184: The code currently uses sanitize_value() (which only
truncates) in build_method_params, allowing secrets/PII (tokens, passwords,
emails, raw user content) to be persisted in AgentTrace; update
build_method_params to perform redaction instead of just truncation: detect
sensitive parameter names (e.g., password, pass, pwd, token, secret, api_key,
authorization, email, ssn, creditcard) and mask or replace their values with a
constant like "[REDACTED]" and also apply pattern-based redaction for values
that look like emails, JWTs, credit-card numbers, or long opaque tokens;
likewise apply the same redaction step to the code that produces
method_return_value (the return-value serialization around lines 324-335) so
return payloads are scrubbed before writing to AgentTrace. Ensure you centralize
the logic (e.g., a redact_sensitive_value helper) and call it from
build_method_params and the return-value handling code so all persisted traces
are redacted.
In
`@cognee/tests/integration/modules/agent_memory/test_agent_memory_integration.py`:
- Around line 78-92: The fixture agent_memory_integration_env mutates global
state (os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] and
cognee.config.data_root_directory / cognee.config.system_root_directory) but
never restores it; update the fixture to snapshot the prior values (store
previous os.environ value or absence and previous data/system root values)
before mutating, then in the teardown (after yield, before/after calling
_reset_engines_and_prune()) restore the original environment key and call
cognee.config.data_root_directory(...) and
cognee.config.system_root_directory(...) with the previous values (or unset the
env var if it was absent); alternatively replace these mutations with pytest
monkeypatch.setenv and monkeypatch.setattr to scope changes automatically around
the yield while keeping engine_setup() and _reset_engines_and_prune() calls
unchanged.
In `@cognee/tests/test_agent_memory_e2e.py`:
- Around line 70-85: The fixture agent_memory_e2e_env mutates global state
(os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] and
cognee.config.data_root_directory/system_root_directory) but doesn't restore it;
update the fixture to snapshot and restore those values (store previous
os.environ value or absence and previous data/system root directory values, then
restore them after the yield) or accept and use pytest's monkeypatch fixture
(use monkeypatch.setenv and monkeypatch.setenv/monkeypatch.setattr for
cognee.config.*) so that state is reverted in teardown alongside the existing
_reset_engines_and_prune/engine_setup calls.
---
Nitpick comments:
In
`@cognee/tests/integration/modules/agent_memory/test_agent_memory_integration.py`:
- Around line 146-167: The test currently monkeypatches
LLMGateway.acreate_structured_output but returns context.memory_context directly
in shared_memory_agent(), so the monkeypatch is unused; either remove the
monkeypatch block entirely, or change shared_memory_agent() to route the memory
through the gateway (call LLMGateway.acreate_structured_output with
text_input=context.memory_context and assert on the gateway's return value) so
the monkeypatched acreate_structured_output is exercised; update references to
get_current_agent_memory_context(), shared_memory_agent, and the
LLMGateway.acreate_structured_output call accordingly.
In `@cognee/tests/test_agent_memory_e2e.py`:
- Around line 1-2: Move the end-to-end test file from
cognee/tests/test_agent_memory_e2e.py into an approved subdirectory under
cognee/tests (e.g., cognee/tests/integration/test_agent_memory_e2e.py); update
any imports or test discovery references that point to the old module path
(search for references to test_agent_memory_e2e or the top-level path
cognee.tests.test_agent_memory_e2e) and ensure the test module lives under the
cognee/tests/** hierarchy so CI and pytest discovery follow the project's
convention.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 5f3cf6b5-6f12-4a40-b33d-2673a44a249f
📒 Files selected for processing (8)
.github/workflows/e2e_tests.ymlcognee/modules/agent_memory/decorator.pycognee/modules/agent_memory/models.pycognee/modules/agent_memory/runtime.pycognee/tests/integration/modules/agent_memory/test_agent_memory_integration.pycognee/tests/test_agent_memory_e2e.pycognee/tests/unit/modules/agent_memory/test_agent_memory.pyexamples/guides/agent_memory_quickstart.py
✅ Files skipped from review due to trivial changes (1)
- examples/guides/agent_memory_quickstart.py
🚧 Files skipped from review as they are similar to previous changes (2)
- cognee/modules/agent_memory/models.py
- cognee/modules/agent_memory/decorator.py
cognee/tests/integration/modules/agent_memory/test_agent_memory_integration.py
Show resolved
Hide resolved
Contributor
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 `@cognee/modules/agent_memory/runtime.py`:
- Around line 356-369: The normalize_search_results function currently treats
Pydantic SearchResult instances as generic objects and falls back to
str(results), losing the nested search_result field; update
normalize_search_results to detect Pydantic models (e.g., check hasattr(results,
"search_result") or isinstance(results, BaseModel)) before the final fallback
and call normalize_search_results(results.search_result) so the structured
search_result content is preserved when normalizing outputs from search() which
returns List[SearchResult].
- Around line 305-308: set_database_global_context_variables is being called
with context.scope.dataset_owner_id which is Optional[UUID] but the function
signature requires a non-optional UUID; replace the optional with a guaranteed
user id fallback by passing context.scope.user.id when
context.scope.dataset_owner_id is None (i.e., call
set_database_global_context_variables with context.scope.dataset_id and
context.scope.dataset_owner_id or context.scope.user.id as the user_id argument)
to avoid the type mismatch.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 0a2f32f7-4d71-4f5e-8b87-47b9989797d7
📒 Files selected for processing (2)
cognee/modules/agent_memory/runtime.pycognee/tests/unit/modules/agent_memory/test_agent_memory.py
✅ Files skipped from review due to trivial changes (1)
- cognee/tests/unit/modules/agent_memory/test_agent_memory.py
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (2)
cognee/modules/agent_memory/runtime.py (2)
305-308:⚠️ Potential issue | 🟠 MajorType mismatch:
dataset_owner_idisOptional[UUID]but the function expectsUUID.
context.scope.dataset_owner_idis typed asOptional[UUID](line 44), butset_database_global_context_variablesexpectsuser_id: UUIDwithoutOptional. PassingNonewould cause a runtime error when access control is enabled.Proposed fix using user.id as fallback
await set_database_global_context_variables( context.scope.dataset_id, - context.scope.dataset_owner_id, + context.scope.dataset_owner_id or context.scope.user.id, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/runtime.py` around lines 305 - 308, context.scope.dataset_owner_id is Optional[UUID] but set_database_global_context_variables requires a non-optional UUID; update the call in runtime.py to pass a guaranteed UUID (e.g., use context.scope.dataset_owner_id if not None else context.user.id) or raise a clear error if neither is available. Locate the call to set_database_global_context_variables and change the argument for user_id to a validated/fallback value derived from context.scope.dataset_owner_id or context.user.id so a None is never passed at runtime.
358-371:⚠️ Potential issue | 🟡 MinorHandle Pydantic
SearchResultmodels to preserve structured data.The
search()function returnsList[SearchResult](PydanticBaseModel). Since Pydantic models are notdictinstances, they fall through tostr(results)(line 371), losing the structuredsearch_resultfield content.Proposed fix to handle objects with search_result attribute
if isinstance(results, dict): if "search_result" in results: return normalize_search_results(results["search_result"]) return json.dumps(sanitize_value(results), ensure_ascii=False) + if hasattr(results, "search_result"): + return normalize_search_results(results.search_result) return str(results)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/runtime.py` around lines 358 - 371, normalize_search_results currently treats Pydantic SearchResult models as generic objects and falls back to str(...), losing structured fields; update normalize_search_results to detect Pydantic models or any object with a "search_result" attribute (or a .dict() method) and extract/normalize that field instead of stringifying the model. Specifically, in normalize_search_results add a branch that checks for hasattr(results, "search_result") and returns normalize_search_results(results.search_result), and/or for objects with a .dict() method call sanitize_value(results.dict()) and json.dumps(...) to preserve structured data; keep existing handling for dict, list, str, None, and the final fallback.
🧹 Nitpick comments (1)
cognee/modules/agent_memory/runtime.py (1)
311-312: Redundantawait asyncio.create_task()pattern.
await asyncio.create_task(coro())creates a task then immediately awaits it, adding overhead without benefit. If the goal is isolated ContextVar scope, consider adding a comment explaining this. Otherwise, simplify to a direct await.Option 1: Direct await (if blocking is acceptable)
try: - await asyncio.create_task(_persist_trace_in_task()) + await _persist_trace_in_task() except Exception as error:Option 2: Fire-and-forget (if non-blocking is intended)
- try: - await asyncio.create_task(_persist_trace_in_task()) - except Exception as error: - logger.warning( - "Agent trace persistence failed for %s: %s", - context.origin_function, - error, - exc_info=False, - ) + task = asyncio.create_task(_persist_trace_in_task()) + task.add_done_callback( + lambda t: logger.warning( + "Agent trace persistence failed for %s: %s", + context.origin_function, + t.exception(), + exc_info=False, + ) if t.exception() else None + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/runtime.py` around lines 311 - 312, Replace the redundant await asyncio.create_task(_persist_trace_in_task()) pattern in runtime.py: either directly await the coroutine by calling await _persist_trace_in_task() if you intend to wait for completion, or call asyncio.create_task(_persist_trace_in_task()) without awaiting it for fire-and-forget behavior; if you used create_task solely to isolate ContextVar scope, keep create_task but add a short comment referencing that intent near the call and ensure the task is not immediately awaited.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@cognee/modules/agent_memory/runtime.py`:
- Around line 305-308: context.scope.dataset_owner_id is Optional[UUID] but
set_database_global_context_variables requires a non-optional UUID; update the
call in runtime.py to pass a guaranteed UUID (e.g., use
context.scope.dataset_owner_id if not None else context.user.id) or raise a
clear error if neither is available. Locate the call to
set_database_global_context_variables and change the argument for user_id to a
validated/fallback value derived from context.scope.dataset_owner_id or
context.user.id so a None is never passed at runtime.
- Around line 358-371: normalize_search_results currently treats Pydantic
SearchResult models as generic objects and falls back to str(...), losing
structured fields; update normalize_search_results to detect Pydantic models or
any object with a "search_result" attribute (or a .dict() method) and
extract/normalize that field instead of stringifying the model. Specifically, in
normalize_search_results add a branch that checks for hasattr(results,
"search_result") and returns normalize_search_results(results.search_result),
and/or for objects with a .dict() method call sanitize_value(results.dict()) and
json.dumps(...) to preserve structured data; keep existing handling for dict,
list, str, None, and the final fallback.
---
Nitpick comments:
In `@cognee/modules/agent_memory/runtime.py`:
- Around line 311-312: Replace the redundant await
asyncio.create_task(_persist_trace_in_task()) pattern in runtime.py: either
directly await the coroutine by calling await _persist_trace_in_task() if you
intend to wait for completion, or call
asyncio.create_task(_persist_trace_in_task()) without awaiting it for
fire-and-forget behavior; if you used create_task solely to isolate ContextVar
scope, keep create_task but add a short comment referencing that intent near the
call and ensure the task is not immediately awaited.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 170e99f5-e7e8-4d66-a2a5-c3547394a3b8
📒 Files selected for processing (1)
cognee/modules/agent_memory/runtime.py
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (2)
cognee/modules/agent_memory/runtime.py (2)
331-345:⚠️ Potential issue | 🟠 MajorPersisted traces still retain potentially sensitive payloads.
The TODO is acknowledged, but current persistence still stores
method_paramsandmethod_return_valuewithout redaction (Lines 339-340). Truncation alone does not mitigate secrets/PII retention risk.Proposed direction
+def redact_sensitive_value(key: Optional[str], value: Any) -> Any: + # mask by sensitive key names + known token/email patterns, then sanitize/truncate + ... + def build_agent_trace(context: AgentMemoryContext) -> AgentTrace: @@ - return AgentTrace( + redacted_params = { + key: redact_sensitive_value(key, value) + for key, value in context.method_params.items() + } + redacted_return = redact_sensitive_value(None, context.method_return_value) + return AgentTrace( @@ - method_params=context.method_params, - method_return_value=sanitize_value(context.method_return_value), + method_params=redacted_params, + method_return_value=redacted_return,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/runtime.py` around lines 331 - 345, build_agent_trace is persisting raw method_params and method_return_value (AgentTrace.method_params and method_return_value) which risks storing secrets/PII; modify build_agent_trace to sanitize/redact these fields before constructing the AgentTrace: run both context.method_params and context.method_return_value through a dedicated redaction routine (e.g., redact_sensitive_data or extend sanitize_value) that removes or masks secrets and PII and then truncate if needed (use truncate_text/MAX_SERIALIZED_VALUE_LENGTH) before assigning to AgentTrace.method_params and AgentTrace.method_return_value so only redacted/truncated payloads are persisted.
84-141:⚠️ Potential issue | 🟠 MajorValidate
memory_top_kat the decorator boundary.
memory_top_kis accepted without type/range checks (Line 138). Invalid values fail later inside search/retrieval and can degrade into silent empty-memory behavior instead of a clear config error.Proposed fix
def validate_agent_memory_config( @@ - if user is not None and not hasattr(user, "id"): + if not isinstance(memory_top_k, int) or isinstance(memory_top_k, bool): + raise CogneeValidationError("memory_top_k must be an integer.", log=False) + if memory_top_k < 1: + raise CogneeValidationError("memory_top_k must be >= 1.", log=False) + if user is not None and not hasattr(user, "id"): raise CogneeValidationError("user must have an id attribute.", log=False) @@ memory_top_k=memory_top_k,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/runtime.py` around lines 84 - 141, The validate_agent_memory_config function does not validate memory_top_k so invalid types or non-positive values can propagate; add a check in validate_agent_memory_config to ensure memory_top_k is an int and >= 1 (or your minimum allowed), and raise CogneeValidationError with a clear message (log=False) if it fails; then pass the validated/normalized memory_top_k into the AgentMemoryConfig return so downstream code always receives a safe integer value.
🧹 Nitpick comments (1)
cognee/modules/agent_memory/runtime.py (1)
187-191: Add a type annotation forfuncinbuild_method_params.
funcis currently untyped (Line 187). Typing it asCallable[..., Any]keeps this public runtime helper consistent with project typing standards.As per coding guidelines: "Add type hints to Python functions" and "Public Python APIs should be type-annotated where practical".Proposed fix
-from typing import Any, Optional +from typing import Any, Callable, Optional @@ -def build_method_params(func, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: +def build_method_params( + func: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any] +) -> dict[str, Any]:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/modules/agent_memory/runtime.py` around lines 187 - 191, The function build_method_params lacks a type annotation for its first parameter; update its signature to type the func parameter as Callable[..., Any] (import Callable and Any from typing) so the public runtime helper adheres to project typing standards and typing checks; ensure you update any relevant imports in the module if needed and run type checks to confirm no new issues.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@cognee/modules/agent_memory/runtime.py`:
- Around line 331-345: build_agent_trace is persisting raw method_params and
method_return_value (AgentTrace.method_params and method_return_value) which
risks storing secrets/PII; modify build_agent_trace to sanitize/redact these
fields before constructing the AgentTrace: run both context.method_params and
context.method_return_value through a dedicated redaction routine (e.g.,
redact_sensitive_data or extend sanitize_value) that removes or masks secrets
and PII and then truncate if needed (use
truncate_text/MAX_SERIALIZED_VALUE_LENGTH) before assigning to
AgentTrace.method_params and AgentTrace.method_return_value so only
redacted/truncated payloads are persisted.
- Around line 84-141: The validate_agent_memory_config function does not
validate memory_top_k so invalid types or non-positive values can propagate; add
a check in validate_agent_memory_config to ensure memory_top_k is an int and >=
1 (or your minimum allowed), and raise CogneeValidationError with a clear
message (log=False) if it fails; then pass the validated/normalized memory_top_k
into the AgentMemoryConfig return so downstream code always receives a safe
integer value.
---
Nitpick comments:
In `@cognee/modules/agent_memory/runtime.py`:
- Around line 187-191: The function build_method_params lacks a type annotation
for its first parameter; update its signature to type the func parameter as
Callable[..., Any] (import Callable and Any from typing) so the public runtime
helper adheres to project typing standards and typing checks; ensure you update
any relevant imports in the module if needed and run type checks to confirm no
new issues.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ddc8b5dd-1f0e-49bf-bbc9-f0da713865bf
📒 Files selected for processing (3)
cognee/modules/agent_memory/runtime.pycognee/tests/integration/modules/agent_memory/test_agent_memory_integration.pycognee/tests/test_agent_memory_e2e.py
🚧 Files skipped from review as they are similar to previous changes (2)
- cognee/tests/integration/modules/agent_memory/test_agent_memory_integration.py
- cognee/tests/test_agent_memory_e2e.py
…-the-simple-agentic-poc-feature
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
cognee/tests/test_agent_memory_e2e.py (1)
70-82:⚠️ Potential issue | 🟠 MajorRestore global config state mutated by the fixture.
Lines 73-74 mutate global Cognee root directories, but teardown at Line 81 does not restore previous values. This can leak state to later tests in the same process.
Proposed fix
+from collections.abc import AsyncIterator + `@pytest_asyncio.fixture` -async def agent_memory_e2e_env(tmp_path): +async def agent_memory_e2e_env(tmp_path: Path) -> AsyncIterator[None]: """Create a clean backend-access-controlled environment for agent-memory e2e tests.""" root = Path(tmp_path) + previous_data_root = cognee.config.data_root_directory() + previous_system_root = cognee.config.system_root_directory() cognee.config.data_root_directory(str(root / "data")) cognee.config.system_root_directory(str(root / "system")) await _reset_engines_and_prune() await engine_setup() - yield - - await _reset_engines_and_prune() + try: + yield + finally: + await _reset_engines_and_prune() + cognee.config.data_root_directory(previous_data_root) + cognee.config.system_root_directory(previous_system_root)Based on learnings: "Applies to cognee/tests/**/test_*.py : Avoid external state in Python tests; rely on test fixtures and CI-provided environment variables".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cognee/tests/test_agent_memory_e2e.py` around lines 70 - 82, The fixture agent_memory_e2e_env mutates global Cognee config via cognee.config.data_root_directory(...) and cognee.config.system_root_directory(...) but never restores the previous values; capture the original values (e.g., old_data_root = cognee.config.data_root_directory() and old_system_root = cognee.config.system_root_directory()) before setting the test paths, then after the yield restore them (call cognee.config.data_root_directory(old_data_root) and cognee.config.system_root_directory(old_system_root)) so the global state is returned to its prior values; ensure these captures/restores wrap the existing await _reset_engines_and_prune()/engine_setup() logic in the same async generator fixture.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cognee/tests/test_agent_memory_e2e.py`:
- Around line 23-30: The cleanup block in the test is swallowing all exceptions;
instead, call get_vector_engine(), attempt disposal via
vector_engine.engine.dispose(close=True) but catch only expected exceptions
(e.g., AttributeError, RuntimeError, or the specific engine-dispose error type)
and log the exception details (or re-raise after logging) rather than using a
bare except: pass; ensure you reference get_vector_engine and
vector_engine.engine.dispose so the handler only suppresses known harmless
errors and surfaces or logs unexpected failures for debugging.
---
Duplicate comments:
In `@cognee/tests/test_agent_memory_e2e.py`:
- Around line 70-82: The fixture agent_memory_e2e_env mutates global Cognee
config via cognee.config.data_root_directory(...) and
cognee.config.system_root_directory(...) but never restores the previous values;
capture the original values (e.g., old_data_root =
cognee.config.data_root_directory() and old_system_root =
cognee.config.system_root_directory()) before setting the test paths, then after
the yield restore them (call cognee.config.data_root_directory(old_data_root)
and cognee.config.system_root_directory(old_system_root)) so the global state is
returned to its prior values; ensure these captures/restores wrap the existing
await _reset_engines_and_prune()/engine_setup() logic in the same async
generator fixture.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 4ab5cb29-3906-4b65-a2fd-80e723b7b9da
📒 Files selected for processing (5)
cognee/modules/agent_memory/decorator.pycognee/modules/agent_memory/runtime.pycognee/tests/test_agent_memory_e2e.pycognee/tests/unit/modules/agent_memory/test_agent_memory.pyexamples/guides/agent_memory_quickstart.py
✅ Files skipped from review due to trivial changes (4)
- cognee/modules/agent_memory/decorator.py
- examples/guides/agent_memory_quickstart.py
- cognee/tests/unit/modules/agent_memory/test_agent_memory.py
- cognee/modules/agent_memory/runtime.py
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
1 participant
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.
Description
Adds
cognee.agent_memoryas a minimal productionized agent-memory decorator for async agent entrypoints.This PR promotes the simple agentic decorator POC into a supported Cognee feature with a narrow first version:
LLMGateway.acreate_structured_output(...)AgentTracepersistence after executiondataset_namescope handling with combined read/write permission checksRelated issue:
Acceptance Criteria
cognee.agent_memorydecorator exists on the public Cognee surfaceLLMGatewaycalls can consume memory context without custom agent logicType of Change
Screenshots
Pre-submission Checklist
CONTRIBUTING.md)DCO Affirmation
I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin.
Summary by CodeRabbit