Skip to main content

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}