langchain_rust/memory/
window_buffer.rs1use 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}