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, MessageType};
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
195        // 0.22.0 H-M2: the window must not open on an orphaned Tool message.
196        // Pruning keeps the newest suffix, so a kept window can start on a
197        // `Tool` / tool RESULT whose matching assistant `tool_calls` fell
198        // outside the window. A standalone tool message is malformed for
199        // OpenAI/Anthropic (tool message without a preceding assistant
200        // tool_calls → 400). Drop any leading Tool messages until the window
201        // opens on a non-Tool message.
202        while kept_messages
203            .first()
204            .is_some_and(|m| matches!(m.message_type, MessageType::Tool { .. }))
205        {
206            kept_messages.remove(0);
207            // The token budget may now be under the limit; that is fine.
208            // Subsequent turns re-run the pruner as history grows.
209            if kept_messages.is_empty() {
210                break;
211            }
212        }
213
214        kept_messages
215    }
216
217    /// 默认 token 计数器:优先 `TiktokenCounter`(BPE 口径,与 `ContextWindow` 一致);
218    /// tiktoken 模型加载失败(离线/缺模型)时优雅降级为字符比估算,
219    /// 使 `new()` 保持不可失败的签名。
220    fn default_token_counter() -> Arc<dyn TokenCounter> {
221        TiktokenCounter::new()
222            .map(|c| Arc::new(c) as Arc<dyn TokenCounter>)
223            .unwrap_or_else(|_| Arc::new(CharRatioCounter::new(4)) as Arc<dyn TokenCounter>)
224    }
225
226    async fn predict_new_summary(&self, new_lines: &str) -> Result<String, MemoryError> {
227        let buffer = self.buffer.clone();
228
229        let prompt = {
230            let template = PromptTemplate::new(&self.summary_prompt);
231            let mut vars: std::collections::HashMap<&str, &str> = std::collections::HashMap::new();
232            vars.insert("summary", buffer.as_str());
233            vars.insert("new_lines", new_lines);
234            template
235                .format(&vars)
236                .unwrap_or_else(|_| self.summary_prompt.clone())
237        };
238
239        let messages = vec![Message::human(&prompt)];
240
241        let result =
242            self.llm.invoke(messages, None).await.map_err(|e| {
243                MemoryError::SaveError(format!("LLM summary generation failed: {}", e))
244            })?;
245
246        Ok(result.content)
247    }
248}
249
250#[async_trait]
251impl<M: BaseChatModel + Send + Sync + 'static> BaseMemory for ConversationSummaryBufferMemory<M>
252where
253    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
254{
255    fn memory_variables(&self) -> Vec<&str> {
256        vec![&self.memory_key]
257    }
258
259    async fn load_memory_variables(
260        &self,
261        _inputs: &HashMap<String, String>,
262    ) -> Result<HashMap<String, Value>, MemoryError> {
263        let mut result = HashMap::new();
264
265        let buffer = self.buffer.clone();
266        let messages = self.chat_memory.messages();
267        let pruned = self.prune_messages(messages);
268
269        if self.return_messages {
270            let mut all_messages = Vec::new();
271
272            if !buffer.is_empty() {
273                all_messages.push(Message::system(&buffer));
274            }
275
276            all_messages.extend(pruned);
277
278            let messages_value: Vec<Value> = all_messages
279                .iter()
280                .map(|m| serde_json::to_value(m).unwrap_or(Value::Null))
281                .collect();
282
283            result.insert(self.memory_key.clone(), Value::Array(messages_value));
284        } else {
285            let mut history = String::new();
286
287            if !buffer.is_empty() {
288                history.push_str(&format!("Summary: {}\n\n", buffer));
289            }
290
291            for msg in &pruned {
292                let role = match msg.message_type {
293                    lc_schema::MessageType::Human => "Human",
294                    lc_schema::MessageType::AI => "AI",
295                    lc_schema::MessageType::System => "System",
296                    lc_schema::MessageType::Tool { .. } => "Tool",
297                };
298                history.push_str(&format!("{}: {}\n", role, msg.content));
299            }
300
301            result.insert(self.memory_key.clone(), Value::String(history));
302        }
303
304        Ok(result)
305    }
306
307    async fn save_context(
308        &mut self,
309        inputs: &HashMap<String, String>,
310        outputs: &HashMap<String, String>,
311    ) -> Result<(), MemoryError> {
312        // P1-1: 与 Buffer/Window 一致——缺失 key 返回 SaveError,不再静默用空串
313        // 存空消息(否则会对空行做无意义摘要,白烧一次 LLM 调用)。
314        let input = inputs.get(&self.input_key).ok_or_else(|| {
315            MemoryError::SaveError(format!("Missing input key '{}'", self.input_key))
316        })?;
317        let output = outputs.get(&self.output_key).ok_or_else(|| {
318            MemoryError::SaveError(format!("Missing output key '{}'", self.output_key))
319        })?;
320
321        self.chat_memory.add_user_message(input);
322        self.chat_memory.add_ai_message(output);
323
324        let messages = self.chat_memory.messages();
325        let total_tokens = messages
326            .iter()
327            .map(|m| self.estimate_tokens(&m.content))
328            .sum::<usize>();
329
330        if total_tokens > self.max_token_limit {
331            let pruned = self.prune_messages(messages);
332
333            let pruned_count = pruned.len();
334
335            if messages.len() > pruned_count {
336                let messages_to_summarize: Vec<&Message> = messages
337                    .iter()
338                    .take(messages.len() - pruned_count)
339                    .collect();
340
341                if !messages_to_summarize.is_empty() {
342                    let new_lines: String = messages_to_summarize
343                        .iter()
344                        .map(|m| {
345                            let role = match m.message_type {
346                                lc_schema::MessageType::Human => "Human",
347                                lc_schema::MessageType::AI => "AI",
348                                lc_schema::MessageType::System => "System",
349                                lc_schema::MessageType::Tool { .. } => "Tool",
350                            };
351                            format!("{}: {}", role, m.content)
352                        })
353                        .collect::<Vec<_>>()
354                        .join("\n");
355
356                    // P2-4: 摘要失败时保留旧摘要、不清空 chat_memory(原始消息
357                    // 留下,下轮 prune 会再次尝试总结);错误记录到 last_summary_error
358                    // 供上层观察,不冒泡打断链。
359                    match self.predict_new_summary(&new_lines).await {
360                        Ok(new_summary) => {
361                            self.buffer = new_summary;
362                            self.last_summary_error = None;
363
364                            self.chat_memory.clear();
365                            for msg in pruned {
366                                if matches!(msg.message_type, lc_schema::MessageType::Human) {
367                                    self.chat_memory.add_user_message(&msg.content);
368                                } else if matches!(msg.message_type, lc_schema::MessageType::AI) {
369                                    self.chat_memory.add_ai_message(&msg.content);
370                                } else if matches!(msg.message_type, lc_schema::MessageType::System)
371                                {
372                                    // H28: Preserve System messages during pruning
373                                    self.chat_memory.add_system_message(&msg.content);
374                                }
375                            }
376                        }
377                        Err(e) => {
378                            self.last_summary_error = Some(e.to_string());
379                            log::warn!(
380                                "ConversationSummaryBufferMemory summarization failed, keeping old summary and original messages for next retry: {}",
381                                e
382                            );
383                        }
384                    }
385                }
386            }
387        }
388
389        Ok(())
390    }
391
392    async fn clear(&mut self) -> Result<(), MemoryError> {
393        self.buffer = String::new();
394        self.chat_memory.clear();
395        self.last_summary_error = None;
396        Ok(())
397    }
398}
399
400/// P0-1: `ConversationSummaryBufferMemory` 实现 `BaseChatMemory`。
401impl<M: BaseChatModel + Send + Sync + 'static> BaseChatMemory for ConversationSummaryBufferMemory<M>
402where
403    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
404{
405    fn messages(&self) -> &[Message] {
406        self.chat_memory.messages()
407    }
408
409    fn add_message(&mut self, message: Message) {
410        self.chat_memory.add_message(message);
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417    use crate::test_support::MockLlm;
418    use lc_providers::{OpenAIChat, OpenAIConfig};
419
420    fn create_test_config() -> OpenAIConfig {
421        OpenAIConfig::default()
422    }
423
424    #[test]
425    fn test_new() {
426        let llm = OpenAIChat::new(create_test_config());
427        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
428            ConversationSummaryBufferMemory::new(llm, 1000);
429
430        assert_eq!(memory.memory_variables(), vec!["history"]);
431        assert_eq!(memory.max_token_limit(), 1000);
432    }
433
434    #[test]
435    fn test_with_options() {
436        let llm = OpenAIChat::new(create_test_config());
437        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
438            ConversationSummaryBufferMemory::new(llm, 500)
439                .with_input_key("question")
440                .with_output_key("answer")
441                .with_memory_key("context")
442                .with_return_messages(true);
443
444        assert_eq!(memory.input_key, "question");
445        assert_eq!(memory.output_key, "answer");
446        assert_eq!(memory.memory_key, "context");
447        assert!(memory.return_messages);
448    }
449
450    #[test]
451    fn test_estimate_tokens_uses_default_counter() {
452        // 默认 TiktokenCounter(BPE 口径);离线时降级 CharRatioCounter。
453        // 两种实现下,较长文本的估算 token 数都严格大于较短文本。
454        let llm = OpenAIChat::new(create_test_config());
455        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
456            ConversationSummaryBufferMemory::new(llm, 1000);
457
458        let text1 = "Hello";
459        let text2 = "Hello World";
460        let text3 = "This is some Chinese text";
461
462        assert!(memory.estimate_tokens(text1) > 0);
463        assert!(memory.estimate_tokens(text2) > memory.estimate_tokens(text1));
464        assert!(memory.estimate_tokens(text3) > 0);
465    }
466
467    #[test]
468    fn test_with_counter_injection() {
469        // 注入 CharRatioCounter(ratio=4):8 个字符估算 2 token,可复现、不依赖 tiktoken。
470        let llm = OpenAIChat::new(create_test_config());
471        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
472            ConversationSummaryBufferMemory::new(llm, 1000)
473                .with_counter(std::sync::Arc::new(CharRatioCounter::new(4)));
474
475        assert_eq!(memory.estimate_tokens("abcdefgh"), 2);
476    }
477
478    #[tokio::test]
479    async fn test_set_summary_and_token_limit() {
480        // P1-3: 持久化回灌摘要 + 预算单一来源。
481        let llm = OpenAIChat::new(create_test_config());
482        let mut memory: ConversationSummaryBufferMemory<OpenAIChat> =
483            ConversationSummaryBufferMemory::new(llm, 1000);
484
485        memory.set_summary("previous summary".to_string());
486        assert_eq!(memory.buffer().await, "previous summary");
487
488        memory.set_max_token_limit(500);
489        assert_eq!(memory.max_token_limit(), 500);
490    }
491
492    #[test]
493    fn test_prune_messages_within_limit() {
494        let llm = OpenAIChat::new(create_test_config());
495        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
496            ConversationSummaryBufferMemory::new(llm, 1000);
497
498        let messages = vec![
499            Message::human("Short message 1"),
500            Message::ai("Short reply 1"),
501        ];
502
503        let pruned = memory.prune_messages(&messages);
504
505        assert_eq!(pruned.len(), 2);
506    }
507
508    /// 0.22.0 H-M2: pruning keeps a newest-suffix window, so the budget can cut
509    /// right after an old assistant `tool_calls`, leaving a free-standing Tool
510    /// message at the front of the kept window. That is malformed for
511    /// OpenAI/Anthropic (400); the pruner must drop leading Tool messages.
512    #[test]
513    fn test_prune_never_opens_on_orphan_tool() {
514        let llm = MockLlm::new(vec![Ok("summary".to_string())]);
515        // CharRatioCounter(4): 40 chars == 10 tokens exactly (4 chars/token).
516        let memory: ConversationSummaryBufferMemory<MockLlm> =
517            ConversationSummaryBufferMemory::new(llm, 40)
518                .with_counter(std::sync::Arc::new(CharRatioCounter::new(4)));
519
520        // #1 large assistant message (would be cut), then Tool → human → AI.
521        // Budget 40 tokens fits only #2+#3+#4 (10+10+10); #1 (250 t) is dropped,
522        // leaving the Tool message as the leading kept message — an orphan.
523        let messages = vec![
524            Message::ai("z".repeat(1000)), // ~250 tokens, pruned
525            Message {
526                content: "r".repeat(40),
527                message_type: MessageType::Tool {
528                    tool_call_id: "c1".into(),
529                },
530                ..Message::human("")
531            },
532            Message::human("m".repeat(40)),
533            Message::ai("f".repeat(40)),
534        ];
535
536        let pruned = memory.prune_messages(&messages);
537
538        // The window must not open on the orphaned Tool message.
539        assert!(
540            !matches!(pruned.first().map(|m| &m.message_type),
541                Some(MessageType::Tool { .. })),
542            "pruned window started on an orphaned Tool message: {:?}",
543            pruned.first().map(|m| &m.message_type)
544        );
545        // The newest user+assistant pair is still preserved.
546        assert!(!pruned.is_empty());
547        assert!(matches!(
548            pruned[0].message_type,
549            MessageType::Human
550        ));
551        assert!(matches!(
552            pruned[1].message_type,
553            MessageType::AI
554        ));
555    }
556
557    #[tokio::test]
558    async fn test_buffer_initial_empty() {
559        let llm = OpenAIChat::new(create_test_config());
560        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
561            ConversationSummaryBufferMemory::new(llm, 1000);
562
563        let buffer = memory.buffer().await;
564        assert!(buffer.is_empty());
565    }
566
567    #[tokio::test]
568    async fn test_load_memory_variables_empty() {
569        let llm = OpenAIChat::new(create_test_config());
570        let memory: ConversationSummaryBufferMemory<OpenAIChat> =
571            ConversationSummaryBufferMemory::new(llm, 1000);
572
573        let vars = memory.load_memory_variables(&HashMap::new()).await.unwrap();
574        let history = vars.get("history").unwrap().as_str().unwrap();
575
576        assert!(history.is_empty());
577    }
578
579    #[tokio::test]
580    async fn test_clear() {
581        let llm = OpenAIChat::new(create_test_config());
582        let mut memory: ConversationSummaryBufferMemory<OpenAIChat> =
583            ConversationSummaryBufferMemory::new(llm, 1000);
584
585        memory.chat_memory.add_user_message("test");
586        memory.chat_memory.add_ai_message("reply");
587
588        memory.buffer = "Test summary".to_string();
589
590        memory.clear().await.unwrap();
591
592        assert!(memory.buffer().await.is_empty());
593        assert_eq!(memory.chat_memory().len(), 0);
594    }
595
596    /// P2-4: 剪枝触发摘要、但摘要 LLM 失败时——保留旧摘要、不清空 chat_memory、
597    /// 记录错误且不冒泡;下轮成功总结后摘要生效、错误清空。
598    #[tokio::test]
599    async fn test_prune_summary_failure_keeps_messages_and_retries() {
600        // MockLlm 按 LIFO 消费:第一次剪枝总结失败,第二次成功。
601        let llm = MockLlm::new(vec![
602            Ok("summary-b".to_string()),
603            Err("summarizer down".to_string()),
604        ]);
605        // CharRatioCounter 保证 token 估算可复现(不依赖 tiktoken 在线)。
606        let mut memory: ConversationSummaryBufferMemory<MockLlm> =
607            ConversationSummaryBufferMemory::new(llm, 5)
608                .with_counter(std::sync::Arc::new(CharRatioCounter::new(4)));
609
610        let long_input =
611            "这是一段足够长的中文消息,用来确保本轮消息总 token 数超过预算并触发剪枝总结逻辑";
612        let inputs = HashMap::from([("input".to_string(), long_input.to_string())]);
613        let outputs = HashMap::from([("output".to_string(), long_input.to_string())]);
614
615        // 第一轮:总 token 超限 -> 触发剪枝总结 -> LLM 失败
616        memory.save_context(&inputs, &outputs).await.unwrap();
617        assert!(memory.buffer().await.is_empty(), "失败时不应覆盖旧摘要");
618        assert!(memory
619            .last_summary_error()
620            .unwrap()
621            .contains("summarizer down"));
622        // 失败时不清空原始消息,下轮 prune 才能重试总结
623        assert_eq!(memory.chat_memory().len(), 2);
624
625        // 第二轮:再次触发剪枝总结 -> 成功,摘要生效、错误清空
626        memory.save_context(&inputs, &outputs).await.unwrap();
627        assert_eq!(memory.buffer().await, "summary-b");
628        assert!(memory.last_summary_error().is_none());
629    }
630}
lc_memory/summary_buffer.rs

lc_memory/
summary_buffer.rs
