Skip to content

Commit c7dfc5b

Browse files
committed
feat(memory): add steel-memory semantic vector search
Adds optional steel-memory integration for semantic memory search using vector embeddings (AllMiniLML6V2). When the 'steel-memory' feature is enabled: - memory_search uses semantic similarity instead of BM25 keyword matching - Memories are stored in .steel-memory/palace.sqlite3 - New add_memory tool for directly adding to the vector index - Workspace files (MEMORY.md, memory/*.md) are auto-indexed The steel-memory feature is optional and included in 'full'. Without it, the existing BM25-based search with temporal decay continues to work. Dependencies: - steel-memory (git) - embedding + SQLite vector storage - fastembed re-exported for embeddings
1 parent 96ba8d7 commit c7dfc5b

File tree

6 files changed

+513
-11
lines changed

6 files changed

+513
-11
lines changed

crates/rustyclaw-core/Cargo.toml

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ default = ["web-tools"]
2020
web-tools = ["dep:scraper", "dep:html2md"]
2121
browser = ["dep:chromiumoxide"]
2222
mcp = ["dep:rmcp", "dep:schemars"]
23+
steel-memory = ["dep:steel-memory"]
2324
matrix = ["chat-system/matrix"]
2425
whatsapp = ["chat-system/whatsapp"]
2526
ssh = ["dep:russh", "dep:russh-keys", "dep:rand_core", "dep:sha2"]
@@ -31,7 +32,7 @@ telegram-cli = []
3132
discord-cli = []
3233
slack-cli = []
3334
all-messengers = ["whatsapp", "signal-cli", "matrix-cli", "telegram-cli", "discord-cli", "slack-cli"]
34-
full = ["web-tools", "browser", "mcp", "all-messengers", "ssh"]
35+
full = ["web-tools", "browser", "mcp", "all-messengers", "ssh", "steel-memory"]
3536

3637
[dependencies]
3738
serde.workspace = true
@@ -94,6 +95,9 @@ chromiumoxide = { workspace = true, optional = true }
9495
rmcp = { workspace = true, optional = true }
9596
schemars = { workspace = true, optional = true }
9697

98+
# Steel Memory - semantic vector search memory palace
99+
steel-memory = { git = "https://github.com/rexlunae/steel-memory", package = "steel-memory", optional = true }
100+
97101
# SSH transport (optional)
98102
russh = { version = "0.52", features = ["async-trait"], optional = true }
99103
russh-keys = { version = "0.44", optional = true }

crates/rustyclaw-core/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ pub mod mcp;
1717
pub mod memory;
1818
pub mod memory_consolidation;
1919
pub mod memory_flush;
20+
#[cfg(feature = "steel-memory")]
21+
pub mod steel_memory;
2022
pub mod messengers;
2123
pub mod mnemo;
2224
pub mod models;
Lines changed: 352 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,352 @@
1+
//! Steel Memory integration for RustyClaw.
2+
//!
3+
//! Provides semantic vector search over agent memories using steel-memory's
4+
//! embedding and SQLite-backed storage.
5+
//!
6+
//! This replaces the BM25-based keyword search with true semantic search,
7+
//! delivering much better recall for natural language queries.
8+
9+
use std::path::{Path, PathBuf};
10+
use std::sync::Arc;
11+
use tokio::sync::Mutex;
12+
use tracing::{debug, info};
13+
14+
use steel_memory::{
15+
fastembed::{EmbeddingModel, InitOptions, TextEmbedding},
16+
storage::vector::VectorStorage,
17+
types::{Drawer, SearchResult as SteelSearchResult},
18+
};
19+
20+
/// A semantic memory index using steel-memory's vector storage.
21+
pub struct SteelMemoryIndex {
22+
/// Path to the vector database
23+
db_path: PathBuf,
24+
/// Path to the palace/workspace directory
25+
palace_path: PathBuf,
26+
/// Embedding model (lazy initialized)
27+
embedding: Arc<Mutex<Option<TextEmbedding>>>,
28+
}
29+
30+
/// A search result from steel-memory.
31+
#[derive(Debug, Clone)]
32+
pub struct SearchResult {
33+
/// The matching memory content
34+
pub content: String,
35+
/// Source file path
36+
pub path: String,
37+
/// Wing (category) of the memory
38+
pub wing: String,
39+
/// Room (subcategory) of the memory
40+
pub room: String,
41+
/// Similarity score (0.0 to 1.0)
42+
pub similarity: f32,
43+
/// Drawer ID for reference
44+
pub id: String,
45+
}
46+
47+
impl From<SteelSearchResult> for SearchResult {
48+
fn from(r: SteelSearchResult) -> Self {
49+
Self {
50+
content: r.drawer.content,
51+
path: r.drawer.source_file,
52+
wing: r.drawer.wing,
53+
room: r.drawer.room,
54+
similarity: r.similarity,
55+
id: r.drawer.id,
56+
}
57+
}
58+
}
59+
60+
impl SteelMemoryIndex {
61+
/// Create a new steel-memory index for a workspace.
62+
///
63+
/// The database will be stored at `workspace/.steel-memory/palace.sqlite3`.
64+
pub fn new(workspace: &Path) -> Result<Self, String> {
65+
let steel_dir = workspace.join(".steel-memory");
66+
std::fs::create_dir_all(&steel_dir)
67+
.map_err(|e| format!("Failed to create .steel-memory directory: {}", e))?;
68+
69+
let db_path = steel_dir.join("palace.sqlite3");
70+
71+
// Initialize storage to create tables
72+
VectorStorage::new(&db_path)
73+
.map_err(|e| format!("Failed to initialize vector storage: {}", e))?;
74+
75+
Ok(Self {
76+
db_path,
77+
palace_path: workspace.to_path_buf(),
78+
embedding: Arc::new(Mutex::new(None)),
79+
})
80+
}
81+
82+
/// Ensure the embedding model is loaded.
83+
async fn ensure_embedding(&self) -> Result<(), String> {
84+
let mut guard: tokio::sync::MutexGuard<'_, Option<TextEmbedding>> = self.embedding.lock().await;
85+
if guard.is_none() {
86+
info!("Loading embedding model (AllMiniLML6V2)...");
87+
let model = tokio::task::spawn_blocking(|| -> Result<TextEmbedding, String> {
88+
TextEmbedding::try_new(InitOptions::new(EmbeddingModel::AllMiniLML6V2))
89+
.map_err(|e| format!("Failed to load embedding model: {}", e))
90+
})
91+
.await
92+
.map_err(|e| format!("Embedding task failed: {}", e))??;
93+
*guard = Some(model);
94+
info!("Embedding model loaded");
95+
}
96+
Ok(())
97+
}
98+
99+
/// Embed text into a vector.
100+
async fn embed(&self, text: &str) -> Result<Vec<f32>, String> {
101+
self.ensure_embedding().await?;
102+
103+
let embedding = self.embedding.clone();
104+
let text_owned = text.to_string();
105+
106+
let result = tokio::task::spawn_blocking(move || -> Result<Vec<f32>, String> {
107+
let mut guard = embedding.blocking_lock();
108+
let model = guard.as_mut().ok_or_else(|| "Embedding model not initialized".to_string())?;
109+
let mut embeddings = model
110+
.embed(vec![text_owned.as_str()], None)
111+
.map_err(|e| format!("Embedding failed: {}", e))?;
112+
Ok(embeddings.remove(0))
113+
})
114+
.await
115+
.map_err(|e| format!("Embedding task failed: {}", e))??;
116+
117+
Ok(result)
118+
}
119+
120+
/// Search for memories matching a query.
121+
pub async fn search(
122+
&self,
123+
query: &str,
124+
max_results: usize,
125+
min_score: Option<f32>,
126+
) -> Result<Vec<SearchResult>, String> {
127+
debug!(query, max_results, "Searching steel-memory");
128+
129+
let query_vec = self.embed(query).await?;
130+
let db_path = self.db_path.clone();
131+
let min_score = min_score.unwrap_or(0.3);
132+
133+
let results = tokio::task::spawn_blocking(move || -> Result<Vec<SteelSearchResult>, String> {
134+
let storage = VectorStorage::new(&db_path)
135+
.map_err(|e| format!("Failed to open storage: {}", e))?;
136+
storage
137+
.search(&query_vec, max_results * 2, None, None) // Over-fetch to filter
138+
.map_err(|e| format!("Search failed: {}", e))
139+
})
140+
.await
141+
.map_err(|e| format!("Search task failed: {}", e))??;
142+
143+
Ok(results
144+
.into_iter()
145+
.filter(|r| r.similarity >= min_score)
146+
.take(max_results)
147+
.map(SearchResult::from)
148+
.collect())
149+
}
150+
151+
/// Add a memory to the index.
152+
pub async fn add_memory(
153+
&self,
154+
content: &str,
155+
wing: &str,
156+
room: &str,
157+
source_file: Option<&str>,
158+
) -> Result<String, String> {
159+
let vec = self.embed(content).await?;
160+
let id = uuid::Uuid::new_v4().to_string();
161+
162+
let drawer = Drawer {
163+
id: id.clone(),
164+
content: content.to_string(),
165+
wing: wing.to_string(),
166+
room: room.to_string(),
167+
source_file: source_file.unwrap_or("rustyclaw").to_string(),
168+
source_mtime: 0,
169+
chunk_index: 0,
170+
added_by: "rustyclaw".to_string(),
171+
filed_at: chrono::Utc::now().to_rfc3339(),
172+
hall: String::new(),
173+
topic: String::new(),
174+
drawer_type: String::new(),
175+
agent: "rustyclaw".to_string(),
176+
date: chrono::Utc::now().format("%Y-%m-%d").to_string(),
177+
importance: 3.0,
178+
};
179+
180+
let db_path = self.db_path.clone();
181+
tokio::task::spawn_blocking(move || -> Result<(), String> {
182+
let storage = VectorStorage::new(&db_path)
183+
.map_err(|e| format!("Failed to open storage: {}", e))?;
184+
storage
185+
.add_drawer(&drawer, &vec)
186+
.map_err(|e| format!("Failed to add drawer: {}", e))
187+
})
188+
.await
189+
.map_err(|e| format!("Add task failed: {}", e))??;
190+
191+
debug!(id = %id, wing, room, "Added memory to steel-memory");
192+
Ok(id)
193+
}
194+
195+
/// Index workspace memory files (MEMORY.md, memory/*.md).
196+
///
197+
/// This reads markdown files and chunks them into the vector database,
198+
/// replacing the BM25 index with semantic embeddings.
199+
pub async fn index_workspace(&self) -> Result<usize, String> {
200+
info!(workspace = %self.palace_path.display(), "Indexing workspace memories");
201+
202+
let mut count = 0;
203+
204+
// Index MEMORY.md
205+
let memory_md = self.palace_path.join("MEMORY.md");
206+
if memory_md.exists() {
207+
count += self.index_file(&memory_md, "MEMORY.md", "memory", "long-term").await?;
208+
}
209+
210+
// Index memory/*.md
211+
let memory_dir = self.palace_path.join("memory");
212+
if memory_dir.exists() && memory_dir.is_dir() {
213+
for entry in std::fs::read_dir(&memory_dir)
214+
.map_err(|e| format!("Failed to read memory dir: {}", e))?
215+
{
216+
let entry = entry.map_err(|e| format!("Failed to read entry: {}", e))?;
217+
let path = entry.path();
218+
if path.extension().map(|e| e == "md").unwrap_or(false) {
219+
let name = path.file_name().unwrap().to_string_lossy();
220+
let relative = format!("memory/{}", name);
221+
222+
// Use date as room if filename is YYYY-MM-DD.md
223+
let room = if name.len() == 13 && name.chars().take(10).all(|c| c.is_ascii_digit() || c == '-') {
224+
name.trim_end_matches(".md").to_string()
225+
} else {
226+
"notes".to_string()
227+
};
228+
229+
count += self.index_file(&path, &relative, "memory", &room).await?;
230+
}
231+
}
232+
}
233+
234+
info!(count, "Indexed memory files");
235+
Ok(count)
236+
}
237+
238+
/// Index a single markdown file.
239+
async fn index_file(
240+
&self,
241+
path: &Path,
242+
relative_path: &str,
243+
wing: &str,
244+
room: &str,
245+
) -> Result<usize, String> {
246+
let content = std::fs::read_to_string(path)
247+
.map_err(|e| format!("Failed to read {}: {}", relative_path, e))?;
248+
249+
let chunks = chunk_markdown(&content);
250+
let mut count = 0;
251+
252+
for chunk in chunks {
253+
if chunk.trim().is_empty() {
254+
continue;
255+
}
256+
257+
// Add memory (deduplication can be added later with content hashing)
258+
self.add_memory(&chunk, wing, room, Some(relative_path)).await?;
259+
count += 1;
260+
}
261+
262+
debug!(path = %relative_path, chunks = count, "Indexed file");
263+
Ok(count)
264+
}
265+
266+
/// Get total number of memories.
267+
pub async fn count(&self) -> Result<usize, String> {
268+
let db_path = self.db_path.clone();
269+
let result = tokio::task::spawn_blocking(move || -> Result<Vec<Drawer>, String> {
270+
let storage = VectorStorage::new(&db_path)
271+
.map_err(|e| format!("Failed to open storage: {}", e))?;
272+
storage
273+
.get_all(None, None, usize::MAX)
274+
.map_err(|e| format!("Failed to count: {}", e))
275+
})
276+
.await
277+
.map_err(|e| format!("Count task failed: {}", e))??;
278+
279+
Ok(result.len())
280+
}
281+
}
282+
283+
/// Chunk markdown content into sections.
284+
fn chunk_markdown(content: &str) -> Vec<String> {
285+
let mut chunks = Vec::new();
286+
let mut current_chunk = String::new();
287+
let mut line_count = 0;
288+
289+
for line in content.lines() {
290+
let is_heading = line.starts_with("## ") || line.starts_with("# ");
291+
292+
// Start new chunk on heading or every ~20 lines
293+
if (is_heading || line_count >= 20) && !current_chunk.trim().is_empty() {
294+
chunks.push(current_chunk.trim().to_string());
295+
current_chunk = String::new();
296+
line_count = 0;
297+
}
298+
299+
current_chunk.push_str(line);
300+
current_chunk.push('\n');
301+
line_count += 1;
302+
}
303+
304+
// Don't forget the last chunk
305+
if !current_chunk.trim().is_empty() {
306+
chunks.push(current_chunk.trim().to_string());
307+
}
308+
309+
chunks
310+
}
311+
312+
#[cfg(test)]
313+
mod tests {
314+
use super::*;
315+
use tempfile::TempDir;
316+
use std::fs;
317+
318+
#[tokio::test]
319+
async fn test_basic_search() {
320+
let dir = TempDir::new().unwrap();
321+
let index = SteelMemoryIndex::new(dir.path()).unwrap();
322+
323+
// Add some memories
324+
index.add_memory("I love programming in Rust", "preferences", "languages", None).await.unwrap();
325+
index.add_memory("Python is great for data science", "preferences", "languages", None).await.unwrap();
326+
index.add_memory("The sky is blue today", "observations", "weather", None).await.unwrap();
327+
328+
// Search for programming
329+
let results = index.search("Rust programming", 5, None).await.unwrap();
330+
assert!(!results.is_empty());
331+
assert!(results[0].content.contains("Rust"));
332+
}
333+
334+
#[tokio::test]
335+
async fn test_index_workspace() {
336+
let dir = TempDir::new().unwrap();
337+
338+
// Create test files
339+
fs::write(dir.path().join("MEMORY.md"), "# Memory\n\nI like cats.").unwrap();
340+
fs::create_dir(dir.path().join("memory")).unwrap();
341+
fs::write(dir.path().join("memory/2026-04-13.md"), "# Today\n\nWent for a walk.").unwrap();
342+
343+
let index = SteelMemoryIndex::new(dir.path()).unwrap();
344+
let count = index.index_workspace().await.unwrap();
345+
346+
assert!(count >= 2);
347+
348+
// Search should find results
349+
let results = index.search("cats", 5, None).await.unwrap();
350+
assert!(!results.is_empty());
351+
}
352+
}

0 commit comments

Comments
 (0)