lc_memory/
summary_buffer.rs

1// lc-memory/src/summary_buffer.rs
2//! Conversation Summary Buffer Memory
3//!
4//! Combines summary and full conversation, balancing token consumption and conversation quality.
5
6use async_trait::async_trait;
7use serde_json::Value;
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use super::base::{BaseChatMemory, BaseMemory, ChatMessageHistory, MemoryError};
12use lc_core::language_models::BaseChatModel;
13use lc_core::language_models::LLMResult;
14use lc_core::runnables::Runnable;
15use lc_core::token_counter::{CharRatioCounter, TiktokenCounter, TokenCounter};
16use lc_prompts::PromptTemplate;
17use lc_schema::Message;
18
19const DEFAULT_SUMMARY_PROMPT: &str =
20    "Progressively summarize the conversation, adding new content to the previous summary.
21
22Current summary:
23{summary}
24
25New lines of conversation:
26{new_lines}
27
28New summary:";
29
30/// Conversation Summary Buffer Memory
31///
32/// Combines summary and full conversation:
33/// - Keeps the last k rounds of full conversation (ensuring fluency)
34/// - Summarizes older conversations (saving tokens)
35///
36/// # Example
37/// ```ignore
38/// use lc_memory::ConversationSummaryBufferMemory;
39/// use lc_providers::OpenAIChat;
40///
41/// let llm = OpenAIChat::new(config);
42/// let memory = ConversationSummaryBufferMemory::new(llm, 5); // Keep last 5 rounds
43///
44/// // After 20 rounds:
45/// // - First 15 rounds -> summary
46/// // - Last 5 rounds -> full conversation
47/// ```
48pub struct ConversationSummaryBufferMemory<M: BaseChatModel> {
49    llm: M,
50
51    /// M67: Removed `Mutex<String>` - &mut self already guarantees exclusive access
52    buffer: String,
53    chat_memory: ChatMessageHistory,
54
55    max_token_limit: usize,
56
57    /// P1-2: 可插拔 token 计数器。默认 `TiktokenCounter`(与 `ContextWindow`
58    /// 同口径,BPE 预算语义统一);可注入 `CharRatioCounter` 保留零依赖快路径。
59    counter: Arc<dyn TokenCounter>,
60
61    input_key: String,
62    output_key: String,
63    memory_key: String,
64
65    summary_prompt: String,
66    return_messages: bool,
67
68    /// P2-4: 最近一次摘要 LLM 失败的原因;成功总结或 `clear()` 后清空。
69    /// 摘要失败时保留旧摘要与原始消息,不清空 `chat_memory`,下轮 prune 重试。
70    last_summary_error: Option<String>,
71}
72
73impl<M: BaseChatModel> ConversationSummaryBufferMemory<M> {
74    pub fn new(llm: M, max_token_limit: usize) -> Self {
75        Self {
76            llm,
77            buffer: String::new(),
78            chat_memory: ChatMessageHistory::new(),
79            max_token_limit,
80            counter: Self::default_token_counter(),
81            input_key: "input".to_string(),
82            output_key: "output".to_string(),
83            memory_key: "history".to_string(),
84            summary_prompt: DEFAULT_SUMMARY_PROMPT.to_string(),
85            return_messages: false,
86            last_summary_error: None,
87        }
88    }
89
90    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
91        self.input_key = key.into();
92        self
93    }
94
95    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
96        self.output_key = key.into();
97        self
98    }
99
100    pub fn with_memory_key(mut self, key: impl Into<String>) -> Self {
101        self.memory_key = key.into();
102        self
103    }
104
105    pub fn with_summary_prompt(mut self, prompt: impl Into<String>) -> Self {
106        self.summary_prompt = prompt.into();
107        self
108    }
109
110    pub fn with_return_messages(mut self, return_messages: bool) -> Self {
111        self.return_messages = return_messages;
112        self
113    }
114
115    /// 注入自定义 token 计数器。
116    ///
117    /// 默认 `TiktokenCounter`(BPE 口径,与 `ContextWindow` 一致);需要零依赖
118    /// 快路径时注入 `CharRatioCounter::new(4)`。
119    pub fn with_counter(mut self, counter: Arc<dyn TokenCounter>) -> Self {
120        self.counter = counter;
121        self
122    }
123
124    /// P1-3: 从持久化存储回灌摘要状态,保证续写会话摘要链连续。
125    pub fn set_summary(&mut self, summary: String) {
126        self.buffer = summary;
127    }
128
129    /// P1-3: 设置 token 预算(持久化 config 的单一来源)。
130    pub fn set_max_token_limit(&mut self, max_token_limit: usize) {
131        self.max_token_limit = max_token_limit;
132    }
133
134    pub fn chat_memory(&self) -> &ChatMessageHistory {
135        &self.chat_memory
136    }
137
138    pub fn chat_memory_mut(&mut self) -> &mut ChatMessageHistory {
139        &mut self.chat_memory
140    }
141
142    pub fn max_token_limit(&self) -> usize {
143        self.max_token_limit
144    }
145
146    pub async fn buffer(&self) -> String {
147        self.buffer.clone()
148    }
149
150    /// P2-4: 最近一次摘要失败的原因(无失败则 `None`)。
151    pub fn last_summary_error(&self) -> Option<&str> {
152        self.last_summary_error.as_deref()
153    }
154
155    /// P1-2: 估算文本 token 数,委托给可插拔计数器(默认 BPE 口径)。
156    fn estimate_tokens(&self, text: &str) -> usize {
157        self.counter.count_tokens(text) as usize
158    }
159
160    fn prune_messages(&self, messages: &[Message]) -> Vec<Message> {
161        let total_tokens = messages
162            .iter()
163            .map(|m| self.estimate_tokens(&m.content))
164            .sum::<usize>();
165
166        if total_tokens <= self.max_token_limit {
167            return messages.to_vec();
168        }
169
170        let mut kept_messages = Vec::new();
171        let mut current_tokens = 0;
172
173        for msg in messages.iter().rev() {
174            let msg_tokens = self.estimate_tokens(&msg.content);
175            if current_tokens + msg_tokens <= self.max_token_limit {
176                kept_messages.push(msg.clone());
177                current_tokens += msg_tokens;
178            } else {
179                break;
180            }
181        }
182
183        kept_messages.reverse();
184        kept_messages
185    }
186
187    /// 默认 token 计数器:优先 `TiktokenCounter`(BPE 口径,与 `ContextWindow` 一致);
188    /// tiktoken 模型加载失败(离线/缺模型)时优雅降级为字符比估算,
189    /// 使 `new()` 保持不可失败的签名。
190    fn default_token_counter() -> Arc<dyn TokenCounter> {
191        TiktokenCounter::new()
192            .map(|c| Arc::new(c) as Arc<dyn TokenCounter>)
193            .unwrap_or_else(|_| Arc::new(CharRatioCounter::new(4)) as Arc<dyn TokenCounter>)
194    }
195
196    async fn predict_new_summary(&self, new_lines: &str) -> Result<String, MemoryError> {
197        let buffer = self.buffer.clone();
198
199        let prompt = {
200            let template = PromptTemplate::new(&self.summary_prompt);
201            let mut vars: std::collections::HashMap<&str, &str> = std::collections::HashMap::new();
202            vars.insert("summary", buffer.as_str());
203            vars.insert("new_lines", new_lines);
204            template
205                .format(&vars)
206                .unwrap_or_else(|_| self.summary_prompt.clone())
207        };
208
209        let messages = vec![Message::human(&prompt)];
210
211        let result =
212            self.llm.invoke(messages, None).await.map_err(|e| {
213                MemoryError::SaveError(format!("LLM summary generation failed: {}", e))
214            })?;
215
216        Ok(result.content)
217    }
218}
219
220#[async_trait]
221impl<M: BaseChatModel + Send + Sync + 'static> BaseMemory for ConversationSummaryBufferMemory<M>
222where
223    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
224{
225    fn memory_variables(&self) -> Vec<&str> {
226        vec![&self.memory_key]
227    }
228
229    async fn load_memory_variables(
230        &self,
231        _inputs: &HashMap<String, String>,
232    ) -> Result<HashMap<String, Value>, MemoryError> {
233        let mut result = HashMap::new();
234
235        let buffer = self.buffer.clone();
236        let messages = self.chat_memory.messages();
237        let pruned = self.prune_messages(messages);
238
239        if self.return_messages {
240            let mut all_messages = Vec::new();
241
242            if !buffer.is_empty() {
243                all_messages.push(Message::system(&buffer));
244            }
245
246            all_messages.extend(pruned);
247
248            let messages_value: Vec<Value> = all_messages
249                .iter()
250                .map(|m| serde_json::to_value(m).unwrap_or(Value::Null))
251                .collect();
252
253            result.insert(self.memory_key.clone(), Value::Array(messages_value));
254        } else {
255            let mut history = String::new();
256
257            if !buffer.is_empty() {
258                history.push_str(&format!("Summary: {}\n\n", buffer));
259            }
260
261            for msg in &pruned {
262                let role = match msg.message_type {
263                    lc_schema::MessageType::Human => "Human",
264                    lc_schema::MessageType::AI => "AI",
265                    lc_schema::MessageType::System => "System",
266                    lc_schema::MessageType::Tool { .. } => "Tool",
267                };
268                history.push_str(&format!("{}: {}\n", role, msg.content));
269            }
270
271            result.insert(self.memory_key.clone(), Value::String(history));
272        }
273
274        Ok(result)
275    }
276
277    async fn save_context(
278        &mut self,
279        inputs: &HashMap<String, String>,
280        outputs: &HashMap<String, String>,
281    ) -> Result<(), MemoryError> {
282        // P1-1: 与 Buffer/Window 一致——缺失 key 返回 SaveError,不再静默用空串
283        // 存空消息(否则会对空行做无意义摘要,白烧一次 LLM 调用)。
284        let input = inputs.get(&self.input_key).ok_or_else(|| {
285            MemoryError::SaveError(format!("Missing input key '{}'", self.input_key))
286        })?;
287        let output = outputs.get(&self.output_key).ok_or_else(|| {
288            MemoryError::SaveError(format!("Missing output key '{}'", self.output_key))
289        })?;
290
291        self.chat_memory.add_user_message(input);
292        self.chat_memory.add_ai_message(output);
293
294        let messages = self.chat_memory.messages();
295        let total_tokens = messages
296            .iter()
297            .map(|m| self.estimate_tokens(&m.content))
298            .sum::<usize>();
299
300        if total_tokens > self.max_token_limit {
301            let pruned = self.prune_messages(messages);
302
303            let pruned_count = pruned.len();
304
305            if messages.len() > pruned_count {
306                let messages_to_summarize: Vec<&Message> = messages
307                    .iter()
308                    .take(messages.len() - pruned_count)
309                    .collect();
310
311                if !messages_to_summarize.is_empty() {
312                    let new_lines: String = messages_to_summarize
313                        .iter()
314                        .map(|m| {
315                            let role = match m.message_type {
316                                lc_schema::MessageType::Human => "Human",
317                                lc_schema::MessageType::AI => "AI",
318                                lc_schema::MessageType::System => "System",
319                                lc_schema::MessageType::Tool { .. } => "Tool",
320                            };
321                            format!("{}: {}", role, m.content)
322                        })
323                        .collect::<Vec<_>>()
324                        .join("\n");
325
326                    // P2-4: 摘要失败时保留旧摘要、不清空 chat_memory(原始消息
327                    // 留下,下轮 prune 会再次尝试总结);错误记录到 last_summary_error
328                    // 供上层观察,不冒泡打断链。
329                    match self.predict_new_summary(&new_lines).await {
330                        Ok(new_summary) => {
331                            self.buffer = new_summary;
332                            self.last_summary_error = None;
333
334                            self.chat_memory.clear();
335                            for msg in pruned {
336                                if matches!(msg.message_type, lc_schema::MessageType::Human) {
337                                    self.chat_memory.add_user_message(&msg.content);
338                                } else if matches!(msg.message_type, lc_schema::MessageType::AI) {
339                                    self.chat_memory.add_ai_message(&msg.content);
340                                } else if matches!(msg.message_type, lc_schema::MessageType::System)
341                                {
342                                    // H28: Preserve System messages during pruning
343                                    self.chat_memory.add_system_message(&msg.content);
344                                }
345                            }
346                        }
347                        Err(e) => {
348                            self.last_summary_error = Some(e.to_string());
349                            log::warn!(
350                                "ConversationSummaryBufferMemory 摘要失败,保留旧摘要与原始消息待下轮重试: {}",
351                                e
352                            );
353                        }
354                    }
355                }
356            }
357        }
358
359        Ok(())
360    }
361
362    async fn clear(&mut self) -> Result<(), MemoryError> {
363        self.buffer = String::new();
364        self.chat_memory.clear();
365        self.last_summary_error = None;
366        Ok(())
367    }
368}
369
370/// P0-1: `ConversationSummaryBufferMemory` 实现 `BaseChatMemory`。
371impl<M: BaseChatModel + Send + Sync + 'static> BaseChatMemory for ConversationSummaryBufferMemory<M>
372where
373    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
374{
375    fn messages(&self) -> &[Message] {
376        self.chat_memory.messages()
377    }
378
379    fn add_message(&mut self, message: Message) {
380        self.chat_memory.add_message(message);
381    }
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387    use crate::test_support::MockLlm;
388    use lc_providers::{OpenAIChat, OpenAIConfig};
389
390    fn create_test_config() -> OpenAIConfig {
391        OpenAIConfig::default()
392    }
393
394    #[test]
395    fn test_new() {
396        let llm = OpenAIChat::new(create_test_config());
397        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
398            ConversationSummaryBufferMemory::new(llm, 1000);
399
400        assert_eq!(memory.memory_variables(), vec!["history"]);
401        assert_eq!(memory.max_token_limit(), 1000);
402    }
403
404    #[test]
405    fn test_with_options() {
406        let llm = OpenAIChat::new(create_test_config());
407        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
408            ConversationSummaryBufferMemory::new(llm, 500)
409                .with_input_key("question")
410                .with_output_key("answer")
411                .with_memory_key("context")
412                .with_return_messages(true);
413
414        assert_eq!(memory.input_key, "question");
415        assert_eq!(memory.output_key, "answer");
416        assert_eq!(memory.memory_key, "context");
417        assert!(memory.return_messages);
418    }
419
420    #[test]
421    fn test_estimate_tokens_uses_default_counter() {
422        // 默认 TiktokenCounter(BPE 口径);离线时降级 CharRatioCounter。
423        // 两种实现下,较长文本的估算 token 数都严格大于较短文本。
424        let llm = OpenAIChat::new(create_test_config());
425        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
426            ConversationSummaryBufferMemory::new(llm, 1000);
427
428        let text1 = "Hello";
429        let text2 = "Hello World";
430        let text3 = "This is some Chinese text";
431
432        assert!(memory.estimate_tokens(text1) > 0);
433        assert!(memory.estimate_tokens(text2) > memory.estimate_tokens(text1));
434        assert!(memory.estimate_tokens(text3) > 0);
435    }
436
437    #[test]
438    fn test_with_counter_injection() {
439        // 注入 CharRatioCounter(ratio=4):8 个字符估算 2 token,可复现、不依赖 tiktoken。
440        let llm = OpenAIChat::new(create_test_config());
441        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
442            ConversationSummaryBufferMemory::new(llm, 1000)
443                .with_counter(std::sync::Arc::new(CharRatioCounter::new(4)));
444
445        assert_eq!(memory.estimate_tokens("abcdefgh"), 2);
446    }
447
448    #[tokio::test]
449    async fn test_set_summary_and_token_limit() {
450        // P1-3: 持久化回灌摘要 + 预算单一来源。
451        let llm = OpenAIChat::new(create_test_config());
452        let mut memory: ConversationSummaryBufferMemory<OpenAIChat> =
453            ConversationSummaryBufferMemory::new(llm, 1000);
454
455        memory.set_summary("previous summary".to_string());
456        assert_eq!(memory.buffer().await, "previous summary");
457
458        memory.set_max_token_limit(500);
459        assert_eq!(memory.max_token_limit(), 500);
460    }
461
462    #[test]
463    fn test_prune_messages_within_limit() {
464        let llm = OpenAIChat::new(create_test_config());
465        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
466            ConversationSummaryBufferMemory::new(llm, 1000);
467
468        let messages = vec![
469            Message::human("Short message 1"),
470            Message::ai("Short reply 1"),
471        ];
472
473        let pruned = memory.prune_messages(&messages);
474
475        assert_eq!(pruned.len(), 2);
476    }
477
478    #[tokio::test]
479    async fn test_buffer_initial_empty() {
480        let llm = OpenAIChat::new(create_test_config());
481        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
482            ConversationSummaryBufferMemory::new(llm, 1000);
483
484        let buffer = memory.buffer().await;
485        assert!(buffer.is_empty());
486    }
487
488    #[tokio::test]
489    async fn test_load_memory_variables_empty() {
490        let llm = OpenAIChat::new(create_test_config());
491        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
492            ConversationSummaryBufferMemory::new(llm, 1000);
493
494        let vars = memory.load_memory_variables(&HashMap::new()).await.unwrap();
495        let history = vars.get("history").unwrap().as_str().unwrap();
496
497        assert!(history.is_empty());
498    }
499
500    #[tokio::test]
501    async fn test_clear() {
502        let llm = OpenAIChat::new(create_test_config());
503        let mut memory: ConversationSummaryBufferMemory<OpenAIChat> =
504            ConversationSummaryBufferMemory::new(llm, 1000);
505
506        memory.chat_memory.add_user_message("test");
507        memory.chat_memory.add_ai_message("reply");
508
509        memory.buffer = "Test summary".to_string();
510
511        memory.clear().await.unwrap();
512
513        assert!(memory.buffer().await.is_empty());
514        assert_eq!(memory.chat_memory().len(), 0);
515    }
516
517    /// P2-4: 剪枝触发摘要、但摘要 LLM 失败时——保留旧摘要、不清空 chat_memory、
518    /// 记录错误且不冒泡;下轮成功总结后摘要生效、错误清空。
519    #[tokio::test]
520    async fn test_prune_summary_failure_keeps_messages_and_retries() {
521        // MockLlm 按 LIFO 消费:第一次剪枝总结失败,第二次成功。
522        let llm = MockLlm::new(vec![
523            Ok("summary-b".to_string()),
524            Err("summarizer down".to_string()),
525        ]);
526        // CharRatioCounter 保证 token 估算可复现(不依赖 tiktoken 在线)。
527        let mut memory: ConversationSummaryBufferMemory<MockLlm> =
528            ConversationSummaryBufferMemory::new(llm, 5)
529                .with_counter(std::sync::Arc::new(CharRatioCounter::new(4)));
530
531        let long_input =
532            "这是一段足够长的中文消息,用来确保本轮消息总 token 数超过预算并触发剪枝总结逻辑";
533        let inputs = HashMap::from([("input".to_string(), long_input.to_string())]);
534        let outputs = HashMap::from([("output".to_string(), long_input.to_string())]);
535
536        // 第一轮:总 token 超限 -> 触发剪枝总结 -> LLM 失败
537        memory.save_context(&inputs, &outputs).await.unwrap();
538        assert!(memory.buffer().await.is_empty(), "失败时不应覆盖旧摘要");
539        assert!(memory
540            .last_summary_error()
541            .unwrap()
542            .contains("summarizer down"));
543        // 失败时不清空原始消息,下轮 prune 才能重试总结
544        assert_eq!(memory.chat_memory().len(), 2);
545
546        // 第二轮:再次触发剪枝总结 -> 成功,摘要生效、错误清空
547        memory.save_context(&inputs, &outputs).await.unwrap();
548        assert_eq!(memory.buffer().await, "summary-b");
549        assert!(memory.last_summary_error().is_none());
550    }
551}
lc_memory/summary_buffer.rs

lc_memory/
summary_buffer.rs
