molo_agent/memory/in_memory.rs
1//! In-memory implementation: stores all messages verbatim.
2//!
3//! [`InMemoryMemory`] is the simplest [`Memory`] implementation: `record`
4//! appends, `context` returns a full clone, with no trimming or summarization —
5//! suitable for small conversations that need no budget control.
6
7use super::{Memory, MemoryError};
8use crate::message::Message;
9
10/// In-memory implementation: stores all messages verbatim.
11///
12/// The simplest [`Memory`] implementation: all messages are kept in memory with
13/// no trimming, no summarization, and no persistence; `context()` returns a
14/// full clone. `Default` constructs an empty session.
15///
16/// # Comparison
17///
18/// Compared to [`WindowMemory`](crate::memory::WindowMemory): the latter trims
19/// to a recent window on retrieval, suiting long conversations (the model's
20/// context window is finite); this implementation has no budget control —
21/// messages only grow — suiting small conversations that need the full history.
22///
23/// # Example
24///
25/// ```rust
26/// # extern crate molo_agent as molo;
27/// # #[tokio::main]
28/// # async fn main() -> Result<(), molo::memory::MemoryError> {
29/// use molo::memory::{InMemoryMemory, Memory};
30///
31/// let mut memory = InMemoryMemory::default();
32/// memory.record(molo::Message::system("setup")).await?;
33/// memory.record(molo::Message::user("hello")).await?;
34/// memory.record(molo::Message::assistant("hello!")).await?;
35///
36/// let context = memory.context().await?;
37/// assert_eq!(context.len(), 3);
38/// assert_eq!(context[0], molo::Message::System("setup".into()));
39/// # Ok(())
40/// # }
41/// ```
42#[derive(Debug, Default)]
43pub struct InMemoryMemory {
44 messages: Vec<Message>,
45}
46
47#[async_trait::async_trait]
48impl Memory for InMemoryMemory {
49 async fn record(&mut self, message: Message) -> Result<(), MemoryError> {
50 self.messages.push(message);
51 Ok(())
52 }
53
54 async fn context(&self) -> Result<Vec<Message>, MemoryError> {
55 Ok(self.messages.clone())
56 }
57}
58
59#[cfg(test)]
60mod tests {
61 use super::*;
62
63 #[tokio::test]
64 async fn record_preserves_order() {
65 let mut memory = InMemoryMemory::default();
66 memory.record(Message::system("setup")).await.unwrap();
67 memory.record(Message::user("hello")).await.unwrap();
68 memory.record(Message::assistant("hello!")).await.unwrap();
69
70 let context = memory.context().await.unwrap();
71 assert_eq!(context.len(), 3);
72 assert_eq!(context[0], Message::System("setup".into()));
73 assert_eq!(context[1], Message::user("hello"));
74 assert_eq!(context[2], Message::assistant("hello!"));
75 }
76
77 #[tokio::test]
78 async fn context_is_a_copy() {
79 let mut memory = InMemoryMemory::default();
80 memory.record(Message::user("a")).await.unwrap();
81
82 // Mutating the returned context does not affect Memory's internals.
83 let mut context = memory.context().await.unwrap();
84 context.push(Message::assistant("b"));
85
86 assert_eq!(memory.context().await.unwrap().len(), 1);
87 }
88}