refactor(memory): introduce a pluggable MemoryBackend trait#1
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ 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 |
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/memory/maintenance.rs (1)
93-105: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftAvoid decaying only the first 1,000 memories per type.
This says it decays all non-identity memories, but
get_by_type(..., 1000)returns only a capped page; larger stores can permanently skip lower-ranked memories. Extend the backend with pagination/streaming or a backend-side decay operation so maintenance covers the full set.🤖 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/memory/maintenance.rs` around lines 93 - 105, The maintenance loop in `maintain_memories` only processes a capped page because `backend.get_by_type(mem_type, 1000)` limits each `MemoryType`, so larger stores will miss later memories. Update the decay path to paginate or stream through all results for each non-identity type, or add a backend-side bulk decay operation, and keep the existing `maintenance_cancelable_op` flow while replacing the fixed-size fetch in `get_by_type` usage.
🤖 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/api/agents.rs`:
- Around line 834-835: The backend store created in the agent setup is losing
the agent identifier, so runtime-created agents can appear as unknown in
backend-scoped context. Update the `sqlite_backend_arc` setup in
`src/api/agents.rs` to preserve the ID by chaining `with_agent_id` on the
`MemoryStore` created by `MemoryStore::new` before passing it into
`crate::memory::sqlite_backend_arc`. Keep the fix localized to the agent
creation path so `SqliteBackend::agent_id()` forwards the correct value.
In `@src/memory/backend.rs`:
- Around line 164-179: The save flow in MemoryBackend::save writes the
structured memory before storing the embedding, so a failure in embeddings.store
can leave a partially persisted record. Update save(...) to own the compensation
logic by rolling back the earlier self.store.save(memory) when the embedding
write fails, using the existing MemoryBackend and embeddings symbols to locate
the change. Keep the FTS ensure logic after a successful embedding write, and
make sure any rollback path preserves the current error return semantics.
- Around line 347-358: The helper uses two different sources for the agent
identity: `sqlite_backend_arc` receives `agent_id`, but
`SqliteBackend::agent_id()` later reads `store.agent_id()`, which can leave
backend labels as "unknown". Update `sqlite_backend_arc` and the
`SqliteBackend::new`/`agent_id()` path so the store and backend share one source
of truth, either by deriving the backend agent ID from the store or by ensuring
the store is initialized with the passed `agent_id`. Make the change in the
`sqlite_backend_arc`/`SqliteBackend` flow so downstream code calling
`backend.agent_id()` always sees the real agent ID.
In `@src/memory/search.rs`:
- Around line 286-306: The association traversal in `search.rs` is relying on
the incidental order returned by `get_associations_for`, which makes downstream
graph/RRF ranking unstable. Make the ordering explicit before the `by_node`
grouping in the `frontier` processing path, either by sorting `all_edges`
deterministically or by tightening the backend/query contract in
`get_associations_for` so it always returns a defined order. Apply the same fix
wherever this traversal logic is duplicated, including the referenced mirrored
block, so `search` and its ranking stay stable across backends.
In `@src/memory/store.rs`:
- Around line 497-548: The batched IN queries in load_many and the association
lookup path bind the entire ID slice at once, which can exceed SQLite’s
parameter limit during large BFS traversals. Update these methods to split the
incoming ids into bounded chunks, run the existing query logic per chunk, and
merge the results before returning so large graph walks remain reliable without
hitting too many SQL variables.
In `@src/tools/memory_save.rs`:
- Around line 256-258: The FTS retry logic is missing from the `set_embedding`
save flow, so memories written through `MemorySave::save` with `embedding: None`
can stay out of FTS after a startup indexing failure. Update the backend-level
embedding persistence path (the `SqliteBackend::set_embedding`/related save
flow) to reuse the same retry behavior currently implemented in
`SqliteBackend::save(..., Some(_))`, and ensure the later embedding update also
triggers the FTS-index retry when needed.
---
Outside diff comments:
In `@src/memory/maintenance.rs`:
- Around line 93-105: The maintenance loop in `maintain_memories` only processes
a capped page because `backend.get_by_type(mem_type, 1000)` limits each
`MemoryType`, so larger stores will miss later memories. Update the decay path
to paginate or stream through all results for each non-identity type, or add a
backend-side bulk decay operation, and keep the existing
`maintenance_cancelable_op` flow while replacing the fixed-size fetch in
`get_by_type` usage.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 95095290-98fe-4665-aca1-6660730707e0
📒 Files selected for processing (15)
src/agent/cortex.rssrc/agent/maintenance.rssrc/api/agents.rssrc/api/memories.rssrc/conversation/context.rssrc/main.rssrc/memory.rssrc/memory/backend.rssrc/memory/lance.rssrc/memory/maintenance.rssrc/memory/search.rssrc/memory/store.rssrc/tools/memory_delete.rssrc/tools/memory_recall.rssrc/tools/memory_save.rs
Put the SQLite + LanceDB memory stack behind a single MemoryBackend trait so the memory layer is pluggable and testable. SqliteBackend is the production implementation; hybrid search and maintenance (decay/prune/merge) now operate over dyn MemoryBackend instead of the concrete store. Behaviour-preserving: the default build is unchanged and all existing tests pass.
aa675e9 to
c757231
Compare
|
Thanks for the reviews. Addressed in the latest push (folded into the single refactor commit):
On the |
Cross-backend fixes raised by the upstream PR reviews, applied to the integration branch: ingest delete-before-clear-progress, merged_memory_content size cap on the empty-winner path, ingest delete path-component guard, agent_id on the SQLite backend store, deterministic (weight-descending) edge ordering in graph traversal, chunked IN-list queries to stay under SQLite's bind limit, and FTS-index ensure on set_embedding. Feature-off: 891 tests green. Feature-on: clippy clean.
Summary
Decouple the long-term memory layer behind a single
MemoryBackendtrait so the storage backend is pluggable and the rest of the app is storage-agnostic. The existing SQLite + LanceDB stack becomesSqliteBackend(the production implementation); hybrid search and maintenance (decay / prune / merge) now operate overdyn MemoryBackendinstead of the concrete store.Behaviour-preserving: the default build is unchanged, no new dependencies are added, and all existing tests pass.
What changes
MemoryBackendtrait (src/memory/backend.rs) — a complete memory interface: structured storage (save/load/update/forget/record_access), vector + full-text search, the association graph (create / query / traverse), and maintenance ops (prune/merge).SqliteBackendimplements the trait over the currentMemoryStore(SQLite) +EmbeddingTable(LanceDB) — the production default.MemorySearch(hybrid search) and the maintenance routines are now generic overArc<dyn MemoryBackend>. Callers (memory tools, cortex, API) go through the trait instead of touching SQLite/LanceDB directly.Why
The memory layer was hard-wired to SQLite + LanceDB across search, maintenance, the tools and the cortex. Putting it behind a trait makes the storage backend swappable and the memory subsystem independently testable, without committing the project to any particular store.
For context: we maintain a SurrealDB-based unified memory backend (structured + vector + graph in a single embedded store) implemented against this trait in a fork. It could be submitted as a dedicated follow-up PR if there's interest — but this PR is intentionally just the seam: no alternative backend is added here.
Notes
cargo test --lib— all green.Summary by CodeRabbit
Bug Fixes
New Features