agentrs_memory/
sliding_window.rs1use std::collections::VecDeque;
2
3use async_trait::async_trait;
4
5use agentrs_core::{Memory, Message, Result, Role};
6
7use crate::SearchableMemory;
8
9#[derive(Debug, Clone)]
11pub struct SlidingWindowMemory {
12 system_message: Option<Message>,
13 messages: VecDeque<Message>,
14 window_size: usize,
15}
16
17impl SlidingWindowMemory {
18 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}