Skip to main content

walrus_memory/
lib.rs

1//! Memory backends for Walrus agents.
2//!
3//! Defines the [`Memory`] trait, [`Embedder`] trait, and concrete implementations:
4//! [`InMemory`] (volatile) and [`SqliteMemory`] (persistent with FTS5 + vector recall).
5//!
6//! Memory is **not chat history**. It is structured knowledge — extracted facts,
7//! user preferences, agent persona — that gets compiled into the system prompt.
8//!
9//! All SQL lives in `sql/*.sql` files, loaded via `include_str!`.
10
11pub use crate::utils::cosine_similarity;
12use anyhow::Result;
13use compact_str::CompactString;
14use serde_json::Value;
15use std::{future::Future, sync::Arc};
16
17mod embedder;
18mod mem;
19mod sqlite;
20pub mod tools;
21mod utils;
22
23pub use embedder::{Embedder, NoEmbedder};
24pub use mem::InMemory;
25pub use sqlite::SqliteMemory;
26
27/// A structured memory entry with metadata and optional embedding.
28#[derive(Debug, Clone, Default)]
29pub struct MemoryEntry {
30    /// Entry key (identity string).
31    pub key: CompactString,
32    /// Entry value (unbounded content).
33    pub value: String,
34    /// Optional structured metadata (JSON).
35    pub metadata: Option<Value>,
36    /// Unix timestamp when the entry was created.
37    pub created_at: u64,
38    /// Unix timestamp when the entry was last accessed.
39    pub accessed_at: u64,
40    /// Number of times the entry has been accessed.
41    pub access_count: u32,
42    /// Optional embedding vector for semantic search.
43    pub embedding: Option<Vec<f32>>,
44}
45
46/// Options controlling memory recall behavior.
47#[derive(Debug, Clone, Default)]
48pub struct RecallOptions {
49    /// Maximum number of results (0 = implementation default).
50    pub limit: usize,
51    /// Filter by creation time range (start, end) in unix seconds.
52    pub time_range: Option<(u64, u64)>,
53    /// Minimum relevance score threshold (0.0–1.0).
54    pub relevance_threshold: Option<f32>,
55}
56
57/// Structured knowledge memory for LLM agents.
58///
59/// Implementations store named key-value pairs that get compiled
60/// into the system prompt via [`compile()`](Memory::compile).
61///
62/// Uses `&self` for all methods — implementations must handle
63/// interior mutability (e.g. via `Mutex`).
64pub trait Memory: Send + Sync {
65    /// Get the value for a key (owned).
66    fn get(&self, key: &str) -> Option<String>;
67
68    /// Get all key-value pairs (owned).
69    fn entries(&self) -> Vec<(String, String)>;
70
71    /// Set (upsert) a key-value pair. Returns the previous value if the key existed.
72    fn set(&self, key: impl Into<String>, value: impl Into<String>) -> Option<String>;
73
74    /// Remove a key. Returns the removed value if it existed.
75    fn remove(&self, key: &str) -> Option<String>;
76
77    /// Compile all entries into a string for system prompt injection.
78    fn compile(&self) -> String {
79        let entries = self.entries();
80        if entries.is_empty() {
81            return String::new();
82        }
83
84        let mut out = String::from("<memory>\n");
85        for (key, value) in &entries {
86            out.push_str(&format!("<{key}>\n"));
87            out.push_str(value);
88            if !value.ends_with('\n') {
89                out.push('\n');
90            }
91            out.push_str(&format!("</{key}>\n"));
92        }
93        out.push_str("</memory>");
94        out
95    }
96
97    /// Store a key-value pair (async). Default delegates to `set`.
98    fn store(
99        &self,
100        key: impl Into<String> + Send,
101        value: impl Into<String> + Send,
102    ) -> impl Future<Output = Result<()>> + Send {
103        self.set(key, value);
104        async { Ok(()) }
105    }
106
107    /// Search for relevant entries (async). Default returns empty.
108    fn recall(
109        &self,
110        _query: &str,
111        _options: RecallOptions,
112    ) -> impl Future<Output = Result<Vec<MemoryEntry>>> + Send {
113        async { Ok(Vec::new()) }
114    }
115
116    /// Compile relevant entries for a query (async). Default delegates to `compile`.
117    fn compile_relevant(&self, _query: &str) -> impl Future<Output = String> + Send {
118        let compiled = self.compile();
119        async move { compiled }
120    }
121}
122
123/// Apply memory to an agent config — appends compiled memory to the system prompt.
124pub fn with_memory(mut config: wcore::AgentConfig, memory: &impl Memory) -> wcore::AgentConfig {
125    let compiled = memory.compile();
126    if !compiled.is_empty() {
127        config.system_prompt = format!("{}\n\n{compiled}", config.system_prompt);
128    }
129    config
130}
131
132impl wcore::Hook for InMemory {
133    fn on_build_agent(&self, config: wcore::AgentConfig) -> wcore::AgentConfig {
134        with_memory(config, self)
135    }
136
137    fn on_register_tools(
138        &self,
139        registry: &mut wcore::ToolRegistry,
140    ) -> impl std::future::Future<Output = ()> + Send {
141        let mem = Arc::new(self.clone());
142        let remember = tools::remember(Arc::clone(&mem));
143        let recall = tools::recall(mem);
144        registry.insert(remember.tool, remember.handler);
145        registry.insert(recall.tool, recall.handler);
146        async {}
147    }
148}