Skip to main content

potato_agent/agents/memory/
windowed.rs

1use super::{Memory, MemoryTurn};
2use potato_type::prompt::MessageNum;
3use std::collections::VecDeque;
4
5/// Sliding-window memory that keeps only the last `capacity` turns.
6#[derive(Debug)]
7pub struct WindowedMemory {
8    capacity: usize,
9    turns: VecDeque<MemoryTurn>,
10}
11
12impl WindowedMemory {
13    pub fn new(capacity: usize) -> Self {
14        Self {
15            capacity,
16            turns: VecDeque::with_capacity(capacity),
17        }
18    }
19}
20
21impl Memory for WindowedMemory {
22    fn push_turn(&mut self, turn: MemoryTurn) {
23        if self.capacity == 0 {
24            return;
25        }
26        while self.turns.len() >= self.capacity {
27            self.turns.pop_front();
28        }
29        self.turns.push_back(turn);
30    }
31
32    fn messages(&self) -> Vec<MessageNum> {
33        let mut msgs = Vec::with_capacity(self.turns.len() * 2);
34        for turn in &self.turns {
35            msgs.push(turn.user.clone());
36            msgs.push(turn.assistant.clone());
37        }
38        msgs
39    }
40
41    fn clear(&mut self) {
42        self.turns.clear();
43    }
44
45    fn len(&self) -> usize {
46        self.turns.len()
47    }
48}