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

lc_memory/
summary_buffer.rs
