Skip to main content

daimon_core/
core_memory.rs

1//! Core (always-in-context) memory trait.
2//!
3//! Unlike [`Memory`](crate::memory::Memory), which holds the rolling
4//! conversation log, [`CoreMemory`] holds a small, bounded set of
5//! self-editable facts (persona, user preferences, standing instructions)
6//! that a host application renders into the system prompt on every turn —
7//! the MemGPT/Letta "core memory" pattern. Built-in implementations live in
8//! the `daimon` facade crate.
9
10use std::future::Future;
11use std::pin::Pin;
12use std::sync::Arc;
13
14use crate::error::Result;
15
16/// A single named, optionally size-limited block of core memory.
17///
18/// `limit` is a character count; `None` means unbounded. Implementations of
19/// [`CoreMemory::put_block`] and [`CoreMemory::append_block`] must reject
20/// writes that would push `value` past `limit`.
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct CoreMemoryBlock {
23    /// Short identifier for the block (e.g. `"persona"`, `"user"`).
24    pub label: String,
25    /// The block's current text content.
26    pub value: String,
27    /// Maximum length of `value` in characters, or `None` for unbounded.
28    pub limit: Option<usize>,
29}
30
31impl CoreMemoryBlock {
32    /// Creates a new block with no size limit.
33    pub fn new(label: impl Into<String>, value: impl Into<String>) -> Self {
34        Self {
35            label: label.into(),
36            value: value.into(),
37            limit: None,
38        }
39    }
40
41    /// Sets a character limit on this block.
42    pub fn with_limit(mut self, limit: usize) -> Self {
43        self.limit = Some(limit);
44        self
45    }
46}
47
48/// Trait for always-in-context core memory backends.
49///
50/// Core memory is distinct from [`Memory`](crate::memory::Memory): it is not
51/// a chronological log, it is a small set of labeled blocks that are
52/// rendered as a whole and re-injected into the system prompt every turn.
53/// The agent itself (via a tool) or the host application can read and edit
54/// blocks between turns.
55pub trait CoreMemory: Send + Sync {
56    /// Returns all blocks currently stored.
57    fn blocks(&self) -> impl Future<Output = Result<Vec<CoreMemoryBlock>>> + Send;
58
59    /// Returns a single block by label, if it exists.
60    fn get_block(
61        &self,
62        label: &str,
63    ) -> impl Future<Output = Result<Option<CoreMemoryBlock>>> + Send {
64        async move { Ok(self.blocks().await?.into_iter().find(|b| b.label == label)) }
65    }
66
67    /// Creates a block (if `block.label` is new) or overwrites the value and
68    /// limit of an existing one. Fails if `block.value` exceeds `block.limit`.
69    fn put_block(&self, block: CoreMemoryBlock) -> impl Future<Output = Result<()>> + Send;
70
71    /// Appends `text` to an existing block's value (creating an unbounded
72    /// block first if `label` is unknown). Fails without modifying the block
73    /// if the result would exceed the block's limit.
74    fn append_block(&self, label: &str, text: &str) -> impl Future<Output = Result<()>> + Send;
75
76    /// Removes a block. Returns `true` if it existed.
77    fn remove_block(&self, label: &str) -> impl Future<Output = Result<bool>> + Send;
78
79    /// Renders all blocks into a single string suitable for splicing into a
80    /// system prompt. Default rendering is `"## {label}\n{value}"` per
81    /// block, joined by blank lines, in the order returned by [`blocks`](CoreMemory::blocks).
82    fn render(&self) -> impl Future<Output = Result<String>> + Send {
83        async move { Ok(render_blocks(&self.blocks().await?)) }
84    }
85}
86
87/// Shared rendering logic used by [`CoreMemory::render`]'s default
88/// implementation and available to custom implementations that want the
89/// same format.
90pub fn render_blocks(blocks: &[CoreMemoryBlock]) -> String {
91    blocks
92        .iter()
93        .map(|b| format!("## {}\n{}", b.label, b.value))
94        .collect::<Vec<_>>()
95        .join("\n\n")
96}
97
98/// Object-safe wrapper for the `CoreMemory` trait, enabling dynamic dispatch
99/// via `Arc<dyn ErasedCoreMemory>`.
100pub trait ErasedCoreMemory: Send + Sync {
101    fn blocks_erased(
102        &self,
103    ) -> Pin<Box<dyn Future<Output = Result<Vec<CoreMemoryBlock>>> + Send + '_>>;
104
105    fn get_block_erased<'a>(
106        &'a self,
107        label: &'a str,
108    ) -> Pin<Box<dyn Future<Output = Result<Option<CoreMemoryBlock>>> + Send + 'a>>;
109
110    fn put_block_erased(
111        &self,
112        block: CoreMemoryBlock,
113    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>>;
114
115    fn append_block_erased<'a>(
116        &'a self,
117        label: &'a str,
118        text: &'a str,
119    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
120
121    fn remove_block_erased<'a>(
122        &'a self,
123        label: &'a str,
124    ) -> Pin<Box<dyn Future<Output = Result<bool>> + Send + 'a>>;
125
126    fn render_erased(&self) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>>;
127}
128
129impl<T: CoreMemory> ErasedCoreMemory for T {
130    fn blocks_erased(
131        &self,
132    ) -> Pin<Box<dyn Future<Output = Result<Vec<CoreMemoryBlock>>> + Send + '_>> {
133        Box::pin(self.blocks())
134    }
135
136    fn get_block_erased<'a>(
137        &'a self,
138        label: &'a str,
139    ) -> Pin<Box<dyn Future<Output = Result<Option<CoreMemoryBlock>>> + Send + 'a>> {
140        Box::pin(self.get_block(label))
141    }
142
143    fn put_block_erased(
144        &self,
145        block: CoreMemoryBlock,
146    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
147        Box::pin(self.put_block(block))
148    }
149
150    fn append_block_erased<'a>(
151        &'a self,
152        label: &'a str,
153        text: &'a str,
154    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
155        Box::pin(self.append_block(label, text))
156    }
157
158    fn remove_block_erased<'a>(
159        &'a self,
160        label: &'a str,
161    ) -> Pin<Box<dyn Future<Output = Result<bool>> + Send + 'a>> {
162        Box::pin(self.remove_block(label))
163    }
164
165    fn render_erased(&self) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
166        Box::pin(self.render())
167    }
168}
169
170/// Shared ownership of core memory via `Arc<dyn ErasedCoreMemory>`.
171pub type SharedCoreMemory = Arc<dyn ErasedCoreMemory>;
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use std::sync::Mutex;
177
178    struct VecCoreMemory(Mutex<Vec<CoreMemoryBlock>>);
179
180    impl CoreMemory for VecCoreMemory {
181        async fn blocks(&self) -> Result<Vec<CoreMemoryBlock>> {
182            Ok(self.0.lock().unwrap().clone())
183        }
184
185        async fn put_block(&self, block: CoreMemoryBlock) -> Result<()> {
186            if let Some(limit) = block.limit
187                && block.value.chars().count() > limit
188            {
189                return Err(crate::error::DaimonError::Other(format!(
190                    "block '{}' exceeds limit of {limit} characters",
191                    block.label
192                )));
193            }
194            let mut blocks = self.0.lock().unwrap();
195            if let Some(existing) = blocks.iter_mut().find(|b| b.label == block.label) {
196                *existing = block;
197            } else {
198                blocks.push(block);
199            }
200            Ok(())
201        }
202
203        async fn append_block(&self, label: &str, text: &str) -> Result<()> {
204            let mut blocks = self.0.lock().unwrap();
205            if let Some(existing) = blocks.iter_mut().find(|b| b.label == label) {
206                let candidate = format!("{}{}", existing.value, text);
207                if let Some(limit) = existing.limit
208                    && candidate.chars().count() > limit
209                {
210                    return Err(crate::error::DaimonError::Other(format!(
211                        "block '{label}' exceeds limit of {limit} characters"
212                    )));
213                }
214                existing.value = candidate;
215            } else {
216                blocks.push(CoreMemoryBlock::new(label, text));
217            }
218            Ok(())
219        }
220
221        async fn remove_block(&self, label: &str) -> Result<bool> {
222            let mut blocks = self.0.lock().unwrap();
223            let before = blocks.len();
224            blocks.retain(|b| b.label != label);
225            Ok(blocks.len() != before)
226        }
227    }
228
229    #[tokio::test]
230    async fn core_memory_is_implementable_from_core_alone() {
231        let mem = VecCoreMemory(Mutex::new(Vec::new()));
232        mem.put_block(CoreMemoryBlock::new("persona", "helpful assistant"))
233            .await
234            .unwrap();
235        mem.append_block("persona", " who is concise")
236            .await
237            .unwrap();
238
239        let block = mem.get_block("persona").await.unwrap().unwrap();
240        assert_eq!(block.value, "helpful assistant who is concise");
241
242        let rendered = mem.render().await.unwrap();
243        assert_eq!(rendered, "## persona\nhelpful assistant who is concise");
244
245        assert!(mem.remove_block("persona").await.unwrap());
246        assert!(mem.get_block("persona").await.unwrap().is_none());
247
248        let shared: SharedCoreMemory = Arc::new(VecCoreMemory(Mutex::new(Vec::new())));
249        shared
250            .put_block_erased(CoreMemoryBlock::new("user", "likes Rust"))
251            .await
252            .unwrap();
253        assert_eq!(shared.blocks_erased().await.unwrap().len(), 1);
254        assert_eq!(shared.render_erased().await.unwrap(), "## user\nlikes Rust");
255    }
256
257    #[tokio::test]
258    async fn put_block_rejects_over_limit_value() {
259        let mem = VecCoreMemory(Mutex::new(Vec::new()));
260        let err = mem
261            .put_block(CoreMemoryBlock::new("persona", "way too long").with_limit(4))
262            .await
263            .unwrap_err();
264        assert!(err.to_string().contains("exceeds limit"));
265    }
266}