daimon_core/
core_memory.rs1use std::future::Future;
11use std::pin::Pin;
12use std::sync::Arc;
13
14use crate::error::Result;
15
16#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct CoreMemoryBlock {
23 pub label: String,
25 pub value: String,
27 pub limit: Option<usize>,
29}
30
31impl CoreMemoryBlock {
32 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 pub fn with_limit(mut self, limit: usize) -> Self {
43 self.limit = Some(limit);
44 self
45 }
46}
47
48pub trait CoreMemory: Send + Sync {
56 fn blocks(&self) -> impl Future<Output = Result<Vec<CoreMemoryBlock>>> + Send;
58
59 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 fn put_block(&self, block: CoreMemoryBlock) -> impl Future<Output = Result<()>> + Send;
70
71 fn append_block(&self, label: &str, text: &str) -> impl Future<Output = Result<()>> + Send;
75
76 fn remove_block(&self, label: &str) -> impl Future<Output = Result<bool>> + Send;
78
79 fn render(&self) -> impl Future<Output = Result<String>> + Send {
83 async move { Ok(render_blocks(&self.blocks().await?)) }
84 }
85}
86
87pub 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
98pub 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
170pub 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}