lc_memory/
summary.rs

1// lc-memory/src/summary.rs
2//! Conversation Summary Memory
3//!
4//! Uses LLM to automatically summarize conversation history, solving the long conversation token explosion problem.
5
6use async_trait::async_trait;
7use serde_json::Value;
8use std::collections::HashMap;
9
10use super::base::{BaseChatMemory, BaseMemory, ChatMessageHistory, MemoryError};
11use lc_core::language_models::BaseChatModel;
12use lc_core::language_models::LLMResult;
13use lc_core::runnables::Runnable;
14use lc_prompts::PromptTemplate;
15use lc_schema::Message;
16
17/// Default summary prompt
18const DEFAULT_SUMMARY_PROMPT: &str = "Progressively summarize the lines of conversation provided, adding onto the previous summary returning a new summary.
19
20EXAMPLE
21Summary of conversation:
22Human: My name is Zhang San, I like programming.
23AI: Hello Zhang San, nice to meet you! You like programming, any particular language?
24Human: I like Rust.
25AI: Rust is a great programming language, focused on safety and performance.
26
27New lines of conversation:
28Human: I also like Python.
29AI: Python is also popular, with concise syntax, suitable for rapid development.
30
31New summary:
32Human Zhang San likes programming, especially Rust and Python. AI discussed the characteristics of these two languages with Zhang San.
33
34END OF EXAMPLE
35
36Current summary:
37{summary}
38
39New lines of conversation:
40{new_lines}
41
42New summary:";
43
44/// Conversation Summary Memory
45///
46/// Uses LLM to automatically summarize conversation history, avoiding overly long context.
47///
48/// # Example
49/// ```ignore
50/// use lc_memory::ConversationSummaryMemory;
51/// use lc_providers::OpenAIChat;
52///
53/// let llm = OpenAIChat::new(config);
54/// let memory = ConversationSummaryMemory::new(llm);
55///
56/// // Automatically generates summary after each conversation round
57/// memory.save_context(&inputs, &outputs).await?;
58///
59/// // Returns summary instead of full history when loading
60/// let vars = memory.load_memory_variables(&HashMap::new()).await?;
61/// ```
62pub struct ConversationSummaryMemory<M: BaseChatModel> {
63    llm: M,
64
65    /// Current summary (M67: removed Mutex - &mut self already guarantees exclusive access)
66    buffer: String,
67
68    /// Chat history (H29: trimmed after each summary to prevent unbounded growth)
69    chat_memory: ChatMessageHistory,
70
71    /// Input key name
72    input_key: String,
73
74    /// Output key name
75    output_key: String,
76
77    /// Memory variable name
78    memory_key: String,
79
80    /// Summary prompt
81    summary_prompt: String,
82
83    /// Whether to return message objects
84    return_messages: bool,
85
86    /// H29: Maximum number of recent message pairs to keep in chat_memory
87    /// after summarization. Older messages are discarded since the summary
88    /// already captures their content. Default: 2 (last turn only).
89    max_recent_turns: usize,
90
91    /// P2-4: 摘要 LLM 失败时累计的未总结增量(本轮 + 之前失败的轮次)。
92    /// 下轮成功总结时并入 `new_lines` 一并总结,保证失败轮次的内容不丢失。
93    pending_lines: String,
94
95    /// P2-4: 最近一次摘要 LLM 失败的原因;成功总结或 `clear()` 后清空。
96    /// 摘要失败不再让 `save_context` 冒泡错误打断链,而是保留旧摘要继续工作。
97    last_summary_error: Option<String>,
98}
99
100impl<M: BaseChatModel> ConversationSummaryMemory<M> {
101    /// Create a new summary memory
102    pub fn new(llm: M) -> Self {
103        Self {
104            llm,
105            buffer: String::new(),
106            chat_memory: ChatMessageHistory::new(),
107            input_key: "input".to_string(),
108            output_key: "output".to_string(),
109            memory_key: "history".to_string(),
110            summary_prompt: DEFAULT_SUMMARY_PROMPT.to_string(),
111            return_messages: false,
112            max_recent_turns: 2,
113            pending_lines: String::new(),
114            last_summary_error: None,
115        }
116    }
117
118    /// Create from existing messages
119    pub fn from_messages(llm: M, messages: Vec<Message>) -> Self {
120        let chat_memory = ChatMessageHistory::from_messages(messages);
121        Self {
122            llm,
123            buffer: String::new(),
124            chat_memory,
125            input_key: "input".to_string(),
126            output_key: "output".to_string(),
127            memory_key: "history".to_string(),
128            summary_prompt: DEFAULT_SUMMARY_PROMPT.to_string(),
129            return_messages: false,
130            max_recent_turns: 2,
131            pending_lines: String::new(),
132            last_summary_error: None,
133        }
134    }
135
136    /// Set input key name
137    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
138        self.input_key = key.into();
139        self
140    }
141
142    /// Set output key name
143    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
144        self.output_key = key.into();
145        self
146    }
147
148    /// Set memory variable name
149    pub fn with_memory_key(mut self, key: impl Into<String>) -> Self {
150        self.memory_key = key.into();
151        self
152    }
153
154    /// Set summary prompt
155    pub fn with_summary_prompt(mut self, prompt: impl Into<String>) -> Self {
156        self.summary_prompt = prompt.into();
157        self
158    }
159
160    /// Set whether to return message objects
161    pub fn with_return_messages(mut self, return_messages: bool) -> Self {
162        self.return_messages = return_messages;
163        self
164    }
165
166    /// H29: Set maximum recent turns to keep in chat_memory after summarization
167    pub fn with_max_recent_turns(mut self, max: usize) -> Self {
168        self.max_recent_turns = max;
169        self
170    }
171
172    /// Get chat history
173    pub fn chat_memory(&self) -> &ChatMessageHistory {
174        &self.chat_memory
175    }
176
177    /// Get current summary
178    pub async fn buffer(&self) -> String {
179        self.buffer.clone()
180    }
181
182    /// P2-4: 最近一次摘要失败的原因(无失败则 `None`)。
183    pub fn last_summary_error(&self) -> Option<&str> {
184        self.last_summary_error.as_deref()
185    }
186
187    /// P2-4: 摘要失败后累计的待总结增量行。
188    pub fn pending_lines(&self) -> &str {
189        &self.pending_lines
190    }
191
192    /// Format new conversation lines
193    fn format_new_lines(&self, input: &str, output: &str) -> String {
194        format!("Human: {}\nAI: {}", input, output)
195    }
196
197    /// Generate new summary
198    async fn predict_new_summary(&self, new_lines: &str) -> Result<String, MemoryError> {
199        let buffer = self.buffer.clone();
200
201        // P2-4: 把之前失败轮次累计的增量并入本次总结,失败轮次的内容不丢失。
202        let mut combined = String::new();
203        if !self.pending_lines.is_empty() {
204            combined.push_str(&self.pending_lines);
205            combined.push('\n');
206        }
207        combined.push_str(new_lines);
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", combined.as_str());
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 ConversationSummaryMemory<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
247        if self.return_messages {
248            let summary_msg = Message::system(&buffer);
249            result.insert(
250                self.memory_key.clone(),
251                serde_json::to_value(&summary_msg).unwrap_or(Value::Null),
252            );
253        } else {
254            result.insert(self.memory_key.clone(), Value::String(buffer));
255        }
256
257        Ok(result)
258    }
259
260    async fn save_context(
261        &mut self,
262        inputs: &HashMap<String, String>,
263        outputs: &HashMap<String, String>,
264    ) -> Result<(), MemoryError> {
265        // P1-1: 与 Buffer/Window 一致——缺失 key 返回 SaveError,不再静默用空串
266        // 存空消息(否则 LLM 会对 "Human: \nAI: " 空行总结,白烧一次调用)。
267        let input = inputs.get(&self.input_key).ok_or_else(|| {
268            MemoryError::SaveError(format!("Missing input key '{}'", self.input_key))
269        })?;
270        let output = outputs.get(&self.output_key).ok_or_else(|| {
271            MemoryError::SaveError(format!("Missing output key '{}'", self.output_key))
272        })?;
273
274        self.chat_memory.add_user_message(input);
275        self.chat_memory.add_ai_message(output);
276
277        let new_lines = self.format_new_lines(input, output);
278
279        // P2-4: 摘要失败时保留旧摘要、记录错误并把本轮增量累计到 pending_lines
280        // 供下轮重试,而不是让错误冒泡打断上层链。
281        let new_summary = match self.predict_new_summary(&new_lines).await {
282            Ok(s) => s,
283            Err(e) => {
284                if !self.pending_lines.is_empty() {
285                    self.pending_lines.push('\n');
286                }
287                self.pending_lines.push_str(&new_lines);
288                self.last_summary_error = Some(e.to_string());
289                log::warn!(
290                    "ConversationSummaryMemory 摘要失败,保留旧摘要待下轮重试: {}",
291                    e
292                );
293                return Ok(());
294            }
295        };
296
297        self.buffer = new_summary;
298        self.pending_lines.clear();
299        self.last_summary_error = None;
300
301        // H29: Trim chat_memory to prevent unbounded growth.
302        // Since the summary already captures all conversation content,
303        // only keep the most recent turns for context continuity.
304        let max_messages = self.max_recent_turns * 2;
305        let current_len = self.chat_memory.len();
306        if current_len > max_messages {
307            let messages = self.chat_memory.messages().to_vec();
308            self.chat_memory.clear();
309            // Preserve System messages and the most recent turns
310            let start = current_len.saturating_sub(max_messages);
311            for msg in messages.iter().take(start) {
312                if matches!(msg.message_type, lc_schema::MessageType::System) {
313                    self.chat_memory.add_system_message(&msg.content);
314                }
315            }
316            for msg in messages.iter().skip(start) {
317                if matches!(msg.message_type, lc_schema::MessageType::Human) {
318                    self.chat_memory.add_user_message(&msg.content);
319                } else if matches!(msg.message_type, lc_schema::MessageType::AI) {
320                    self.chat_memory.add_ai_message(&msg.content);
321                } else if matches!(msg.message_type, lc_schema::MessageType::System) {
322                    self.chat_memory.add_system_message(&msg.content);
323                }
324            }
325        }
326
327        Ok(())
328    }
329
330    async fn clear(&mut self) -> Result<(), MemoryError> {
331        self.buffer = String::new();
332        self.chat_memory.clear();
333        self.pending_lines.clear();
334        self.last_summary_error = None;
335        Ok(())
336    }
337}
338
339/// P0-1: `ConversationSummaryMemory` 实现 `BaseChatMemory`。
340impl<M: BaseChatModel + Send + Sync + 'static> BaseChatMemory for ConversationSummaryMemory<M>
341where
342    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
343{
344    fn messages(&self) -> &[Message] {
345        self.chat_memory.messages()
346    }
347
348    fn add_message(&mut self, message: Message) {
349        self.chat_memory.add_message(message);
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356    use crate::test_support::MockLlm;
357    use lc_providers::{OpenAIChat, OpenAIConfig};
358
359    fn create_test_config() -> OpenAIConfig {
360        OpenAIConfig {
361            api_key: "sk-test".to_string(),
362            base_url: "https://api.openai.com/v1".to_string(),
363            model: "gpt-3.5-turbo".to_string(),
364            streaming: false,
365            ..Default::default()
366        }
367    }
368
369    #[test]
370    fn test_new() {
371        let llm = OpenAIChat::new(create_test_config());
372        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
373
374        assert_eq!(memory.memory_variables(), vec!["history"]);
375    }
376
377    #[test]
378    fn test_with_options() {
379        let llm = OpenAIChat::new(create_test_config());
380        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm)
381            .with_input_key("question")
382            .with_output_key("answer")
383            .with_memory_key("context");
384
385        assert_eq!(memory.input_key, "question");
386        assert_eq!(memory.output_key, "answer");
387        assert_eq!(memory.memory_key, "context");
388    }
389
390    #[test]
391    fn test_from_messages() {
392        let llm = OpenAIChat::new(create_test_config());
393        let messages = vec![Message::human("Hello"), Message::ai("Hello!")];
394        let memory: ConversationSummaryMemory<OpenAIChat> =
395            ConversationSummaryMemory::from_messages(llm, messages);
396
397        assert_eq!(memory.chat_memory().len(), 2);
398    }
399
400    #[test]
401    fn test_format_new_lines() {
402        let llm = OpenAIChat::new(create_test_config());
403        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
404
405        let new_lines = memory.format_new_lines("Hello", "Hello!");
406        assert_eq!(new_lines, "Human: Hello\nAI: Hello!");
407    }
408
409    #[tokio::test]
410    async fn test_buffer_initial_empty() {
411        let llm = OpenAIChat::new(create_test_config());
412        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
413
414        let buffer = memory.buffer().await;
415        assert!(buffer.is_empty());
416    }
417
418    #[tokio::test]
419    async fn test_load_memory_variables_empty() {
420        let llm = OpenAIChat::new(create_test_config());
421        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
422
423        let vars = memory.load_memory_variables(&HashMap::new()).await.unwrap();
424        let history = vars.get("history").unwrap().as_str().unwrap();
425
426        assert!(history.is_empty());
427    }
428
429    #[tokio::test]
430    async fn test_clear() {
431        let llm = OpenAIChat::new(create_test_config());
432        let mut memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
433
434        memory.chat_memory.add_user_message("test");
435        memory.chat_memory.add_ai_message("reply");
436
437        memory.buffer = "Test summary".to_string();
438
439        memory.clear().await.unwrap();
440
441        assert!(memory.buffer().await.is_empty());
442        assert_eq!(memory.chat_memory().len(), 0);
443    }
444
445    /// P2-4: 摘要 LLM 失败时保留旧摘要、记录错误、不让 `save_context` 冒泡;
446    /// 下轮成功总结时把失败轮次补进新摘要。
447    #[tokio::test]
448    async fn test_summary_failure_keeps_old_summary_and_retries() {
449        // MockLlm 按 LIFO 消费:先失败,后成功。
450        let llm = MockLlm::new(vec![
451            Ok("final summary".to_string()),
452            Err("summarizer down".to_string()),
453        ]);
454        let mut memory: ConversationSummaryMemory<MockLlm> = ConversationSummaryMemory::new(llm);
455
456        let inputs = HashMap::from([("input".to_string(), "你好".to_string())]);
457        let outputs = HashMap::from([("output".to_string(), "你好!".to_string())]);
458
459        // 第一轮:摘要失败 -> save_context 返回 Ok,旧摘要保留,错误被记录
460        memory.save_context(&inputs, &outputs).await.unwrap();
461        assert!(memory.buffer().await.is_empty(), "失败时不应覆盖旧摘要");
462        assert!(memory
463            .last_summary_error()
464            .unwrap()
465            .contains("summarizer down"));
466        assert!(memory.pending_lines().contains("Human: 你好"));
467
468        // 第二轮:摘要成功 -> 新摘要生效,错误与增量清空
469        memory.save_context(&inputs, &outputs).await.unwrap();
470        assert_eq!(memory.buffer().await, "final summary");
471        assert!(memory.last_summary_error().is_none());
472        assert!(memory.pending_lines().is_empty());
473    }
474}
lc_memory/summary.rs

lc_memory/
summary.rs
