Skip to main content

agentrs_memory/
sliding_window.rs

1use std::collections::VecDeque;
2
3use async_trait::async_trait;
4
5use agentrs_core::{Memory, Message, Result, Role};
6
7use crate::SearchableMemory;
8
9/// Memory backend that keeps a fixed number of recent non-system messages.
10#[derive(Debug, Clone)]
11pub struct SlidingWindowMemory {
12    system_message: Option<Message>,
13    messages: VecDeque<Message>,
14    window_size: usize,
15}
16
17impl SlidingWindowMemory {
18    /// Creates a new sliding-window memory.
19    pub fn new(window_size: usize) -> Self {
20        Self {
21            system_message: None,
22            messages: VecDeque::new(),
23            window_size,
24        }
25    }
26}
27
28#[async_trait]
29impl Memory for SlidingWindowMemory {
30    async fn store(&mut self, _key: &str, value: Message) -> Result<()> {
31        if matches!(value.role, Role::System) {
32            self.system_message = Some(value);
33            return Ok(());
34        }
35
36        self.messages.push_back(value);
37        while self.messages.len() > self.window_size {
38            self.messages.pop_front();
39        }
40
41        Ok(())
42    }
43
44    async fn retrieve(&self, query: &str, limit: usize) -> Result<Vec<Message>> {
45        let query = query.to_lowercase();
46        Ok(self
47            .messages
48            .iter()
49            .filter(|message| message.text_content().to_lowercase().contains(&query))
50            .take(limit)
51            .cloned()
52            .collect())
53    }
54
55    async fn history(&self) -> Result<Vec<Message>> {
56        let mut history = Vec::with_capacity(self.messages.len() + 1);
57        if let Some(system_message) = &self.system_message {
58            history.push(system_message.clone());
59        }
60        history.extend(self.messages.iter().cloned());
61        Ok(history)
62    }
63
64    async fn clear(&mut self) -> Result<()> {
65        self.messages.clear();
66        Ok(())
67    }
68}
69
70#[async_trait]
71impl SearchableMemory for SlidingWindowMemory {
72    async fn token_count(&self) -> Result<usize> {
73        Ok(self
74            .messages
75            .iter()
76            .map(|message| message.text_content().chars().count() / 4)
77            .sum())
78    }
79}