langchain_rust/memory/
window_buffer.rs

1use std::sync::Arc;
2
3use tokio::sync::Mutex;
4
5use crate::schemas::{memory::BaseMemory, messages::Message};
6
7pub struct WindowBufferMemory {
8    window_size: usize,
9    messages: Vec<Message>,
10}
11
12impl Default for WindowBufferMemory {
13    fn default() -> Self {
14        Self::new(10)
15    }
16}
17
18impl WindowBufferMemory {
19    pub fn new(window_size: usize) -> Self {
20        Self {
21            messages: Vec::new(),
22            window_size,
23        }
24    }
25}
26
27impl Into<Arc<dyn BaseMemory>> for WindowBufferMemory {
28    fn into(self) -> Arc<dyn BaseMemory> {
29        Arc::new(self)
30    }
31}
32
33impl Into<Arc<Mutex<dyn BaseMemory>>> for WindowBufferMemory {
34    fn into(self) -> Arc<Mutex<dyn BaseMemory>> {
35        Arc::new(Mutex::new(self))
36    }
37}
38
39impl BaseMemory for WindowBufferMemory {
40    fn messages(&self) -> Vec<Message> {
41        self.messages.clone()
42    }
43    fn add_message(&mut self, message: Message) {
44        if self.messages.len() >= self.window_size {
45            self.messages.remove(0);
46        }
47        self.messages.push(message);
48    }
49    fn clear(&mut self) {
50        self.messages.clear();
51    }
52}