langchain_rust/memory/
window_buffer.rs

1use std::collections::VecDeque;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use tokio::sync::Mutex;
6
7use crate::schemas::{
8    memory::{BaseMemory, MemoryError},
9    messages::Message,
10};
11
12/// A sliding-window conversation memory that retains the last `window_size` messages.
13///
14/// Internally uses a [`VecDeque`] so that evicting the oldest message on overflow
15/// is O(1) rather than the O(n) `Vec::remove(0)` it replaces.
16pub struct WindowBufferMemory {
17    window_size: usize,
18    messages: VecDeque<Message>,
19}
20
21impl Default for WindowBufferMemory {
22    fn default() -> Self {
23        Self::new(10)
24    }
25}
26
27impl WindowBufferMemory {
28    pub fn new(window_size: usize) -> Self {
29        Self {
30            messages: VecDeque::new(),
31            window_size,
32        }
33    }
34}
35
36impl Into<Arc<dyn BaseMemory>> for WindowBufferMemory {
37    fn into(self) -> Arc<dyn BaseMemory> {
38        Arc::new(self)
39    }
40}
41
42impl Into<Arc<Mutex<dyn BaseMemory>>> for WindowBufferMemory {
43    fn into(self) -> Arc<Mutex<dyn BaseMemory>> {
44        Arc::new(Mutex::new(self))
45    }
46}
47
48#[async_trait]
49impl BaseMemory for WindowBufferMemory {
50    async fn messages(&self) -> Result<Vec<Message>, MemoryError> {
51        Ok(self.messages.iter().cloned().collect())
52    }
53
54    async fn add_message(&mut self, message: Message) -> Result<(), MemoryError> {
55        if self.window_size == 0 {
56            return Ok(());
57        }
58        if self.messages.len() >= self.window_size {
59            self.messages.pop_front();
60        }
61        self.messages.push_back(message);
62        Ok(())
63    }
64
65    async fn clear(&mut self) -> Result<(), MemoryError> {
66        self.messages.clear();
67        Ok(())
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use crate::schemas::MessageType;
75
76    #[tokio::test]
77    async fn keeps_only_last_window_size_messages() {
78        let mut memory = WindowBufferMemory::new(2);
79        memory.add_user_message(&"one").await.unwrap();
80        memory.add_ai_message(&"two").await.unwrap();
81        memory.add_user_message(&"three").await.unwrap();
82
83        let messages = memory.messages().await.unwrap();
84        assert_eq!(messages.len(), 2);
85        assert_eq!(messages[0].content, "two");
86        assert_eq!(messages[0].message_type, MessageType::AIMessage);
87        assert_eq!(messages[1].content, "three");
88        assert_eq!(messages[1].message_type, MessageType::HumanMessage);
89    }
90
91    #[tokio::test]
92    async fn default_window_is_ten() {
93        let mut memory = WindowBufferMemory::default();
94        for i in 0..15 {
95            memory.add_user_message(&i).await.unwrap();
96        }
97        let messages = memory.messages().await.unwrap();
98        assert_eq!(messages.len(), 10);
99        assert_eq!(messages.first().unwrap().content, "5");
100        assert_eq!(messages.last().unwrap().content, "14");
101    }
102
103    #[tokio::test]
104    async fn zero_window_stores_nothing_and_does_not_panic() {
105        let mut memory = WindowBufferMemory::new(0);
106        memory.add_user_message(&"hello").await.unwrap();
107        assert!(memory.messages().await.unwrap().is_empty());
108    }
109
110    #[tokio::test]
111    async fn clear_and_to_string() {
112        let mut memory = WindowBufferMemory::new(5);
113        memory.add_user_message(&"hi").await.unwrap();
114        memory.add_ai_message(&"hello").await.unwrap();
115        assert_eq!(memory.to_string().await.unwrap(), "human: hi\nai: hello");
116        memory.clear().await.unwrap();
117        assert!(memory.messages().await.unwrap().is_empty());
118        assert_eq!(memory.to_string().await.unwrap(), "");
119    }
120}
langchain_rust/memory/window_buffer.rs

langchain_rust/memory/
window_buffer.rs
