langchain_rust/memory/
simple_memory.rs

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

langchain_rust/memory/
simple_memory.rs
