langchain_rust/memory/
simple_memory.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use tokio::sync::Mutex;
5
6use crate::schemas::{
7    memory::{BaseMemory, MemoryError},
8    messages::Message,
9};
10
11pub struct SimpleMemory {
12    messages: Vec<Message>,
13}
14
15impl SimpleMemory {
16    pub fn new() -> Self {
17        Self {
18            messages: Vec::new(),
19        }
20    }
21}
22
23impl Default for SimpleMemory {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl Into<Arc<dyn BaseMemory>> for SimpleMemory {
30    fn into(self) -> Arc<dyn BaseMemory> {
31        Arc::new(self)
32    }
33}
34
35impl Into<Arc<Mutex<dyn BaseMemory>>> for SimpleMemory {
36    fn into(self) -> Arc<Mutex<dyn BaseMemory>> {
37        Arc::new(Mutex::new(self))
38    }
39}
40
41#[async_trait]
42impl BaseMemory for SimpleMemory {
43    async fn messages(&self) -> Result<Vec<Message>, MemoryError> {
44        Ok(self.messages.clone())
45    }
46
47    async fn add_message(&mut self, message: Message) -> Result<(), MemoryError> {
48        self.messages.push(message);
49        Ok(())
50    }
51
52    async fn clear(&mut self) -> Result<(), MemoryError> {
53        self.messages.clear();
54        Ok(())
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61    use crate::schemas::MessageType;
62
63    #[tokio::test]
64    async fn stores_all_messages_in_order() {
65        let mut memory = SimpleMemory::new();
66        memory.add_user_message(&"question").await.unwrap();
67        memory.add_ai_message(&"answer").await.unwrap();
68        memory
69            .add_message(Message::new_system_message("sys"))
70            .await
71            .unwrap();
72
73        let messages = memory.messages().await.unwrap();
74        let types: Vec<_> = messages.iter().map(|m| m.message_type.clone()).collect();
75        assert_eq!(
76            types,
77            vec![
78                MessageType::HumanMessage,
79                MessageType::AIMessage,
80                MessageType::SystemMessage
81            ]
82        );
83        assert_eq!(
84            memory.to_string().await.unwrap(),
85            "human: question\nai: answer\nsystem: sys"
86        );
87    }
88
89    #[tokio::test]
90    async fn clear_empties_memory() {
91        let mut memory = SimpleMemory::new();
92        memory.add_user_message(&"x").await.unwrap();
93        memory.clear().await.unwrap();
94        assert!(memory.messages().await.unwrap().is_empty());
95    }
96
97    #[tokio::test]
98    async fn converts_into_shared_mutex() {
99        let memory: Arc<Mutex<dyn BaseMemory>> = SimpleMemory::new().into();
100        memory
101            .lock()
102            .await
103            .add_user_message(&"shared")
104            .await
105            .unwrap();
106        assert_eq!(memory.lock().await.messages().await.unwrap().len(), 1);
107    }
108}
langchain_rust/memory/simple_memory.rs

langchain_rust/memory/
simple_memory.rs
