Add AgentMemory preparation - #523
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:
WalkthroughThe change introduces a preparing phase to Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
🛠️ Fix failing CI checks 💡
🧪 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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/draive/agents/state.py (1)
136-148: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winInserting
preparingbeforemetasilently rebinds positionalmetaarguments.Any existing
AgentMemory(recalling, remembering, meta)call now passesmetaaspreparing, producing a memory whose prepare hook is aMetaobject (failing only later, at call time). Make the optional params keyword-only.♻️ Proposed fix
def __init__( self, recalling: AgentMemoryRecalling, remembering: AgentMemoryRemembering, + *, preparing: AgentMemoryPreparing | None = None, meta: Meta = Meta.empty, ) -> None:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/agents/state.py` around lines 136 - 148, Update AgentMemory.__init__ so the optional preparing and meta parameters are keyword-only, preventing existing positional AgentMemory(recalling, remembering, meta) calls from rebinding meta as preparing; preserve the current defaults and initialization behavior in the super().__init__ call.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/draive/agents/agent.py`:
- Around line 213-217: Update the thread lookup before memory.prepare in the
respond flow to use the same sentinel-guarded lookup as the memory step helpers,
rather than ctx.state(AgentThread). Preserve the existing AgentThread typing and
pass the guarded thread to prepare so missing bound state cannot create a
throwaway thread.
In `@src/draive/agents/state.py`:
- Around line 106-111: Update the recall method to refresh the recalled thread’s
LRU recency by moving its entry to the end of the memory ordering before
returning the context. Preserve the existing initial_context fallback and
context combination behavior.
---
Outside diff comments:
In `@src/draive/agents/state.py`:
- Around line 136-148: Update AgentMemory.__init__ so the optional preparing and
meta parameters are keyword-only, preventing existing positional
AgentMemory(recalling, remembering, meta) calls from rebinding meta as
preparing; preserve the current defaults and initialization behavior in the
super().__init__ call.
🪄 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: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: eec26769-93e1-4dbe-a696-ef6f3fc3a83c
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (12)
docs/guides/Agents.mdpyproject.tomlsrc/draive/agents/agent.pysrc/draive/agents/state.pysrc/draive/agents/types.pysrc/draive/anthropic/config.pysrc/draive/anthropic/messages.pysrc/draive/models/__init__.pysrc/draive/ollama/config.pysrc/draive/postgres/agent_memory.pytests/test_agent_memory.pytests/test_postgres_agent_memory.py
💤 Files with no reviewable changes (1)
- src/draive/models/init.py
8e97da3 to
d1fb698
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/draive/agents/agent.py`:
- Around line 213-216: Update the missing-thread check in the surrounding agent
flow to import and raise AgentException from draive.agents.types instead of
ValueError, preserving the existing error message and AgentThread validation
behavior.
🪄 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: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 01b0a13d-6912-4627-a8b1-12a69829ca8e
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (12)
docs/guides/Agents.mdpyproject.tomlsrc/draive/agents/agent.pysrc/draive/agents/state.pysrc/draive/agents/types.pysrc/draive/anthropic/config.pysrc/draive/anthropic/messages.pysrc/draive/models/__init__.pysrc/draive/ollama/config.pysrc/draive/postgres/agent_memory.pytests/test_agent_memory.pytests/test_postgres_agent_memory.py
💤 Files with no reviewable changes (1)
- src/draive/models/init.py
d1fb698 to
843837d
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/draive/agents/state.py (1)
142-154: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winInserting
preparingbeforemetasilently changes the positional signature.Any existing
AgentMemory(recalling, remembering, some_meta)call now bindssome_metatopreparing, which is typed as a callable and will only fail later at call time. Making the new (and ideally all) parameters keyword-only prevents that.♻️ Proposed change
def __init__( self, recalling: AgentMemoryRecalling, remembering: AgentMemoryRemembering, + *, preparing: AgentMemoryPreparing | None = None, meta: Meta = Meta.empty, ) -> None:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/agents/state.py` around lines 142 - 154, Make the AgentMemory constructor parameters in __init__ keyword-only, including recalling, remembering, preparing, and meta, so existing positional calls cannot bind values to the wrong parameter. Preserve the current defaults and initialization behavior when arguments are supplied by keyword.src/draive/postgres/agent_memory.py (1)
24-35: 🚀 Performance & Scalability | 🔵 TrivialConsider a retention/pruning path for unbounded snapshot growth.
Every turn inserts a full-context snapshot and nothing is ever deleted, so table and index size grow with turns × context size per thread. A scheduled retention job (keep N latest snapshots per
(agent_uri, thread_id), or age-based deletion) plus monitoring on table bloat would keep recall latency and storage predictable in production.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/postgres/agent_memory.py` around lines 24 - 35, The immutable snapshot persistence flow needs a retention mechanism to prevent unbounded history growth. Add a scheduled pruning path for snapshots grouped by agent_uri and thread_id, retaining either the configured number of latest snapshots or those within the configured age window; preserve the latest-snapshot recall behavior and add monitoring for table or index bloat.src/draive/agents/agent.py (1)
184-188: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocstrings still describe the pre-rewrite
generativeimplementation.generativenow runs its own completion/tool loop and callsmemory.prepare/recall/rememberdirectly, so both references to the old composition are inaccurate.
src/draive/agents/agent.py#L184-L188: replace theStep.looping_completion(...)backing with the explicit completion-and-tools loop wrapped as a singleStep.src/draive/agents/agent.py#L478-L484: drop the "asgenerativeandfrom_skilldo" clause, since they do not composeprepare_step/recall_step/remember_step.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/agents/agent.py` around lines 184 - 188, The docstrings in src/draive/agents/agent.py at lines 184-188 and 478-484 are outdated. Update the generative documentation to describe its explicit completion-and-tools loop wrapped in a single Step instead of Step.looping_completion(...); also remove the “as generative and from_skill do” clause from the prepare_step/recall_step/remember_step documentation, with no direct code change required beyond these docstring updates.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/guides/Agents.md`:
- Around line 193-198: Revise the context-storage guarantee in the documentation
near the AgentThread memory description to apply only to the built-in
AgentMemory.volatile and PostgresAgentMemory backends. Avoid implying that
arbitrary recalling and remembering callables passed to AgentMemory provide
latest-snapshot behavior, unless explicitly documenting that behavior as a
required custom-backend contract.
In `@docs/guides/Postgres.md`:
- Around line 251-253: Update the persistence description in the Postgres guide
to scope “write-only” specifically to the remember path, while preserving the
statements about immutable snapshots, recall reading the latest snapshot, and
unbounded snapshot history.
---
Outside diff comments:
In `@src/draive/agents/agent.py`:
- Around line 184-188: The docstrings in src/draive/agents/agent.py at lines
184-188 and 478-484 are outdated. Update the generative documentation to
describe its explicit completion-and-tools loop wrapped in a single Step instead
of Step.looping_completion(...); also remove the “as generative and from_skill
do” clause from the prepare_step/recall_step/remember_step documentation, with
no direct code change required beyond these docstring updates.
In `@src/draive/agents/state.py`:
- Around line 142-154: Make the AgentMemory constructor parameters in __init__
keyword-only, including recalling, remembering, preparing, and meta, so existing
positional calls cannot bind values to the wrong parameter. Preserve the current
defaults and initialization behavior when arguments are supplied by keyword.
In `@src/draive/postgres/agent_memory.py`:
- Around line 24-35: The immutable snapshot persistence flow needs a retention
mechanism to prevent unbounded history growth. Add a scheduled pruning path for
snapshots grouped by agent_uri and thread_id, retaining either the configured
number of latest snapshots or those within the configured age window; preserve
the latest-snapshot recall behavior and add monitoring for table or index bloat.
🪄 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: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: d637528f-0a4f-44ac-84f9-8ac5c563d9b4
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (16)
docs/guides/Agents.mddocs/guides/Postgres.mdpyproject.tomlsrc/draive/__init__.pysrc/draive/agents/__init__.pysrc/draive/agents/agent.pysrc/draive/agents/state.pysrc/draive/agents/types.pysrc/draive/anthropic/config.pysrc/draive/anthropic/messages.pysrc/draive/models/__init__.pysrc/draive/ollama/config.pysrc/draive/postgres/agent_memory.pytests/test_agent_memory.pytests/test_agents.pytests/test_postgres_agent_memory.py
💤 Files with no reviewable changes (1)
- src/draive/models/init.py
843837d to
b0bb5f3
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/draive/agents/agent.py (1)
184-188: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winStale
Returnsdocstring.
generativeno longer delegates toStep.looping_completion(...); it runs an explicit streamed completion-and-tools loop.♻️ Proposed change
Returns ------- Self - Agent instance backed by ``Step.looping_completion(...)``. + Agent instance backed by an explicit streamed completion and tool + handling loop. """🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/agents/agent.py` around lines 184 - 188, Update the Returns documentation for generative to describe its explicit streamed completion-and-tools loop instead of claiming it delegates to Step.looping_completion(...), while preserving the documented Self return type and Agent instance behavior.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/draive/agents/agent.py`:
- Around line 478-484: Correct the Agent execution notes to state that
generative and from_skill invoke memory.prepare, recall, and remember directly
within their step bodies, rather than composing prepare_step, recall_step, and
remember_step. Apply the same wording correction in the Agents guide
documentation while preserving the statement that memory is not applied
implicitly.
---
Outside diff comments:
In `@src/draive/agents/agent.py`:
- Around line 184-188: Update the Returns documentation for generative to
describe its explicit streamed completion-and-tools loop instead of claiming it
delegates to Step.looping_completion(...), while preserving the documented Self
return type and Agent instance behavior.
🪄 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: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 9e5b35ae-288a-491e-b9a7-1efc81175135
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (16)
docs/guides/Agents.mddocs/guides/Postgres.mdpyproject.tomlsrc/draive/__init__.pysrc/draive/agents/__init__.pysrc/draive/agents/agent.pysrc/draive/agents/state.pysrc/draive/agents/types.pysrc/draive/anthropic/config.pysrc/draive/anthropic/messages.pysrc/draive/models/__init__.pysrc/draive/ollama/config.pysrc/draive/postgres/agent_memory.pytests/test_agent_memory.pytests/test_agents.pytests/test_postgres_agent_memory.py
💤 Files with no reviewable changes (1)
- src/draive/models/init.py
44f3bd4 to
719cb28
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/draive/agents/state.py (1)
138-150: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake
preparingkeyword-only.Adding
preparingbeforemetachanges the third positional argument frommetato the newpreparingcallable;AgentMemory(recall, remember, meta_value)would bind the wrong argument. Enforcing keyword-only avoids this silent misuse while keeping current construction sites unaffected.♻️ Proposed change
def __init__( self, recalling: AgentMemoryRecalling, remembering: AgentMemoryRemembering, + *, preparing: AgentMemoryPreparing | None = None, meta: Meta = Meta.empty, ) -> None:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/agents/state.py` around lines 138 - 150, Update the AgentMemory.__init__ signature so preparing is keyword-only, preserving meta as the third positional parameter and leaving existing keyword-based preparing construction unchanged.src/draive/postgres/agent_memory.py (1)
134-158: 🩺 Stability & Availability | 🔵 Trivial | ⚖️ Poor tradeoffConsider wrapping Postgres errors into a typed domain exception at this boundary.
recall/rememberlet_recall/_rememberpropagate raw Postgres exceptions (also reflected in the genericRaises: Exceptiondocstring at Lines 130-132). The repo already has domain exception types (e.g.AgentException, used elsewhere in this PR's test suite for scope errors) — wrapping provider errors here with operation/thread context would aid debugging and match guideline expectations.As per coding guidelines, "Translate provider/SDK errors into typed domain exceptions" and "Wrap third-party exceptions at boundaries and include actionable context (
provider,operation, identifiers) while redacting sensitive payloads."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/postgres/agent_memory.py` around lines 134 - 158, The recall and remember boundary methods currently expose raw Postgres exceptions. Wrap provider failures from _recall and _remember in the repository’s established typed domain exception, such as AgentException, including the provider, operation, and thread identifier in the context while excluding sensitive payloads; update the related Raises documentation to describe the typed exception.Source: Coding guidelines
src/draive/agents/agent.py (1)
184-188: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winStale
Returnsdocstring references the removedStep.looping_completion(...)pipeline.The implementation at lines 200-336 replaces
Step.looping_completion(...)with a manual streamed completion/tool loop, but theReturnssection still describes the old backing mechanism.✏️ Proposed fix
Returns ------- Self - Agent instance backed by ``Step.looping_completion(...)``. + Agent instance backed by an explicit streamed completion and + tool-handling loop.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/agents/agent.py` around lines 184 - 188, Update the Returns section of the Agent method docstring to describe the current manual streamed completion/tool loop implementation instead of referencing the removed Step.looping_completion(...) pipeline. Keep the documented return type and Agent instance behavior unchanged.
♻️ Duplicate comments (1)
src/draive/agents/agent.py (1)
213-214: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winStill raises
ValueError, notAgentException— past fix appears reverted.This is the exact code from a prior review round flagged as needing
AgentException(marked "✅ Addressed in commit 843837d"), yet the current diff still raises a bareValueError.AgentMemory's step helpers raiseAgentExceptionfor this same condition, anddocs/guides/Agents.mddocumentsAgentExceptionas the failure mode when noAgentThreadis bound — usingValueErrorhere means the identical failure surfaces with two different exception types depending on entry point.As per coding guidelines, "Translate provider/SDK errors into typed domain exceptions" and "Don't raise bare
Exception; preserve meaningful context."🐛 Proposed fix
if not ctx.contains_state(AgentThread): - raise ValueError("AgentThread not specified") + raise AgentException( + "AgentThread is not available in the current context -" + f" agent {identity.name} requires an active agent thread bound in scope" + )Requires importing
AgentExceptionfromdraive.agents.types.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/draive/agents/agent.py` around lines 213 - 214, In the AgentThread validation near contains_state, replace the bare ValueError with AgentException and import AgentException from draive.agents.types, preserving the existing “AgentThread not specified” context.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/draive/agents/state.py`:
- Around line 69-75: The documented LRU contract and recall behavior are
inconsistent. Update the recall implementation to refresh the accessed entry’s
recency via the existing memory ordering mechanism, preserving eviction behavior
when threads_limit is set; use the symbols governing recall and memory access
rather than weakening the docstring claim.
- Around line 472-475: Update the AgentException message in the AgentThread
context check to replace “memory require” with “memory operations require,”
leaving the rest of the error message unchanged.
In `@src/draive/postgres/agent_memory.py`:
- Around line 106-109: Rename the static factory method in PostgresAgentMemory
from instance to prepare so it matches the documented and tested API. Affected
sites: src/draive/postgres/agent_memory.py lines 106-109 require the definition
change; src/draive/postgres/agent_memory.py line 53 and
tests/test_postgres_agent_memory.py lines 64, 90, 121, and 173 already use
prepare and require no direct changes.
---
Outside diff comments:
In `@src/draive/agents/agent.py`:
- Around line 184-188: Update the Returns section of the Agent method docstring
to describe the current manual streamed completion/tool loop implementation
instead of referencing the removed Step.looping_completion(...) pipeline. Keep
the documented return type and Agent instance behavior unchanged.
In `@src/draive/agents/state.py`:
- Around line 138-150: Update the AgentMemory.__init__ signature so preparing is
keyword-only, preserving meta as the third positional parameter and leaving
existing keyword-based preparing construction unchanged.
In `@src/draive/postgres/agent_memory.py`:
- Around line 134-158: The recall and remember boundary methods currently expose
raw Postgres exceptions. Wrap provider failures from _recall and _remember in
the repository’s established typed domain exception, such as AgentException,
including the provider, operation, and thread identifier in the context while
excluding sensitive payloads; update the related Raises documentation to
describe the typed exception.
---
Duplicate comments:
In `@src/draive/agents/agent.py`:
- Around line 213-214: In the AgentThread validation near contains_state,
replace the bare ValueError with AgentException and import AgentException from
draive.agents.types, preserving the existing “AgentThread not specified”
context.
🪄 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: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: c1f43752-78b0-4f6a-a5e8-ad3c882e3551
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (15)
docs/guides/Agents.mddocs/guides/Postgres.mdpyproject.tomlsrc/draive/agents/__init__.pysrc/draive/agents/agent.pysrc/draive/agents/state.pysrc/draive/agents/types.pysrc/draive/anthropic/config.pysrc/draive/anthropic/messages.pysrc/draive/models/__init__.pysrc/draive/ollama/config.pysrc/draive/postgres/agent_memory.pytests/test_agent_memory.pytests/test_agents.pytests/test_postgres_agent_memory.py
💤 Files with no reviewable changes (1)
- src/draive/models/init.py
719cb28 to
91c1913
Compare
1 participant
No description provided.