lellm_agent/memory/
short_term.rs

1//! 短期记忆 — 基于 VecDeque 的对话历史窗口。
2
3use std::collections::VecDeque;
4
5use lellm_core::Message;
6
7/// 短期记忆 — 环形缓冲区，默认 200 条容量。
8pub struct ShortTermMemory {
9    messages: VecDeque<Message>,
10    capacity: usize,
11}
12
13impl ShortTermMemory {
14    pub fn new(capacity: usize) -> Self {
15        Self {
16            messages: VecDeque::with_capacity(capacity),
17            capacity,
18        }
19    }
20
21    pub fn push(&mut self, msg: Message) {
22        if self.messages.len() >= self.capacity {
23            self.messages.pop_front();
24        }
25        self.messages.push_back(msg);
26    }
27
28    /// 获取最近 n 条消息
29    pub fn recent(&self, n: usize) -> Vec<Message> {
30        if n >= self.messages.len() {
31            self.messages.iter().cloned().collect()
32        } else {
33            self.messages
34                .iter()
35                .skip(self.messages.len() - n)
36                .cloned()
37                .collect()
38        }
39    }
40
41    /// 获取所有消息
42    pub fn messages(&self) -> Vec<Message> {
43        self.messages.iter().cloned().collect()
44    }
45
46    pub fn len(&self) -> usize {
47        self.messages.len()
48    }
49
50    pub fn is_empty(&self) -> bool {
51        self.messages.is_empty()
52    }
53
54    pub fn clear(&mut self) {
55        self.messages.clear();
56    }
57}
58
59impl Default for ShortTermMemory {
60    fn default() -> Self {
61        Self::new(200)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn test_short_term_memory_push_and_messages() {
71        let mut memory = ShortTermMemory::new(10);
72        memory.push(Message::User {
73            content: lellm_core::text_block("hello".to_string()),
74        });
75        assert_eq!(memory.len(), 1);
76        assert!(!memory.is_empty());
77    }
78
79    #[test]
80    fn test_short_term_memory_capacity() {
81        let mut memory = ShortTermMemory::new(2);
82        for i in 0..5 {
83            memory.push(Message::User {
84                content: lellm_core::text_block(format!("msg{}", i)),
85            });
86        }
87        assert_eq!(memory.len(), 2);
88        let recent = memory.recent(2);
89        assert_eq!(recent.len(), 2);
90        assert_eq!(recent[0].extract_text(), "msg3");
91        assert_eq!(recent[1].extract_text(), "msg4");
92    }
93
94    #[test]
95    fn test_short_term_memory_clear() {
96        let mut memory = ShortTermMemory::new(10);
97        memory.push(Message::User {
98            content: lellm_core::text_block("hello".to_string()),
99        });
100        memory.clear();
101        assert!(memory.is_empty());
102    }
103
104    #[test]
105    fn test_short_term_memory_default_capacity() {
106        let mut memory = ShortTermMemory::default();
107        for i in 0..205 {
108            memory.push(Message::User {
109                content: lellm_core::text_block(format!("msg{}", i)),
110            });
111        }
112        assert_eq!(memory.len(), 200);
113    }
114}
lellm_agent/memory/short_term.rs

lellm_agent/memory/
short_term.rs
