Skip to main content

abu_agent/memory/
slidingwindow.rs

1use std::{collections::VecDeque, convert::Infallible}; 
2use abu_base::chat::ChatMessage;
3use super::Memory;
4
5pub struct SliceWindowMemory {
6    history: VecDeque<ChatMessage>, 
7    window_size: usize,
8}
9
10impl SliceWindowMemory {
11    pub fn new(window_size: usize) -> Self {
12        Self {
13            history: VecDeque::with_capacity(window_size),
14            window_size,
15        }
16    }
17
18    pub fn window_size(&self) -> usize {
19        self.window_size
20    }
21}
22
23impl Memory for SliceWindowMemory {
24    type Error = Infallible;
25
26    async fn add(&mut self, user_input: &str, ai_response: &str) -> Result<(), Self::Error> {
27        while self.history.len() / 2 > self.window_size {
28            self.history.pop_front(); // Pop User
29            self.history.pop_front(); // Pop Assistant
30        }
31        self.history.push_back(ChatMessage::user(user_input));
32        self.history.push_back(ChatMessage::assistant(ai_response, []));
33        Ok(())
34    }
35
36    async fn search(&self, _query: &str) -> Result<Vec<ChatMessage>, Self::Error> {
37        Ok(self.history.iter().cloned().collect())
38    }
39
40    async fn clear(&mut self) -> Result<(), Self::Error> {
41        self.history.clear();
42        Ok(())
43    }
44}