diff --git a/src/agent/cortex.rs b/src/agent/cortex.rs index c560a759a..faa78f7a3 100644 --- a/src/agent/cortex.rs +++ b/src/agent/cortex.rs @@ -2468,9 +2468,8 @@ async fn run_cortex_loop( tokio::sync::watch::channel(false); maintenance_task = Some(tokio::spawn(async move { memory_maintenance::run_maintenance_with_cancel( - memory_search.store(), - memory_search.embedding_table(), - memory_search.embedding_model_arc(), + memory_search.backend().clone(), + memory_search.embedding_model_arc().clone(), &maintenance_config, maintenance_cancel_rx, ) @@ -4530,8 +4529,7 @@ async fn run_association_pass( let max_per_pass = cortex_config.association_max_per_pass; let is_backfill = since.is_none(); - let store = deps.memory_search.store(); - let embedding_table = deps.memory_search.embedding_table(); + let store = deps.memory_search.backend(); // Get the memories to process let memories = match fetch_memories_for_association(&deps.sqlite_pool, since).await { @@ -4555,7 +4553,7 @@ async fn run_association_pass( } // Find similar memories via embedding search - let similar = match embedding_table + let similar = match store .find_similar(memory_id, similarity_threshold, 10) .await { diff --git a/src/agent/maintenance.rs b/src/agent/maintenance.rs index 871533bd0..be17f4add 100644 --- a/src/agent/maintenance.rs +++ b/src/agent/maintenance.rs @@ -87,9 +87,8 @@ async fn run_maintenance_for_agent(deps: &AgentDeps) -> anyhow::Result<()> { }; let memory_search = &deps.memory_search; let report = crate::memory::maintenance::run_maintenance( - memory_search.store(), - memory_search.embedding_table(), - memory_search.embedding_model_arc(), + memory_search.backend().clone(), + memory_search.embedding_model_arc().clone(), &config, ) .await diff --git a/src/api/agents.rs b/src/api/agents.rs index 3ccc530df..22c4cfe79 100644 --- a/src/api/agents.rs +++ b/src/api/agents.rs @@ -830,23 +830,14 @@ pub async fn create_agent_internal( .clone() }; - let memory_store = crate::memory::MemoryStore::new(db.sqlite.clone()); - let embedding_table = crate::memory::EmbeddingTable::open_or_create(&db.lance) - .await - .map_err(|error| { - tracing::error!(%error, agent_id = %agent_id, "failed to init embeddings"); - format!("failed to init embeddings: {error}") - })?; - - if let Err(error) = embedding_table.ensure_fts_index().await { - tracing::warn!(%error, agent_id = %agent_id, "failed to create FTS index"); - } - - let memory_search = std::sync::Arc::new(crate::memory::MemorySearch::new( - memory_store, - embedding_table, - embedding_model, - )); + let backend: std::sync::Arc = { + let memory_store = crate::memory::MemoryStore::with_agent_id(db.sqlite.clone(), &agent_id); + crate::memory::sqlite_backend_arc(memory_store, &db.lance, &agent_id) + .await + .map_err(|e| format!("failed to init memory backend: {e}"))? + }; + let memory_search = + std::sync::Arc::new(crate::memory::MemorySearch::new(backend, embedding_model)); let task_store = state .task_store .load() diff --git a/src/api/memories.rs b/src/api/memories.rs index e8441b055..d29ec3f19 100644 --- a/src/api/memories.rs +++ b/src/api/memories.rs @@ -146,7 +146,7 @@ pub(super) async fn list_memories( ) -> Result, StatusCode> { let searches = state.memory_searches.load(); let memory_search = searches.get(&query.agent_id).ok_or(StatusCode::NOT_FOUND)?; - let store = memory_search.store(); + let store = memory_search.backend(); let limit = query.limit.min(200); let sort = parse_sort(&query.sort); @@ -232,7 +232,7 @@ pub(super) async fn memory_graph( ) -> Result, StatusCode> { let searches = state.memory_searches.load(); let memory_search = searches.get(&query.agent_id).ok_or(StatusCode::NOT_FOUND)?; - let store = memory_search.store(); + let store = memory_search.backend(); let limit = query.limit.min(500); let sort = parse_sort(&query.sort); @@ -290,7 +290,7 @@ pub(super) async fn memory_graph_neighbors( ) -> Result, StatusCode> { let searches = state.memory_searches.load(); let memory_search = searches.get(&query.agent_id).ok_or(StatusCode::NOT_FOUND)?; - let store = memory_search.store(); + let store = memory_search.backend(); let depth = query.depth.min(3); let exclude_ids: Vec = query diff --git a/src/conversation/context.rs b/src/conversation/context.rs index cdc412a9f..e7a1f2371 100644 --- a/src/conversation/context.rs +++ b/src/conversation/context.rs @@ -1,8 +1,6 @@ //! Context assembly: prompt + identity + memories + status. -use crate::agent::status::StatusBlock; use crate::error::Result; -use crate::memory::MemoryStore; /// Assembled context ready for injection into LLM. #[derive(Debug, Clone)] @@ -13,61 +11,10 @@ pub struct AssembledContext { pub conversation_history: String, } -/// Build context for a channel. -pub async fn build_channel_context( - base_prompt: &str, - memory_store: &MemoryStore, - status_block: &StatusBlock, - _conversation_id: &str, -) -> Result { - let mut context = String::new(); - - // Base channel prompt - context.push_str(base_prompt); - context.push_str("\n\n"); - - // Add status block - let status = status_block.render(); - if !status.is_empty() { - context.push_str("## Current Status\n\n"); - context.push_str(&status); - context.push('\n'); - } - - // Add identity memories (always included) - let identity_memories = memory_store - .get_by_type(crate::memory::types::MemoryType::Identity, 10) - .await?; - - if !identity_memories.is_empty() { - context.push_str("## Identity\n\n"); - for memory in identity_memories { - context.push_str(&format!("- {}\n", memory.content)); - } - context.push('\n'); - } - - // Add high-importance memories - let important_memories = memory_store.get_high_importance(0.8, 5).await?; - let non_identity: Vec<_> = important_memories - .into_iter() - .filter(|m| m.memory_type != crate::memory::types::MemoryType::Identity) - .collect(); - - if !non_identity.is_empty() { - context.push_str("## Key Context\n\n"); - for memory in non_identity { - context.push_str(&format!( - "- [{}] {}\n", - memory.memory_type, - memory.content.lines().next().unwrap_or(&memory.content) - )); - } - context.push('\n'); - } - - Ok(context) -} +// NOTE: `build_channel_context` was removed — it had no callers and took a +// concrete `&MemoryStore`, which would bypass the `MemoryBackend` abstraction +// if ever wired in (followups #13). Context building for channels should go +// through `Arc` instead. /// Build minimal context for a branch. pub async fn build_branch_context(base_prompt: &str) -> Result { diff --git a/src/main.rs b/src/main.rs index da784382e..ef204740d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2884,23 +2884,15 @@ async fn initialize_agents( }; // Per-agent memory system - let memory_store = - spacebot::memory::MemoryStore::with_agent_id(db.sqlite.clone(), &agent_config.id); let project_store = global_project_store.clone(); - let embedding_table = spacebot::memory::EmbeddingTable::open_or_create(&db.lance) - .await - .with_context(|| { - format!("failed to init embeddings for agent '{}'", agent_config.id) - })?; - - // Ensure FTS index exists for full-text search queries - if let Err(error) = embedding_table.ensure_fts_index().await { - tracing::warn!(%error, agent = %agent_config.id, "failed to create FTS index"); - } + let backend: Arc = { + let memory_store = + spacebot::memory::MemoryStore::with_agent_id(db.sqlite.clone(), &agent_config.id); + spacebot::memory::sqlite_backend_arc(memory_store, &db.lance, &agent_config.id).await? + }; let memory_search = Arc::new(spacebot::memory::MemorySearch::new( - memory_store, - embedding_table, + backend, embedding_model.clone(), )); diff --git a/src/memory.rs b/src/memory.rs index fb18163aa..40c8ac320 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -1,5 +1,6 @@ //! Memory storage and retrieval system. +pub mod backend; pub mod embedding; pub mod lance; pub mod maintenance; @@ -8,6 +9,7 @@ pub mod store; pub mod types; pub mod working; +pub use backend::{MemoryBackend, SqliteBackend, sqlite_backend_arc}; pub use embedding::EmbeddingModel; pub use lance::EmbeddingTable; pub use search::{MemorySearch, SearchConfig, SearchMode, SearchSort, curate_results}; diff --git a/src/memory/backend.rs b/src/memory/backend.rs new file mode 100644 index 000000000..84d39f9d6 --- /dev/null +++ b/src/memory/backend.rs @@ -0,0 +1,543 @@ +//! Pluggable memory storage backend trait and implementations. + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::error::Result; +use crate::memory::search::SearchSort; +use crate::memory::types::{Association, Memory, MemoryType}; + +/// Trait for pluggable memory storage backends. +/// +/// A backend encapsulates both structured memory storage (e.g. SQLite) and +/// vector/FTS search (e.g. LanceDB) behind a single async interface. +/// All provided methods delegate to the underlying stores. +#[async_trait] +pub trait MemoryBackend: Send + Sync + std::fmt::Debug { + /// The agent ID this backend is scoped to (empty string if none). + fn agent_id(&self) -> &str; + + // ── CRUD ───────────────────────────────────────────────────────────────── + + /// Persist a new memory and optionally its embedding. + /// + /// Not guaranteed atomic across the structured store and the embedding/vector + /// store: an implementation may commit the memory row and then fail to store + /// the embedding. Callers that need both-or-neither semantics must compensate + /// (e.g. delete the orphaned row) — see `tools::memory_save`. + async fn save(&self, memory: &Memory, embedding: Option<&[f32]>) -> Result<()>; + + /// Store or replace the embedding for an existing memory. + async fn set_embedding(&self, memory: &Memory, embedding: &[f32]) -> Result<()>; + + /// Hard-delete a memory and its embedding. + async fn delete(&self, id: &str) -> Result<()>; + + /// Load a memory by ID; returns `None` if not found. + async fn load(&self, id: &str) -> Result>; + + /// Update an existing memory's metadata/content. + async fn update(&self, memory: &Memory) -> Result<()>; + + /// Soft-delete: mark as forgotten. Returns `true` if the record was modified. + async fn forget(&self, id: &str) -> Result; + + /// Record an access, updating `last_accessed_at` and `access_count`. + async fn record_access(&self, id: &str) -> Result<()>; + + // ── Filtered retrieval ──────────────────────────────────────────────────── + + /// Get up to `limit` memories of the given type, ordered by importance. + async fn get_by_type(&self, t: MemoryType, limit: i64) -> Result>; + + /// Get up to `limit` memories with importance ≥ `threshold`. + async fn get_high_importance(&self, threshold: f32, limit: i64) -> Result>; + + /// Get up to `limit` memories sorted by `sort`, optionally filtered by type. + async fn get_sorted( + &self, + sort: SearchSort, + limit: i64, + t: Option, + ) -> Result>; + + // ── Associations ────────────────────────────────────────────────────────── + + /// Create (or upsert) an association between two memories. + async fn create_association(&self, a: &Association) -> Result<()>; + + /// Get all associations for a memory (incoming + outgoing). + async fn get_associations(&self, id: &str) -> Result>; + + /// Get associations where both endpoints are within `ids`. + async fn get_associations_between(&self, ids: &[String]) -> Result>; + + /// All associations incident to ANY of `ids` (either endpoint). Empty → empty. + async fn get_associations_for(&self, ids: &[String]) -> Result>; + + /// Batch-load memories by id (order unspecified; missing ids omitted; INCLUDES forgotten). + /// Empty → empty. + async fn load_many(&self, ids: &[String]) -> Result>; + + /// Delete all associations referencing `id`. Returns the number deleted. + async fn delete_associations_for_memory(&self, id: &str) -> Result; + + /// Traverse the graph up to `depth` hops from `id`, excluding `exclude`. + /// + /// Returns `(neighbor_memories, edges)`. + async fn get_neighbors( + &self, + id: &str, + depth: u32, + exclude: &[String], + ) -> Result<(Vec, Vec)>; + + // ── Search ──────────────────────────────────────────────────────────────── + + /// Approximate nearest-neighbour search. Returns `(memory_id, distance)`. + async fn vector_search(&self, q: &[f32], limit: usize) -> Result>; + + /// Full-text search. Returns `(memory_id, score)`. + async fn text_search(&self, q: &str, limit: usize) -> Result>; + + /// Find memories similar to the one at `id`. Returns `(memory_id, similarity)`. + async fn find_similar( + &self, + id: &str, + threshold: f32, + limit: usize, + ) -> Result>; + + // ── Maintenance ─────────────────────────────────────────────────────────── + + /// Delete all non-identity memories with importance below `threshold` that + /// were created before `older_than`. Returns the number of memories deleted. + async fn prune_below( + &self, + threshold: f32, + older_than: chrono::DateTime, + ) -> Result; + + /// Merge `loser_id` into `survivor_id`, updating the survivor's content and + /// optionally its embedding. The loser is soft-deleted (forgotten) and its + /// associations are rewired to the survivor. + async fn merge( + &self, + survivor_id: &str, + loser_id: &str, + new_content: &str, + new_embedding: Option<&[f32]>, + ) -> Result<()>; +} + +// ── SqliteBackend ───────────────────────────────────────────────────────────── + +/// `MemoryBackend` implementation backed by SQLite (via `MemoryStore`) and +/// LanceDB (via `EmbeddingTable`). This is the production default. +#[derive(Clone)] +pub struct SqliteBackend { + store: Arc, + embeddings: crate::memory::lance::EmbeddingTable, +} + +impl std::fmt::Debug for SqliteBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SqliteBackend") + .field("store", &self.store) + .field("embeddings", &"") + .finish() + } +} + +impl SqliteBackend { + /// Create a new `SqliteBackend` wrapping the given store and embedding table. + pub fn new( + store: Arc, + embeddings: crate::memory::lance::EmbeddingTable, + ) -> Self { + Self { store, embeddings } + } +} + +#[async_trait] +impl MemoryBackend for SqliteBackend { + fn agent_id(&self) -> &str { + self.store.agent_id() + } + + async fn save(&self, memory: &Memory, embedding: Option<&[f32]>) -> Result<()> { + self.store.save(memory).await?; + if let Some(emb) = embedding { + self.embeddings + .store(&memory.id, &memory.content, emb) + .await?; + // Ensure the FTS index exists once content has been written. The + // index can fail to create at startup on an empty table; this + // per-save retry (idempotent — no-ops once it exists) is how FTS + // starts working. Dropping it would leave text_search permanently + // falling back to vector+graph. A failure here is non-fatal. + if let Err(error) = self.embeddings.ensure_fts_index().await { + tracing::warn!(%error, "failed to ensure FTS index after memory save"); + } + } + Ok(()) + } + + async fn set_embedding(&self, memory: &Memory, embedding: &[f32]) -> Result<()> { + // EmbeddingTable::store is an append, not an upsert — delete any + // existing vector first so "set" replaces rather than duplicates. + self.embeddings.delete(&memory.id).await?; + self.embeddings + .store(&memory.id, &memory.content, embedding) + .await?; + // Ensure the FTS index exists once content has been written. Callers that + // save with `embedding: None` and set it here (e.g. the memory_save tool) + // would otherwise never trigger index creation, leaving text_search + // permanently falling back to vector+graph. Idempotent; non-fatal on error. + if let Err(error) = self.embeddings.ensure_fts_index().await { + tracing::warn!(%error, "failed to ensure FTS index after set_embedding"); + } + Ok(()) + } + + async fn delete(&self, id: &str) -> Result<()> { + self.store.delete(id).await?; // FK ON DELETE CASCADE drops edges + self.embeddings.delete(id).await?; // Lance row has no FK — delete explicitly + Ok(()) + } + + async fn load(&self, id: &str) -> Result> { + self.store.load(id).await + } + + async fn update(&self, m: &Memory) -> Result<()> { + self.store.update(m).await + } + + async fn forget(&self, id: &str) -> Result { + self.store.forget(id).await + } + + async fn record_access(&self, id: &str) -> Result<()> { + self.store.record_access(id).await + } + + async fn get_by_type(&self, t: MemoryType, limit: i64) -> Result> { + self.store.get_by_type(t, limit).await + } + + async fn get_high_importance(&self, th: f32, limit: i64) -> Result> { + self.store.get_high_importance(th, limit).await + } + + async fn get_sorted( + &self, + sort: SearchSort, + limit: i64, + t: Option, + ) -> Result> { + self.store.get_sorted(sort, limit, t).await + } + + async fn create_association(&self, a: &Association) -> Result<()> { + self.store.create_association(a).await + } + + async fn get_associations(&self, id: &str) -> Result> { + self.store.get_associations(id).await + } + + async fn get_associations_between(&self, ids: &[String]) -> Result> { + self.store.get_associations_between(ids).await + } + + async fn get_associations_for(&self, ids: &[String]) -> Result> { + self.store.get_associations_for(ids).await + } + + async fn load_many(&self, ids: &[String]) -> Result> { + self.store.load_many(ids).await + } + + async fn delete_associations_for_memory(&self, id: &str) -> Result { + self.store.delete_associations_for_memory(id).await + } + + async fn get_neighbors( + &self, + id: &str, + depth: u32, + exclude: &[String], + ) -> Result<(Vec, Vec)> { + self.store.get_neighbors(id, depth, exclude).await + } + + async fn vector_search(&self, q: &[f32], limit: usize) -> Result> { + self.embeddings.vector_search(q, limit).await + } + + async fn text_search(&self, q: &str, limit: usize) -> Result> { + self.embeddings.text_search(q, limit).await + } + + async fn find_similar(&self, id: &str, th: f32, limit: usize) -> Result> { + self.embeddings.find_similar(id, th, limit).await + } + + async fn prune_below( + &self, + threshold: f32, + older_than: chrono::DateTime, + ) -> Result { + // Mirror maintenance.rs:165-179 exactly: one unbounded SQL SELECT, then + // delete each (self.delete also drops the Lance embedding — the + // intentional orphan-cleanup behaviour change). + use sqlx::Row as _; + let rows = sqlx::query( + "SELECT id FROM memories WHERE importance < ? AND memory_type != 'identity' AND created_at < ?", + ) + .bind(threshold) + .bind(older_than) + .fetch_all(self.store.pool()) + .await + .map_err(|e| crate::error::DbError::Query(e.to_string()))?; + + let mut n = 0u64; + for row in rows { + let id: String = row + .try_get("id") + .map_err(|e| crate::error::DbError::Query(e.to_string()))?; + self.delete(&id).await?; + n += 1; + } + Ok(n) + } + + async fn merge( + &self, + survivor_id: &str, + loser_id: &str, + new_content: &str, + new_embedding: Option<&[f32]>, + ) -> Result<()> { + // merge_memories_atomic(updated_survivor, loser) internally forgets the + // loser and rewires its associations onto the survivor — DO NOT set + // forgotten here. We only build the updated survivor. + let mut survivor = self.store.load(survivor_id).await?.ok_or_else(|| { + crate::error::DbError::Query(format!("merge survivor {survivor_id} not found")) + })?; + let loser = self.store.load(loser_id).await?.ok_or_else(|| { + crate::error::DbError::Query(format!("merge loser {loser_id} not found")) + })?; + + survivor.content = new_content.to_string(); + survivor.updated_at = chrono::Utc::now(); + + self.store.merge_memories_atomic(&survivor, &loser).await?; + + // Embedding fix-up, replicating maintenance.rs::merge_pair exactly: + // EmbeddingTable::store is an append, so delete the survivor's old + // vector before re-storing, and drop the loser's vector. + if let Some(emb) = new_embedding { + self.embeddings.delete(survivor_id).await?; + self.embeddings.store(survivor_id, new_content, emb).await?; + } + self.embeddings.delete(loser_id).await?; + + Ok(()) + } +} + +// ── Construction helper ─────────────────────────────────────────────────────── + +/// Build a `SqliteBackend` wrapped in an `Arc`. +/// +/// Opens (or creates) the LanceDB embedding table, attempts to ensure the FTS +/// index (non-fatal on failure), then wraps the result. The caller is +/// responsible for constructing the `MemoryStore` with the appropriate scope +/// (e.g. `MemoryStore::with_agent_id` vs `MemoryStore::new`). +pub async fn sqlite_backend_arc( + store: Arc, + lance: &lancedb::Connection, + agent_id: &str, +) -> crate::error::Result> { + let embeddings = crate::memory::lance::EmbeddingTable::open_or_create(lance) + .await + .map_err(|e| crate::error::DbError::LanceConnect(format!("agent '{agent_id}': {e}")))?; + if let Err(error) = embeddings.ensure_fts_index().await { + tracing::warn!(%error, agent = %agent_id, "failed to ensure FTS index"); + } + Ok(Arc::new(SqliteBackend::new(store, embeddings))) +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory::lance::EmbeddingTable; + use crate::memory::store::MemoryStore; + use crate::memory::types::{Memory, MemoryType}; + + const DIM: usize = 384; + + // Returns the backend AND the TempDir guard — keep the guard alive for the + // test's duration (dropping it deletes the Lance directory). Mirrors the + // construction in src/memory/search.rs:577 and maintenance.rs:531. + async fn sqlite_backend() -> (SqliteBackend, tempfile::TempDir) { + let store = MemoryStore::connect_in_memory().await; + let dir = tempfile::tempdir().unwrap(); + let conn = lancedb::connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let embeddings = EmbeddingTable::open_or_create(&conn).await.unwrap(); + (SqliteBackend::new(store, embeddings), dir) + } + + #[tokio::test] + async fn save_load_delete_roundtrips_through_trait() { + let (be, _dir) = sqlite_backend().await; + let be: &dyn MemoryBackend = &be; + let m = Memory::new("the sky is blue", MemoryType::Fact); + let emb = vec![0.1_f32; DIM]; + + be.save(&m, Some(&emb)).await.unwrap(); + let loaded = be.load(&m.id).await.unwrap().expect("memory present"); + assert_eq!(loaded.content, "the sky is blue"); + + be.delete(&m.id).await.unwrap(); + assert!(be.load(&m.id).await.unwrap().is_none()); + // embedding gone from Lance too: vector search returns nothing for it + let hits = be.vector_search(&emb, 5).await.unwrap(); + assert!(hits.iter().all(|(id, _)| id != &m.id)); + } + + #[tokio::test] + async fn merge_moves_content_and_forgets_loser() { + let (be, _dir) = sqlite_backend().await; + let mut a = Memory::new("cats are mammals", MemoryType::Fact); + a.importance = 0.9; + let b = Memory::new("cats are animals", MemoryType::Fact); + be.save(&a, Some(&vec![0.1; DIM])).await.unwrap(); + be.save(&b, Some(&vec![0.2; DIM])).await.unwrap(); + + be.merge( + &a.id, + &b.id, + "cats are mammals\n\ncats are animals", + Some(&vec![0.15; DIM]), + ) + .await + .unwrap(); + + let survivor = be.load(&a.id).await.unwrap().unwrap(); + assert!(survivor.content.contains("mammals") && survivor.content.contains("animals")); + // merge_memories_atomic forgets the loser internally — the caller never + // sets it. Verify the function did so. + let loser = be.load(&b.id).await.unwrap().unwrap(); + assert!(loser.forgotten); + // The loser's embedding must be gone (no stale vector for it). + assert!( + be.vector_search(&vec![0.2; DIM], 5) + .await + .unwrap() + .iter() + .all(|(id, _)| id != &b.id) + ); + } + + #[tokio::test] + async fn prune_below_skips_identity_and_recent() { + let (be, _dir) = sqlite_backend().await; + let mut low = Memory::new("trivia", MemoryType::Fact); + low.importance = 0.1; + low.created_at = chrono::Utc::now() - chrono::Duration::days(30); + let mut ident = Memory::new("my name is X", MemoryType::Identity); + ident.importance = 0.1; + ident.created_at = low.created_at; + be.save(&low, None).await.unwrap(); + be.save(&ident, None).await.unwrap(); + + let cut = chrono::Utc::now() - chrono::Duration::days(7); + let n = be.prune_below(0.5, cut).await.unwrap(); + assert_eq!(n, 1); + assert!(be.load(&low.id).await.unwrap().is_none()); + assert!(be.load(&ident.id).await.unwrap().is_some()); // identity preserved + } + + #[tokio::test] + async fn get_associations_for_returns_incident_edges() { + let (be, _dir) = sqlite_backend().await; + use crate::memory::types::RelationType; + let a = Memory::new("a", MemoryType::Fact); + let b = Memory::new("b", MemoryType::Fact); + let c = Memory::new("c", MemoryType::Fact); + for m in [&a, &b, &c] { + be.save(m, None).await.unwrap(); + } + be.create_association(&Association::new(&a.id, &b.id, RelationType::RelatedTo)) + .await + .unwrap(); + be.create_association(&Association::new(&b.id, &c.id, RelationType::RelatedTo)) + .await + .unwrap(); + // incident to {a}: only a→b + let e = be + .get_associations_for(std::slice::from_ref(&a.id)) + .await + .unwrap(); + assert_eq!(e.len(), 1); + // incident to {a, c}: a→b (a is endpoint) and b→c (c is endpoint) + let e2 = be + .get_associations_for(&[a.id.clone(), c.id.clone()]) + .await + .unwrap(); + assert_eq!(e2.len(), 2); + assert!(be.get_associations_for(&[]).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn load_many_returns_present_memories() { + let (be, _dir) = sqlite_backend().await; + let a = Memory::new("x", MemoryType::Fact); + be.save(&a, None).await.unwrap(); + let got = be + .load_many(&[a.id.clone(), "missing".into()]) + .await + .unwrap(); + assert_eq!(got.len(), 1); + assert_eq!(got[0].id, a.id); + assert!(be.load_many(&[]).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn load_many_includes_forgotten() { + let (be, _dir) = sqlite_backend().await; + let a = Memory::new("forgotten fact", MemoryType::Fact); + be.save(&a, None).await.unwrap(); + be.forget(&a.id).await.unwrap(); + let got = be.load_many(std::slice::from_ref(&a.id)).await.unwrap(); + assert_eq!(got.len(), 1, "load_many must return forgotten rows"); + assert!( + got[0].forgotten, + "the returned row must be marked forgotten" + ); + } + + #[tokio::test] + async fn set_embedding_replaces_does_not_duplicate() { + let (be, _dir) = sqlite_backend().await; + let m = Memory::new("the sky is blue", MemoryType::Fact); + be.save(&m, Some(&vec![0.1_f32; DIM])).await.unwrap(); + + // Re-embed the same memory; store is an append, so set_embedding must + // delete the old vector first or KNN would surface two rows for one id. + be.set_embedding(&m, &vec![0.9_f32; DIM]).await.unwrap(); + + let hits = be.vector_search(&vec![0.9_f32; DIM], 10).await.unwrap(); + let occurrences = hits.iter().filter(|(id, _)| id == &m.id).count(); + assert_eq!(occurrences, 1, "set_embedding must replace, not duplicate"); + } +} diff --git a/src/memory/lance.rs b/src/memory/lance.rs index 382bedefc..9b17c2e0f 100644 --- a/src/memory/lance.rs +++ b/src/memory/lance.rs @@ -9,7 +9,7 @@ use std::sync::Arc; /// Schema constants for the embeddings table. const TABLE_NAME: &str = "memory_embeddings"; -const EMBEDDING_DIM: i32 = 384; // all-MiniLM-L6-v2 dimension +pub const EMBEDDING_DIM: i32 = 384; // all-MiniLM-L6-v2 dimension /// LanceDB table for memory embeddings with HNSW index and FTS. pub struct EmbeddingTable { diff --git a/src/memory/maintenance.rs b/src/memory/maintenance.rs index ffb0b39bb..4a080bd4f 100644 --- a/src/memory/maintenance.rs +++ b/src/memory/maintenance.rs @@ -1,11 +1,10 @@ //! Memory maintenance: decay, prune, merge, reindex. use crate::error::Result; -use crate::memory::{EmbeddingModel, EmbeddingTable, Memory, MemoryStore, MemoryType}; +use crate::memory::backend::MemoryBackend; +use crate::memory::{EmbeddingModel, Memory, MemoryType}; use anyhow::Context; -use sqlx::Row; -use sqlx::sqlite::SqliteRow; use tokio::sync::watch; use std::collections::HashSet; @@ -43,29 +42,20 @@ impl Default for MaintenanceConfig { /// Run maintenance tasks on the memory store. pub async fn run_maintenance( - memory_store: &MemoryStore, - embedding_table: &EmbeddingTable, - embedding_model: &Arc, + backend: Arc, + embedding_model: Arc, config: &MaintenanceConfig, ) -> Result { let (_maintenance_cancel_tx, maintenance_cancel_rx) = watch::channel(false); - run_maintenance_with_cancel( - memory_store, - embedding_table, - embedding_model, - config, - maintenance_cancel_rx, - ) - .await + run_maintenance_with_cancel(backend, embedding_model, config, maintenance_cancel_rx).await } /// Run maintenance tasks with a cancellation signal. /// /// The signal allows maintenance to exit quickly when the caller decides to stop it. pub async fn run_maintenance_with_cancel( - memory_store: &MemoryStore, - embedding_table: &EmbeddingTable, - embedding_model: &Arc, + backend: Arc, + embedding_model: Arc, config: &MaintenanceConfig, mut maintenance_cancel_rx: watch::Receiver, ) -> Result { @@ -78,12 +68,11 @@ pub async fn run_maintenance_with_cancel( #[allow(clippy::field_reassign_with_default)] { report.decayed = - apply_decay(memory_store, config.decay_rate, &mut maintenance_cancel_rx).await?; - report.pruned = prune_memories(memory_store, config, &mut maintenance_cancel_rx).await?; + apply_decay(&backend, config.decay_rate, &mut maintenance_cancel_rx).await?; + report.pruned = prune_memories(&backend, config, &mut maintenance_cancel_rx).await?; report.merged = merge_similar_memories( - memory_store, - embedding_table, - embedding_model, + &backend, + &embedding_model, config.merge_similarity_threshold, &mut maintenance_cancel_rx, ) @@ -95,7 +84,7 @@ pub async fn run_maintenance_with_cancel( /// Apply importance decay based on recency and access patterns. async fn apply_decay( - memory_store: &MemoryStore, + backend: &Arc, decay_rate: f32, maintenance_cancel_rx: &mut watch::Receiver, ) -> Result { @@ -111,11 +100,9 @@ async fn apply_decay( let mut decayed_count = 0; for mem_type in all_types { - let memories = maintenance_cancelable_op( - maintenance_cancel_rx, - memory_store.get_by_type(mem_type, 1000), - ) - .await?; + let memories = + maintenance_cancelable_op(maintenance_cancel_rx, backend.get_by_type(mem_type, 1000)) + .await?; for mut memory in memories { check_maintenance_cancellation(maintenance_cancel_rx).await?; @@ -139,8 +126,7 @@ async fn apply_decay( if (new_importance - memory.importance).abs() > 0.01 { memory.importance = new_importance.clamp(0.0, 1.0); memory.updated_at = now; - maintenance_cancelable_op(maintenance_cancel_rx, memory_store.update(&memory)) - .await?; + maintenance_cancelable_op(maintenance_cancel_rx, backend.update(&memory)).await?; decayed_count += 1; } } @@ -151,7 +137,7 @@ async fn apply_decay( /// Prune memories that have fallen below the importance threshold. async fn prune_memories( - memory_store: &MemoryStore, + backend: &Arc, config: &MaintenanceConfig, maintenance_cancel_rx: &mut watch::Receiver, ) -> Result { @@ -161,44 +147,23 @@ async fn prune_memories( let min_age = chrono::Duration::days(config.min_age_days); let cutoff_date = now - min_age; - // Get all memories below threshold that are old enough - let candidates: Vec = maintenance_cancelable_op( + let n = maintenance_cancelable_op( maintenance_cancel_rx, - sqlx::query( - r#" - SELECT id FROM memories - WHERE importance < ? - AND memory_type != 'identity' - AND created_at < ? - "#, - ) - .bind(config.prune_threshold) - .bind(cutoff_date) - .fetch_all(memory_store.pool()), + backend.prune_below(config.prune_threshold, cutoff_date), ) .await?; - let mut pruned_count = 0; - - for row in candidates { - let id: String = row.try_get("id")?; - check_maintenance_cancellation(maintenance_cancel_rx).await?; - maintenance_cancelable_op(maintenance_cancel_rx, memory_store.delete(&id)).await?; - pruned_count += 1; - } - - Ok(pruned_count) + Ok(n as usize) } /// Merge near-duplicate memories. async fn merge_similar_memories( - memory_store: &MemoryStore, - embedding_table: &EmbeddingTable, + backend: &Arc, embedding_model: &Arc, similarity_threshold: f32, maintenance_cancel_rx: &mut watch::Receiver, ) -> Result { - let memory_ids = fetch_candidate_memory_ids(memory_store, maintenance_cancel_rx).await?; + let memory_ids = fetch_candidate_memory_ids(backend, maintenance_cancel_rx).await?; if memory_ids.is_empty() { return Ok(0); } @@ -217,7 +182,7 @@ async fn merge_similar_memories( } let Some(source_memory) = - maintenance_cancelable_op(maintenance_cancel_rx, memory_store.load(&source_id)).await? + maintenance_cancelable_op(maintenance_cancel_rx, backend.load(&source_id)).await? else { continue; }; @@ -231,7 +196,7 @@ async fn merge_similar_memories( let similar = maintenance_cancelable_op( maintenance_cancel_rx, - embedding_table.find_similar( + backend.find_similar( &source_memory.id, similarity_threshold, MAX_MAINTENANCE_SIMILAR_CANDIDATES, @@ -262,7 +227,7 @@ async fn merge_similar_memories( } let Some(candidate_memory) = - maintenance_cancelable_op(maintenance_cancel_rx, memory_store.load(&candidate_id)) + maintenance_cancelable_op(maintenance_cancel_rx, backend.load(&candidate_id)) .await? else { continue; @@ -273,8 +238,7 @@ async fn merge_similar_memories( let (winner, loser) = choose_merge_pair(&active_survivor, &candidate_memory); let merged_survivor = merge_pair( - memory_store, - embedding_table, + backend, embedding_model, &winner, &loser, @@ -300,7 +264,7 @@ async fn merge_similar_memories( Ok(merged_count) } -fn choose_merge_pair(first: &Memory, second: &Memory) -> (Memory, Memory) { +pub fn choose_merge_pair(first: &Memory, second: &Memory) -> (Memory, Memory) { let first_wins = first.importance > second.importance || (first.importance == second.importance && first.id < second.id); @@ -311,7 +275,7 @@ fn choose_merge_pair(first: &Memory, second: &Memory) -> (Memory, Memory) { } } -fn merged_memory_content(winner: String, loser: &str) -> String { +pub fn merged_memory_content(winner: String, loser: &str) -> String { let winner_trimmed = winner.trim_end(); let loser_trimmed = loser.trim_end(); @@ -334,8 +298,7 @@ fn merged_memory_content(winner: String, loser: &str) -> String { } async fn merge_pair( - memory_store: &MemoryStore, - embedding_table: &EmbeddingTable, + backend: &Arc, embedding_model: &Arc, survivor: &Memory, merged: &Memory, @@ -343,62 +306,48 @@ async fn merge_pair( ) -> Result { check_maintenance_cancellation(maintenance_cancel_rx).await?; - let mut updated_survivor = survivor.clone(); - updated_survivor.content = merged_memory_content(updated_survivor.content, &merged.content); - updated_survivor.updated_at = chrono::Utc::now(); + let content = merged_memory_content(survivor.content.clone(), &merged.content); - maintenance_cancelable_op( - maintenance_cancel_rx, - memory_store.merge_memories_atomic(&updated_survivor, merged), - ) - .await?; + let embedding = + maintenance_cancelable_op(maintenance_cancel_rx, embedding_model.embed_one(&content)) + .await?; - let updated_survivor_embedding = maintenance_cancelable_op( - maintenance_cancel_rx, - embedding_model.embed_one(&updated_survivor.content), - ) - .await?; maintenance_cancelable_op( maintenance_cancel_rx, - embedding_table.delete(&updated_survivor.id), + backend.merge(&survivor.id, &merged.id, &content, Some(&embedding)), ) .await?; - maintenance_cancelable_op( - maintenance_cancel_rx, - embedding_table.store( - &updated_survivor.id, - &updated_survivor.content, - &updated_survivor_embedding, - ), - ) - .await?; - maintenance_cancelable_op(maintenance_cancel_rx, embedding_table.delete(&merged.id)).await?; + + // Return the updated survivor with the new content so the caller can + // chain additional merges against the right state. + let mut updated_survivor = survivor.clone(); + updated_survivor.content = content; + updated_survivor.updated_at = chrono::Utc::now(); Ok(updated_survivor) } async fn fetch_candidate_memory_ids( - memory_store: &MemoryStore, + backend: &Arc, maintenance_cancel_rx: &mut watch::Receiver, ) -> Result> { check_maintenance_cancellation(maintenance_cancel_rx).await?; - let rows: Vec = maintenance_cancelable_op( + use crate::memory::search::SearchSort; + let memories = maintenance_cancelable_op( maintenance_cancel_rx, - sqlx::query( - "SELECT id FROM memories WHERE forgotten = 0 ORDER BY importance DESC, created_at DESC, id ASC LIMIT ?", - ) - .bind(MAX_MAINTENANCE_MERGE_SOURCE_MEMORIES) - .fetch_all(memory_store.pool()), + backend.get_sorted( + SearchSort::Importance, + MAX_MAINTENANCE_MERGE_SOURCE_MEMORIES, + None, + ), ) .await .with_context(|| "failed to fetch candidate memories for maintenance")?; - let ids: Vec = rows + let ids: Vec = memories .into_iter() - .map(|row| { - let memory_id: String = row.get("id"); - memory_id - }) + .filter(|m| !m.forgotten) + .map(|m| m.id) .collect(); Ok(ids) @@ -483,7 +432,8 @@ pub struct MaintenanceReport { #[cfg(test)] mod tests { use super::*; - use crate::memory::{Association, RelationType}; + use crate::memory::backend::{MemoryBackend, SqliteBackend}; + use crate::memory::types::{Association, RelationType}; use std::sync::{Arc, OnceLock}; use tempfile::tempdir; use tokio::time::Duration; @@ -500,41 +450,41 @@ mod tests { })) } + async fn sqlite_backend() -> (Arc, tempfile::TempDir) { + let store = crate::memory::MemoryStore::connect_in_memory().await; + let dir = tempdir().unwrap(); + let conn = lancedb::connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let embeddings = crate::memory::EmbeddingTable::open_or_create(&conn) + .await + .unwrap(); + let backend: Arc = Arc::new(SqliteBackend::new(store, embeddings)); + (backend, dir) + } + async fn create_memory_with_embedding( - store: &MemoryStore, - embedding_table: &crate::memory::lance::EmbeddingTable, + backend: &Arc, content: &str, memory_type: MemoryType, importance: f32, embedding: Vec, - ) -> Memory { - let memory = Memory::new(content, memory_type).with_importance(importance); - store.save(&memory).await.expect("failed to save memory"); - - embedding_table - .store(&memory.id, &memory.content, &embedding) + ) -> crate::memory::Memory { + let memory = crate::memory::Memory::new(content, memory_type).with_importance(importance); + backend + .save(&memory, Some(&embedding)) .await - .expect("failed to store embedding"); - + .expect("failed to save memory"); memory } #[tokio::test] async fn merges_near_duplicate_memories_and_transfers_associations() { - let store = MemoryStore::connect_in_memory().await; - - let dir = tempdir().expect("failed to create temp dir"); - let lance_conn = lancedb::connect(dir.path().to_str().expect("temp path")) - .execute() - .await - .expect("failed to connect to lancedb"); - let embedding_table = crate::memory::EmbeddingTable::open_or_create(&lance_conn) - .await - .expect("failed to create embedding table"); + let (backend, _dir) = sqlite_backend().await; let survivor = create_memory_with_embedding( - &store, - &embedding_table, + &backend, "rust memory maintenance", MemoryType::Fact, 0.9, @@ -543,8 +493,7 @@ mod tests { .await; let duplicate = create_memory_with_embedding( - &store, - &embedding_table, + &backend, "rust memory maintenance updated", MemoryType::Fact, 0.4, @@ -553,8 +502,7 @@ mod tests { .await; let related = create_memory_with_embedding( - &store, - &embedding_table, + &backend, "related memory", MemoryType::Fact, 0.7, @@ -562,7 +510,7 @@ mod tests { ) .await; - store + backend .create_association(&Association::new( &duplicate.id, &related.id, @@ -571,7 +519,7 @@ mod tests { .await .expect("failed to create related-to association"); - store + backend .create_association(&Association::new( &related.id, &duplicate.id, @@ -588,13 +536,13 @@ mod tests { }; let embedding_model = shared_embedding_model(); - let report = run_maintenance(&store, &embedding_table, &embedding_model, &config) + let report = run_maintenance(Arc::clone(&backend), Arc::clone(&embedding_model), &config) .await .expect("maintenance should succeed"); assert_eq!(report.merged, 1); - let updated_survivor = store + let updated_survivor = backend .load(&survivor.id) .await .expect("failed to load survivor") @@ -604,20 +552,20 @@ mod tests { // The loser ("rust memory maintenance updated") is NOT appended. assert_eq!(updated_survivor.content, "rust memory maintenance"); - let forgotten_duplicate = store + let forgotten_duplicate = backend .load(&duplicate.id) .await .expect("failed to load duplicate") .expect("duplicate should still exist"); assert!(forgotten_duplicate.forgotten); - let duplicate_embeddings = embedding_table + let duplicate_embeddings = backend .find_similar(&duplicate.id, 0.0, 10) .await .expect("failed to search for missing duplicate embeddings"); assert!(duplicate_embeddings.is_empty()); - let survivor_associations = store + let survivor_associations = backend .get_associations(&survivor.id) .await .expect("failed to fetch survivor associations"); @@ -641,7 +589,7 @@ mod tests { .any(|assoc| assoc.source_id == related.id && assoc.target_id == survivor.id) ); - let duplicate_associations = store + let duplicate_associations = backend .get_associations(&duplicate.id) .await .expect("failed to load duplicate associations"); @@ -659,20 +607,10 @@ mod tests { #[tokio::test] async fn merges_multiple_duplicates_into_one_survivor_in_single_pass() { - let store = MemoryStore::connect_in_memory().await; - - let dir = tempdir().expect("failed to create temp dir"); - let lance_conn = lancedb::connect(dir.path().to_str().expect("temp path")) - .execute() - .await - .expect("failed to connect to lancedb"); - let embedding_table = crate::memory::EmbeddingTable::open_or_create(&lance_conn) - .await - .expect("failed to create embedding table"); + let (backend, _dir) = sqlite_backend().await; let survivor = create_memory_with_embedding( - &store, - &embedding_table, + &backend, "durable rust maintenance note", MemoryType::Fact, 0.9, @@ -681,8 +619,7 @@ mod tests { .await; let duplicate_a = create_memory_with_embedding( - &store, - &embedding_table, + &backend, "durable rust maintenance note update A", MemoryType::Fact, 0.6, @@ -690,8 +627,7 @@ mod tests { ) .await; let duplicate_b = create_memory_with_embedding( - &store, - &embedding_table, + &backend, "durable rust maintenance note update B", MemoryType::Fact, 0.5, @@ -699,34 +635,22 @@ mod tests { ) .await; - let related_a = create_memory_with_embedding( - &store, - &embedding_table, - "related A", - MemoryType::Fact, - 0.7, - { + let related_a = + create_memory_with_embedding(&backend, "related A", MemoryType::Fact, 0.7, { let mut embedding = vec![0.0; 384]; embedding[0] = 1.0; embedding - }, - ) - .await; - let related_b = create_memory_with_embedding( - &store, - &embedding_table, - "related B", - MemoryType::Fact, - 0.7, - { + }) + .await; + let related_b = + create_memory_with_embedding(&backend, "related B", MemoryType::Fact, 0.7, { let mut embedding = vec![0.0; 384]; embedding[1] = 1.0; embedding - }, - ) - .await; + }) + .await; - store + backend .create_association(&Association::new( &duplicate_a.id, &related_a.id, @@ -734,7 +658,7 @@ mod tests { )) .await .expect("failed to create duplicate_a association"); - store + backend .create_association(&Association::new( &related_b.id, &duplicate_b.id, @@ -745,9 +669,8 @@ mod tests { let embedding_model = shared_embedding_model(); let report = run_maintenance( - &store, - &embedding_table, - &embedding_model, + Arc::clone(&backend), + Arc::clone(&embedding_model), &MaintenanceConfig { prune_threshold: 0.2, decay_rate: 0.05, @@ -760,7 +683,7 @@ mod tests { assert_eq!(report.merged, 2); - let refreshed_survivor = store + let refreshed_survivor = backend .load(&survivor.id) .await .expect("failed to load survivor") @@ -770,21 +693,21 @@ mod tests { assert_eq!(refreshed_survivor.content, "durable rust maintenance note"); for duplicate_id in [&duplicate_a.id, &duplicate_b.id] { - let duplicate = store + let duplicate = backend .load(duplicate_id) .await .expect("failed to load duplicate") .expect("duplicate should exist"); assert!(duplicate.forgotten); - let duplicate_embeddings = embedding_table + let duplicate_embeddings = backend .find_similar(duplicate_id, 0.0, 10) .await .expect("failed to search duplicate embeddings"); assert!(duplicate_embeddings.is_empty()); } - let survivor_associations = store + let survivor_associations = backend .get_associations(&survivor.id) .await .expect("failed to load survivor associations"); @@ -825,23 +748,13 @@ mod tests { #[tokio::test] async fn run_maintenance_with_cancel_stops_when_cancel_requested() { - let store = MemoryStore::connect_in_memory().await; - - let dir = tempdir().expect("failed to create temp dir"); - let lance_conn = lancedb::connect(dir.path().to_str().expect("temp path")) - .execute() - .await - .expect("failed to connect to lancedb"); - let embedding_table = crate::memory::EmbeddingTable::open_or_create(&lance_conn) - .await - .expect("failed to create embedding table"); + let (backend, _dir) = sqlite_backend().await; let (_cancel_tx, maintenance_cancel_rx) = tokio::sync::watch::channel(true); let embedding_model = shared_embedding_model(); let result = run_maintenance_with_cancel( - &store, - &embedding_table, - &embedding_model, + Arc::clone(&backend), + Arc::clone(&embedding_model), &MaintenanceConfig::default(), maintenance_cancel_rx, ) @@ -865,15 +778,7 @@ mod tests { #[tokio::test] async fn run_maintenance_rejects_invalid_configuration_ranges() { - let store = MemoryStore::connect_in_memory().await; - let dir = tempdir().expect("failed to create temp dir"); - let lance_conn = lancedb::connect(dir.path().to_str().expect("temp path")) - .execute() - .await - .expect("failed to connect to lancedb"); - let embedding_table = crate::memory::EmbeddingTable::open_or_create(&lance_conn) - .await - .expect("failed to create embedding table"); + let (backend, _dir) = sqlite_backend().await; let invalid_config = MaintenanceConfig { prune_threshold: 0.2, @@ -883,8 +788,12 @@ mod tests { }; let embedding_model = shared_embedding_model(); - let result = - run_maintenance(&store, &embedding_table, &embedding_model, &invalid_config).await; + let result = run_maintenance( + Arc::clone(&backend), + Arc::clone(&embedding_model), + &invalid_config, + ) + .await; assert!(result.is_err(), "expected invalid config to fail"); assert!( result diff --git a/src/memory/search.rs b/src/memory/search.rs index 296686c81..f3f2fe765 100644 --- a/src/memory/search.rs +++ b/src/memory/search.rs @@ -1,8 +1,9 @@ //! Memory search: hybrid (vector + FTS + RRF + graph), temporal, importance, and typed queries. use crate::error::Result; +use crate::memory::EmbeddingModel; +use crate::memory::backend::MemoryBackend; use crate::memory::types::{Memory, MemorySearchResult, MemoryType, RelationType}; -use crate::memory::{EmbeddingModel, EmbeddingTable, MemoryStore}; use std::collections::HashMap; use std::sync::Arc; @@ -35,16 +36,14 @@ pub enum SearchSort { /// Bundles all memory search dependencies. pub struct MemorySearch { - store: Arc, - embedding_table: EmbeddingTable, + backend: Arc, embedding_model: Arc, } impl Clone for MemorySearch { fn clone(&self) -> Self { Self { - store: Arc::clone(&self.store), - embedding_table: self.embedding_table.clone(), + backend: Arc::clone(&self.backend), embedding_model: Arc::clone(&self.embedding_model), } } @@ -53,38 +52,23 @@ impl Clone for MemorySearch { impl std::fmt::Debug for MemorySearch { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("MemorySearch") - .field("store", &self.store) + .field("backend", &self.backend) .finish_non_exhaustive() } } impl MemorySearch { /// Create a new MemorySearch instance. - pub fn new( - store: Arc, - embedding_table: EmbeddingTable, - embedding_model: Arc, - ) -> Self { + pub fn new(backend: Arc, embedding_model: Arc) -> Self { Self { - store, - embedding_table, + backend, embedding_model, } } - /// Get a reference to the memory store. - pub fn store(&self) -> &MemoryStore { - &self.store - } - - /// Get a reference to the embedding table. - pub fn embedding_table(&self) -> &EmbeddingTable { - &self.embedding_table - } - - /// Get a reference to the embedding model. - pub fn embedding_model(&self) -> &EmbeddingModel { - &self.embedding_model + /// Get a reference to the backend. + pub fn backend(&self) -> &Arc { + &self.backend } /// Get a shared handle to the embedding model (for async embed_one). @@ -92,6 +76,11 @@ impl MemorySearch { &self.embedding_model } + /// Get the agent ID this search instance is scoped to. + pub fn agent_id(&self) -> &str { + self.backend.agent_id() + } + /// Unified search entry point. Dispatches to the appropriate strategy /// based on `config.mode`. pub async fn search( @@ -108,7 +97,7 @@ impl MemorySearch { #[cfg(feature = "metrics")] { - let agent_id = self.store.agent_id(); + let agent_id = self.backend.agent_id(); let agent_label = if agent_id.is_empty() { "unknown" } else { @@ -131,7 +120,7 @@ impl MemorySearch { config: &SearchConfig, ) -> Result> { let memories = self - .store + .backend .get_sorted(sort, config.max_results as i64, config.memory_type) .await?; @@ -173,13 +162,13 @@ impl MemorySearch { // FTS requires an inverted index. If the index doesn't exist yet (empty // table, first run) this will fail — fall back to vector + graph search. match self - .embedding_table + .backend .text_search(query, config.max_results_per_source) .await { Ok(fts_matches) => { for (memory_id, score) in fts_matches { - if let Some(memory) = self.store.load(&memory_id).await? + if let Some(memory) = self.backend.load(&memory_id).await? && !memory.forgotten { fts_results.push(ScoredMemory { @@ -197,14 +186,14 @@ impl MemorySearch { // 2. Vector similarity search via LanceDB let query_embedding = self.embedding_model.embed_one(query).await?; match self - .embedding_table + .backend .vector_search(&query_embedding, config.max_results_per_source) .await { Ok(vector_matches) => { for (memory_id, distance) in vector_matches { let similarity = 1.0 - distance; - if let Some(memory) = self.store.load(&memory_id).await? + if let Some(memory) = self.backend.load(&memory_id).await? && !memory.forgotten { vector_results.push(ScoredMemory { @@ -221,7 +210,7 @@ impl MemorySearch { // 3. Graph traversal from high-importance memories // Get identity and high-importance memories as starting points - let seed_memories = self.store.get_high_importance(0.8, 20).await?; + let seed_memories = self.backend.get_high_importance(0.8, 20).await?; for seed in seed_memories { // Check if seed is semantically related to query via simple keyword matching @@ -266,71 +255,140 @@ impl MemorySearch { Ok(results) } - /// Traverse the memory graph to find related memories (iterative to avoid async recursion). + /// Traverse the memory graph to find related memories (level-batched BFS). + /// + /// Uses two queries per BFS level — `get_associations_for` (all edges incident + /// to the current frontier) and `load_many` (batch-load all new neighbours) — + /// instead of one query per node + one per neighbour (was O(nodes) N+1, now + /// O(depth)×2). + /// + /// **Ordering / determinism note.** + /// The single collection pass below iterates `frontier` in its current order + /// and updates `visited` inline, so the first frontier node that reaches a + /// given neighbour wins (cross-frontier-node first-seen is deterministic). + /// Intra-node edge order (multiple edges from the *same* frontier node) is made + /// deterministic by sorting the fetched edges by descending weight (tie-broken by + /// the stable association id) before grouping, so traversal and ranking do not + /// depend on the backend's incidental row order or our IN-list chunking. async fn traverse_graph( &self, start_id: &str, max_depth: usize, results: &mut Vec, ) -> Result<()> { - use std::collections::VecDeque; - - // Queue of (memory_id, current_depth) - let mut queue: VecDeque<(String, usize)> = VecDeque::new(); let mut visited: std::collections::HashSet = std::collections::HashSet::new(); - - queue.push_back((start_id.to_string(), 0)); visited.insert(start_id.to_string()); - while let Some((current_id, depth)) = queue.pop_front() { - if depth > max_depth { - continue; + let mut frontier: Vec = vec![start_id.to_string()]; + let mut depth = 0usize; + + while !frontier.is_empty() && depth <= max_depth { + // One query for all edges incident to this level's frontier. Sort by + // descending weight (stronger associations first), tie-broken by the + // stable association id, so grouping, first-seen neighbour selection and + // ranking are deterministic and independent of the backend's incidental + // row order or our IN-list chunking. + let mut all_edges = self.backend.get_associations_for(&frontier).await?; + all_edges.sort_by(|a, b| { + b.weight + .partial_cmp(&a.weight) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.id.cmp(&b.id)) + }); + + // Group edges by the frontier node they are incident to (preferring + // the source endpoint when both are in the frontier). An edge whose + // BOTH endpoints are in the frontier is inert: every node enters the + // frontier only after being inserted into `visited`, so neither + // endpoint can be a new neighbour. Single-incident edges are the only + // ones that produce neighbours. + let mut by_node: std::collections::HashMap< + &str, + Vec<&crate::memory::types::Association>, + > = std::collections::HashMap::new(); + for edge in &all_edges { + let source_in = frontier.iter().any(|f| f == &edge.source_id); + let target_in = frontier.iter().any(|f| f == &edge.target_id); + if source_in { + by_node.entry(&edge.source_id).or_default().push(edge); + } else if target_in { + by_node.entry(&edge.target_id).or_default().push(edge); + } } - let associations = self.store.get_associations(¤t_id).await?; - - for assoc in associations { - // Get the related memory - let related_id = if assoc.source_id == current_id { - &assoc.target_id - } else { - &assoc.source_id - }; - - if visited.contains(related_id) { - continue; + // SINGLE collection pass: consult AND update `visited` inline. + // Iterating `frontier` in order ensures the first frontier node + // that reaches a neighbour claims it (deterministic cross-node case). + let mut new: Vec<(String, RelationType, f32)> = Vec::new(); + for fnode in &frontier { + if let Some(edges) = by_node.get(fnode.as_str()) { + for edge in edges.iter() { + // The neighbour is the endpoint that is NOT this frontier node. + let neighbor_id = if edge.source_id == *fnode { + &edge.target_id + } else { + &edge.source_id + }; + if visited.contains(neighbor_id) { + continue; + } + // Mark visited BEFORE load_many — mirrors the original + // `visited.insert` before `store.load` in the old code. + // Forgotten/missing neighbours are still marked visited + // and never reconsidered even if load_many omits them. + visited.insert(neighbor_id.clone()); + new.push((neighbor_id.clone(), edge.relation_type, edge.weight)); + } } - visited.insert(related_id.clone()); + } - if let Some(memory) = self.store.load(related_id).await? { + if new.is_empty() { + break; + } + + // One query to batch-load all new neighbours. + // load_many returns forgotten rows — the `forgotten` check is done + // in Rust below, AFTER marking visited (exact parity with original). + let new_ids: Vec = new.iter().map(|(id, _, _)| id.clone()).collect(); + let loaded: std::collections::HashMap = self + .backend + .load_many(&new_ids) + .await? + .into_iter() + .map(|m| (m.id.clone(), m)) + .collect(); + + let mut next_frontier: Vec = Vec::new(); + + for (nid, rel, weight) in new { + if let Some(memory) = loaded.get(&nid) { if memory.forgotten { + // Visited already inserted; skip scoring and expansion. continue; } - // Score based on relation type and weight - let type_multiplier = match assoc.relation_type { + let type_multiplier = match rel { RelationType::Updates => 1.5, RelationType::CausedBy | RelationType::ResultOf => 1.3, RelationType::RelatedTo => 1.0, RelationType::Contradicts => 0.5, RelationType::PartOf => 0.8, }; - - let score = memory.importance as f64 * assoc.weight as f64 * type_multiplier; - + let score = memory.importance as f64 * weight as f64 * type_multiplier; results.push(ScoredMemory { memory: memory.clone(), score, }); - - // Add to queue for RelatedTo and PartOf relations - if matches!( - assoc.relation_type, - RelationType::RelatedTo | RelationType::PartOf - ) { - queue.push_back((related_id.clone(), depth + 1)); + // Re-expand only RelatedTo / PartOf (same rule as original). + if matches!(rel, RelationType::RelatedTo | RelationType::PartOf) { + next_frontier.push(nid); } } + // Missing from load_many (not in DB): visited already inserted, + // score silently skipped — same as the original `load` returning None. } + + frontier = next_frontier; + depth += 1; } Ok(()) @@ -442,6 +500,8 @@ pub fn curate_results( #[cfg(test)] mod tests { use super::*; + use crate::memory::backend::SqliteBackend; + use crate::memory::lance::EmbeddingTable; use crate::memory::types::MemoryType; use chrono::{Duration, Utc}; @@ -581,7 +641,8 @@ mod tests { .unwrap(); let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); - let search = MemorySearch::new(store, embedding_table, embedding_model); + let backend = Arc::new(SqliteBackend::new(store, embedding_table)); + let search = MemorySearch::new(backend, embedding_model); let config = SearchConfig { mode: SearchMode::Recent, @@ -609,7 +670,8 @@ mod tests { .unwrap(); let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); - let search = MemorySearch::new(store, embedding_table, embedding_model); + let backend = Arc::new(SqliteBackend::new(store, embedding_table)); + let search = MemorySearch::new(backend, embedding_model); let config = SearchConfig { mode: SearchMode::Important, @@ -634,7 +696,8 @@ mod tests { .unwrap(); let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); - let search = MemorySearch::new(store, embedding_table, embedding_model); + let backend = Arc::new(SqliteBackend::new(store, embedding_table)); + let search = MemorySearch::new(backend, embedding_model); let config = SearchConfig { mode: SearchMode::Typed, @@ -659,7 +722,8 @@ mod tests { .unwrap(); let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); - let search = MemorySearch::new(store, embedding_table, embedding_model); + let backend = Arc::new(SqliteBackend::new(store, embedding_table)); + let search = MemorySearch::new(backend, embedding_model); let config = SearchConfig { mode: SearchMode::Typed, @@ -671,4 +735,158 @@ mod tests { let results = search.search("", &config).await.unwrap(); assert!(results.is_empty()); } + + // ── Characterization test for traverse_graph (Task D2) ──────────────────── + // + // Builds a deterministic graph and asserts the EXACT Vec + // (ids + scores) that the original N+1 BFS produced, now verified against + // the new level-batched implementation. + // + // Graph topology (max_depth = 1): + // + // start --RelatedTo, w=0.8--> A (importance=0.8) + // start --RelatedTo, w=0.7--> D (importance=0.6) + // start --RelatedTo, w=0.5--> F (importance=0.4, FORGOTTEN) + // + // A --RelatedTo, w=0.9 --> B (importance=0.7) re-expands + // A --Contradicts,w=0.6 --> C (importance=0.9) scored, NOT re-expanded + // A --RelatedTo, w=0.75--> E (importance=0.5) A reaches E first (frontier order) + // + // D --Updates, w=0.7 --> E (importance=0.5) E already visited from A → skip + // + // B --RelatedTo, w=0.8 --> G (importance=0.3) NOT reached: B is depth 2 > max_depth=1 + // + // Cases covered: + // ✓ RelatedTo chain re-expands (start→A→B) + // ✓ Contradicts scored but NOT re-expanded (A→C) + // ✓ Forgotten neighbour marked visited, not scored (start→F) + // ✓ Two DIFFERENT frontier nodes (A, D) reach E at the same level; + // A wins because it appears first in the frontier (deterministic cross-node) + // ✓ max_depth bound: B (depth 2) is not expanded → G never scored + // + // Golden scores ((f32_importance as f64) × (f32_weight as f64) × type_multiplier): + // A: 0.8f32 × 0.8f32 × 1.0 (RelatedTo) + // D: 0.6f32 × 0.7f32 × 1.0 (RelatedTo) + // B: 0.7f32 × 0.9f32 × 1.0 (RelatedTo) + // C: 0.9f32 × 0.6f32 × 0.5 (Contradicts) + // E: 0.5f32 × 0.75f32× 1.0 (RelatedTo — A reaches E first) + // + // Expected BFS push order: [A, D, B, C, E] + + async fn build_traverse_graph_fixture() -> (MemorySearch, String, tempfile::TempDir) { + use crate::memory::types::Association; + + let store = crate::memory::MemoryStore::connect_in_memory().await; + + macro_rules! save_mem { + ($content:expr, $imp:expr) => {{ + let m = Memory::new($content, MemoryType::Fact).with_importance($imp); + store.save(&m).await.unwrap(); + m + }}; + } + + let start_mem = save_mem!("start node", 1.0_f32); + let a = save_mem!("node A", 0.8_f32); + let b = save_mem!("node B", 0.7_f32); + let c = save_mem!("node C", 0.9_f32); + let d = save_mem!("node D", 0.6_f32); + let e = save_mem!("node E", 0.5_f32); + let f = save_mem!("node F forgotten", 0.4_f32); + store.forget(&f.id).await.unwrap(); + let _g = save_mem!("node G (unreachable)", 0.3_f32); + + // Edge weights are distinct per source node, so the descending-weight sort in + // traverse_graph yields a deterministic order regardless of insertion/row order. + let edge_specs: &[(&str, &str, RelationType, f32)] = &[ + (&start_mem.id, &a.id, RelationType::RelatedTo, 0.8), + (&start_mem.id, &d.id, RelationType::RelatedTo, 0.7), + (&start_mem.id, &f.id, RelationType::RelatedTo, 0.5), + (&a.id, &b.id, RelationType::RelatedTo, 0.9), + (&a.id, &c.id, RelationType::Contradicts, 0.6), + (&a.id, &e.id, RelationType::RelatedTo, 0.75), + (&d.id, &e.id, RelationType::Updates, 0.7), + (&b.id, &_g.id, RelationType::RelatedTo, 0.8), + ]; + for (src, tgt, rel, weight) in edge_specs { + store + .create_association(&Association::new(*src, *tgt, *rel).with_weight(*weight)) + .await + .unwrap(); + } + + let lance_dir = tempfile::tempdir().unwrap(); + let lance_conn = lancedb::connect(lance_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let embedding_table = EmbeddingTable::open_or_create(&lance_conn).await.unwrap(); + let embedding_model = Arc::new(EmbeddingModel::new(lance_dir.path()).unwrap()); + let backend = Arc::new(SqliteBackend::new(Arc::clone(&store), embedding_table)); + let search = MemorySearch::new(backend, embedding_model); + + (search, start_mem.id, lance_dir) + } + + fn assert_traverse_golden(results: &[ScoredMemory]) { + // Scores are (f32_importance as f64) * (f32_weight as f64) * type_multiplier, + // using the same f32 intermediates as the production code to get exact values. + fn score(imp: f32, weight: f32, mul: f64) -> f64 { + (imp as f64) * (weight as f64) * mul + } + // Expected BFS push order: [A, D, B, E, C]. Within a frontier node, edges are + // processed by descending weight, so among A's neighbours B (0.9) precedes + // E (0.75) precedes C (0.6). + let expected: &[(&str, f64)] = &[ + ("node A", score(0.8, 0.8, 1.0)), // RelatedTo, edge weight 0.8 + ("node D", score(0.6, 0.7, 1.0)), // RelatedTo, edge weight 0.7 + ("node B", score(0.7, 0.9, 1.0)), // RelatedTo, edge weight 0.9 + ("node E", score(0.5, 0.75, 1.0)), // RelatedTo, edge weight 0.75 (A reaches E first) + ("node C", score(0.9, 0.6, 0.5)), // Contradicts, edge weight 0.6 + ]; + assert_eq!( + results.len(), + expected.len(), + "result count mismatch: got {:?}", + results + .iter() + .map(|r| &r.memory.content) + .collect::>() + ); + for (i, (r, (exp_content, exp_score))) in results.iter().zip(expected.iter()).enumerate() { + assert_eq!( + r.memory.content, *exp_content, + "position {i}: expected content {exp_content:?}, got {:?}", + r.memory.content + ); + assert!( + (r.score - exp_score).abs() < 1e-9, + "position {i} ({exp_content}): expected score {exp_score}, got {}", + r.score + ); + } + assert!( + results + .iter() + .all(|r| r.memory.content != "node F forgotten"), + "forgotten node F must not appear in results" + ); + assert!( + results + .iter() + .all(|r| r.memory.content != "node G (unreachable)"), + "depth-bounded node G must not appear in results" + ); + } + + #[tokio::test] + async fn traverse_graph_batched_matches_golden() { + let (search, start_id, _dir) = build_traverse_graph_fixture().await; + let mut results: Vec = Vec::new(); + search + .traverse_graph(&start_id, 1, &mut results) + .await + .unwrap(); + assert_traverse_golden(&results); + } } diff --git a/src/memory/store.rs b/src/memory/store.rs index b2f537587..47c9b5909 100644 --- a/src/memory/store.rs +++ b/src/memory/store.rs @@ -457,33 +457,117 @@ impl MemoryStore { return Ok(Vec::new()); } - // Build a parameterized IN clause. SQLite handles this fine for - // the sizes we deal with (up to ~500 IDs). - let placeholders: String = memory_ids.iter().map(|_| "?").collect::>().join(","); - let query_str = format!( - "SELECT id, source_id, target_id, relation_type, weight, created_at \ - FROM associations \ - WHERE source_id IN ({placeholders}) AND target_id IN ({placeholders})" - ); - - let mut query = sqlx::query(&query_str); - // Bind once for source_id IN, once for target_id IN - for id in memory_ids { - query = query.bind(id); + // Chunk the IN-list to stay under SQLite's bind-parameter limit on large + // sets. Query `source_id IN (chunk)` per chunk (one bind per id) and filter + // `target_id` against the full set in Rust, so cross-chunk pairs (source in + // one chunk, target in another) are not missed. + const CHUNK: usize = 400; + let id_set: std::collections::HashSet<&str> = + memory_ids.iter().map(String::as_str).collect(); + let mut seen: std::collections::HashSet = std::collections::HashSet::new(); + let mut out = Vec::new(); + + for chunk in memory_ids.chunks(CHUNK) { + let placeholders: String = chunk.iter().map(|_| "?").collect::>().join(","); + let query_str = format!( + "SELECT id, source_id, target_id, relation_type, weight, created_at \ + FROM associations WHERE source_id IN ({placeholders})" + ); + let mut query = sqlx::query(&query_str); + for id in chunk { + query = query.bind(id); + } + let rows = query + .fetch_all(&self.pool) + .await + .context("failed to get associations between memory set")?; + for row in rows { + let assoc = row_to_association(&row); + if id_set.contains(assoc.target_id.as_str()) && seen.insert(assoc.id.clone()) { + out.push(assoc); + } + } } - for id in memory_ids { - query = query.bind(id); + + Ok(out) + } + + /// All associations incident to ANY of `ids` (either endpoint). + /// Returns only associations where `source_id IN ids OR target_id IN ids`. + /// Empty `ids` → empty result. + pub async fn get_associations_for(&self, ids: &[String]) -> Result> { + if ids.is_empty() { + return Ok(Vec::new()); } - let rows = query - .fetch_all(&self.pool) - .await - .context("failed to get associations between memory set")?; + // Chunk the IN-list to stay under SQLite's bind-parameter limit. An edge + // incident to ids in two different chunks is matched twice, so dedup by + // association id. + const CHUNK: usize = 400; + let mut seen: std::collections::HashSet = std::collections::HashSet::new(); + let mut out = Vec::new(); + + for chunk in ids.chunks(CHUNK) { + let placeholders: String = chunk.iter().map(|_| "?").collect::>().join(","); + let query_str = format!( + "SELECT id, source_id, target_id, relation_type, weight, created_at \ + FROM associations \ + WHERE source_id IN ({placeholders}) OR target_id IN ({placeholders})" + ); + let mut query = sqlx::query(&query_str); + for id in chunk { + query = query.bind(id); + } + for id in chunk { + query = query.bind(id); + } + let rows = query + .fetch_all(&self.pool) + .await + .context("failed to get associations for id set")?; + for row in rows { + let assoc = row_to_association(&row); + if seen.insert(assoc.id.clone()) { + out.push(assoc); + } + } + } - Ok(rows - .into_iter() - .map(|row| row_to_association(&row)) - .collect()) + Ok(out) + } + + /// Batch-load memories by ID. Missing IDs are silently omitted. + /// Order is unspecified. Forgotten memories ARE included — callers must + /// check `memory.forgotten` themselves. + /// Empty `ids` → empty result. + pub async fn load_many(&self, ids: &[String]) -> Result> { + if ids.is_empty() { + return Ok(Vec::new()); + } + + // Chunk the IN-list to stay under SQLite's bind-parameter limit. + const CHUNK: usize = 400; + let mut out = Vec::new(); + + for chunk in ids.chunks(CHUNK) { + let placeholders: String = chunk.iter().map(|_| "?").collect::>().join(","); + let query_str = format!( + "SELECT id, content, memory_type, importance, created_at, updated_at, \ + last_accessed_at, access_count, source, channel_id, forgotten \ + FROM memories WHERE id IN ({placeholders})" + ); + let mut query = sqlx::query(&query_str); + for id in chunk { + query = query.bind(id); + } + let rows = query + .fetch_all(&self.pool) + .await + .context("failed to batch-load memories")?; + out.extend(rows.into_iter().map(|row| row_to_memory(&row))); + } + + Ok(out) } /// Get neighbors of a memory: all associations plus the connected memories. diff --git a/src/tools/memory_delete.rs b/src/tools/memory_delete.rs index de27f7c92..a62fd7498 100644 --- a/src/tools/memory_delete.rs +++ b/src/tools/memory_delete.rs @@ -85,7 +85,7 @@ impl Tool for MemoryDeleteTool { } async fn call(&self, args: Self::Args) -> std::result::Result { - let store = self.memory_search.store(); + let store = self.memory_search.backend(); // Verify the memory exists first let memory = store diff --git a/src/tools/memory_recall.rs b/src/tools/memory_recall.rs index 42bce96e7..b10116b22 100644 --- a/src/tools/memory_recall.rs +++ b/src/tools/memory_recall.rs @@ -229,7 +229,7 @@ impl Tool for MemoryRecallTool { let curated = curate_results(&search_results, args.max_results); - let store = self.memory_search.store(); + let store = self.memory_search.backend(); let mut memories = Vec::new(); for result in &curated { @@ -256,7 +256,7 @@ impl Tool for MemoryRecallTool { #[cfg(feature = "metrics")] { - let agent_id = self.memory_search.store().agent_id(); + let agent_id = self.memory_search.backend().agent_id(); let agent_label = if agent_id.is_empty() { "unknown" } else { @@ -320,7 +320,7 @@ pub async fn memory_recall( .map_err(|e| crate::error::AgentError::Other(anyhow::anyhow!(e)))?; // Convert back to Memory type for backward compatibility - let store = memory_search.store(); + let store = memory_search.backend(); let mut memories = Vec::new(); for mem_out in output.memories { diff --git a/src/tools/memory_save.rs b/src/tools/memory_save.rs index 8193d7f50..e658efcdf 100644 --- a/src/tools/memory_save.rs +++ b/src/tools/memory_save.rs @@ -253,9 +253,9 @@ impl Tool for MemorySaveTool { } // Save to SQLite database - let store = self.memory_search.store(); + let store = self.memory_search.backend(); store - .save(&memory) + .save(&memory, None) .await .map_err(|e| MemorySaveError(format!("Failed to save memory: {e}")))?; @@ -319,7 +319,7 @@ impl Tool for MemorySaveTool { Err(embed_err) => { if let Err(assoc_err) = self .memory_search - .store() + .backend() .delete_associations_for_memory(&memory.id) .await { @@ -329,7 +329,7 @@ impl Tool for MemorySaveTool { "compensating association delete failed after embedding generation error" ); } - if let Err(del_err) = self.memory_search.store().delete(&memory.id).await { + if let Err(del_err) = self.memory_search.backend().delete(&memory.id).await { tracing::error!( memory_id = %memory.id, %del_err, @@ -344,8 +344,8 @@ impl Tool for MemorySaveTool { match self .memory_search - .embedding_table() - .store(&memory.id, &args.content, &embedding) + .backend() + .set_embedding(&memory, &embedding) .await { Ok(()) => { @@ -356,7 +356,7 @@ impl Tool for MemorySaveTool { Err(embed_err) => { if let Err(assoc_err) = self .memory_search - .store() + .backend() .delete_associations_for_memory(&memory.id) .await { @@ -366,7 +366,7 @@ impl Tool for MemorySaveTool { "compensating association delete failed after embedding store error" ); } - if let Err(del_err) = self.memory_search.store().delete(&memory.id).await { + if let Err(del_err) = self.memory_search.backend().delete(&memory.id).await { tracing::error!( memory_id = %memory.id, %del_err, @@ -379,17 +379,6 @@ impl Tool for MemorySaveTool { } } - // Ensure the FTS index exists so full_text_search queries work. - // Safe to call repeatedly — no-ops if the index already exists. - if let Err(error) = self - .memory_search - .embedding_table() - .ensure_fts_index() - .await - { - tracing::warn!(%error, "failed to ensure FTS index after memory save"); - } - if let Some(event_context) = &self.event_context && event_context.memory_event_tx.receiver_count() > 0 { @@ -429,7 +418,7 @@ impl Tool for MemorySaveTool { #[cfg(feature = "metrics")] { - let agent_id = self.memory_search.store().agent_id(); + let agent_id = self.memory_search.backend().agent_id(); let agent_label = if agent_id.is_empty() { "unknown" } else {