rskit_agent/memory/
sliding_window.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use rskit_errors::{AppError, ErrorCode};
5use rskit_llm::types::Message;
6
7use super::Memory;
8
9// ── SlidingWindowMemory ─────────────────────────────────────────────────────
10
11/// Wraps any [`Memory`] and limits the stored history to the last
12/// `max_messages` non-system entries (always preserving a leading system
13/// message in addition to that limit).
14pub struct SlidingWindowMemory {
15    inner: Arc<dyn Memory>,
16    max_messages: usize,
17}
18
19impl SlidingWindowMemory {
20    /// Create a sliding-window wrapper.
21    ///
22    /// `max_messages` is the number of non-system messages to retain and must
23    /// be at least 1.
24    pub fn new(inner: Arc<dyn Memory>, max_messages: usize) -> Result<Self, AppError> {
25        if max_messages == 0 {
26            return Err(AppError::new(
27                ErrorCode::InvalidInput,
28                "max_messages must be at least 1",
29            ));
30        }
31        Ok(Self {
32            inner,
33            max_messages,
34        })
35    }
36
37    /// Trim a message list: keep the system prompt (if present) plus the last
38    /// `max_messages` non-system messages.
39    fn trim(&self, messages: &[Message]) -> Vec<Message> {
40        if messages.len() <= self.max_messages {
41            return messages.to_vec();
42        }
43
44        let mut result = Vec::with_capacity(self.max_messages + 1);
45
46        let has_system = matches!(messages.first(), Some(Message::System(_)));
47        if has_system {
48            result.push(messages[0].clone());
49        }
50
51        let non_system = if has_system { &messages[1..] } else { messages };
52        let start = non_system.len().saturating_sub(self.max_messages);
53        result.extend_from_slice(&non_system[start..]);
54
55        result
56    }
57}
58
59#[async_trait]
60impl Memory for SlidingWindowMemory {
61    async fn load(&self, session_id: &str) -> Result<Vec<Message>, AppError> {
62        let messages = self.inner.load(session_id).await?;
63        Ok(self.trim(&messages))
64    }
65
66    async fn save(&self, session_id: &str, messages: &[Message]) -> Result<(), AppError> {
67        let trimmed = self.trim(messages);
68        self.inner.save(session_id, &trimmed).await
69    }
70
71    async fn append(&self, session_id: &str, messages: &[Message]) -> Result<(), AppError> {
72        self.inner.append(session_id, messages).await?;
73        // Re-load, trim, and persist.
74        let all = self.inner.load(session_id).await?;
75        let trimmed = self.trim(&all);
76        self.inner.save(session_id, &trimmed).await
77    }
78
79    async fn clear(&self, session_id: &str) -> Result<(), AppError> {
80        self.inner.clear(session_id).await
81    }
82}
83
84// ── Tests ───────────────────────────────────────────────────────────────────
rskit_agent/memory/sliding_window.rs

rskit_agent/memory/
sliding_window.rs
