potato_agent/agents/memory/
windowed.rs1use super::{Memory, MemoryTurn};
2use potato_type::prompt::MessageNum;
3use std::collections::VecDeque;
4
5#[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}