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::{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
92impl<M: BaseChatModel> ConversationSummaryMemory<M> {
93    /// Create a new summary memory
94    pub fn new(llm: M) -> Self {
95        Self {
96            llm,
97            buffer: String::new(),
98            chat_memory: ChatMessageHistory::new(),
99            input_key: "input".to_string(),
100            output_key: "output".to_string(),
101            memory_key: "history".to_string(),
102            summary_prompt: DEFAULT_SUMMARY_PROMPT.to_string(),
103            return_messages: false,
104            max_recent_turns: 2,
105        }
106    }
107
108    /// Create from existing messages
109    pub fn from_messages(llm: M, messages: Vec<Message>) -> Self {
110        let chat_memory = ChatMessageHistory::from_messages(messages);
111        Self {
112            llm,
113            buffer: String::new(),
114            chat_memory,
115            input_key: "input".to_string(),
116            output_key: "output".to_string(),
117            memory_key: "history".to_string(),
118            summary_prompt: DEFAULT_SUMMARY_PROMPT.to_string(),
119            return_messages: false,
120            max_recent_turns: 2,
121        }
122    }
123
124    /// Set input key name
125    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
126        self.input_key = key.into();
127        self
128    }
129
130    /// Set output key name
131    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
132        self.output_key = key.into();
133        self
134    }
135
136    /// Set memory variable name
137    pub fn with_memory_key(mut self, key: impl Into<String>) -> Self {
138        self.memory_key = key.into();
139        self
140    }
141
142    /// Set summary prompt
143    pub fn with_summary_prompt(mut self, prompt: impl Into<String>) -> Self {
144        self.summary_prompt = prompt.into();
145        self
146    }
147
148    /// Set whether to return message objects
149    pub fn with_return_messages(mut self, return_messages: bool) -> Self {
150        self.return_messages = return_messages;
151        self
152    }
153
154    /// H29: Set maximum recent turns to keep in chat_memory after summarization
155    pub fn with_max_recent_turns(mut self, max: usize) -> Self {
156        self.max_recent_turns = max;
157        self
158    }
159
160    /// Get chat history
161    pub fn chat_memory(&self) -> &ChatMessageHistory {
162        &self.chat_memory
163    }
164
165    /// Get current summary
166    pub async fn buffer(&self) -> String {
167        self.buffer.clone()
168    }
169
170    /// Format new conversation lines
171    fn format_new_lines(&self, input: &str, output: &str) -> String {
172        format!("Human: {}\nAI: {}", input, output)
173    }
174
175    /// Generate new summary
176    async fn predict_new_summary(&self, new_lines: &str) -> Result<String, MemoryError> {
177        let buffer = self.buffer.clone();
178
179        let prompt = {
180            let template = PromptTemplate::new(&self.summary_prompt);
181            let mut vars: std::collections::HashMap<&str, &str> = std::collections::HashMap::new();
182            vars.insert("summary", buffer.as_str());
183            vars.insert("new_lines", new_lines);
184            template
185                .format(&vars)
186                .unwrap_or_else(|_| self.summary_prompt.clone())
187        };
188
189        let messages = vec![Message::human(&prompt)];
190
191        let result =
192            self.llm.invoke(messages, None).await.map_err(|e| {
193                MemoryError::SaveError(format!("LLM summary generation failed: {}", e))
194            })?;
195
196        Ok(result.content)
197    }
198}
199
200#[async_trait]
201impl<M: BaseChatModel + Send + Sync + 'static> BaseMemory for ConversationSummaryMemory<M>
202where
203    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
204{
205    fn memory_variables(&self) -> Vec<&str> {
206        vec![&self.memory_key]
207    }
208
209    async fn load_memory_variables(
210        &self,
211        _inputs: &HashMap<String, String>,
212    ) -> Result<HashMap<String, Value>, MemoryError> {
213        let mut result = HashMap::new();
214
215        let buffer = self.buffer.clone();
216
217        if self.return_messages {
218            let summary_msg = Message::system(&buffer);
219            result.insert(
220                self.memory_key.clone(),
221                serde_json::to_value(&summary_msg).unwrap_or(Value::Null),
222            );
223        } else {
224            result.insert(self.memory_key.clone(), Value::String(buffer));
225        }
226
227        Ok(result)
228    }
229
230    async fn save_context(
231        &mut self,
232        inputs: &HashMap<String, String>,
233        outputs: &HashMap<String, String>,
234    ) -> Result<(), MemoryError> {
235        let empty = String::new();
236        let input = inputs.get(&self.input_key).unwrap_or(&empty);
237        let output = outputs.get(&self.output_key).unwrap_or(&empty);
238
239        self.chat_memory.add_user_message(input);
240        self.chat_memory.add_ai_message(output);
241
242        let new_lines = self.format_new_lines(input, output);
243        let new_summary = self.predict_new_summary(&new_lines).await?;
244
245        self.buffer = new_summary;
246
247        // H29: Trim chat_memory to prevent unbounded growth.
248        // Since the summary already captures all conversation content,
249        // only keep the most recent turns for context continuity.
250        let max_messages = self.max_recent_turns * 2;
251        let current_len = self.chat_memory.len();
252        if current_len > max_messages {
253            let messages = self.chat_memory.messages().to_vec();
254            self.chat_memory.clear();
255            // Preserve System messages and the most recent turns
256            let start = current_len.saturating_sub(max_messages);
257            for msg in messages.iter().take(start) {
258                if matches!(msg.message_type, lc_schema::MessageType::System) {
259                    self.chat_memory.add_system_message(&msg.content);
260                }
261            }
262            for msg in messages.iter().skip(start) {
263                if matches!(msg.message_type, lc_schema::MessageType::Human) {
264                    self.chat_memory.add_user_message(&msg.content);
265                } else if matches!(msg.message_type, lc_schema::MessageType::AI) {
266                    self.chat_memory.add_ai_message(&msg.content);
267                } else if matches!(msg.message_type, lc_schema::MessageType::System) {
268                    self.chat_memory.add_system_message(&msg.content);
269                }
270            }
271        }
272
273        Ok(())
274    }
275
276    async fn clear(&mut self) -> Result<(), MemoryError> {
277        self.buffer = String::new();
278        self.chat_memory.clear();
279        Ok(())
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286    use lc_providers::{OpenAIChat, OpenAIConfig};
287
288    fn create_test_config() -> OpenAIConfig {
289        OpenAIConfig {
290            api_key: "sk-test".to_string(),
291            base_url: "https://api.openai.com/v1".to_string(),
292            model: "gpt-3.5-turbo".to_string(),
293            streaming: false,
294            ..Default::default()
295        }
296    }
297
298    #[test]
299    fn test_new() {
300        let llm = OpenAIChat::new(create_test_config());
301        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
302
303        assert_eq!(memory.memory_variables(), vec!["history"]);
304    }
305
306    #[test]
307    fn test_with_options() {
308        let llm = OpenAIChat::new(create_test_config());
309        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm)
310            .with_input_key("question")
311            .with_output_key("answer")
312            .with_memory_key("context");
313
314        assert_eq!(memory.input_key, "question");
315        assert_eq!(memory.output_key, "answer");
316        assert_eq!(memory.memory_key, "context");
317    }
318
319    #[test]
320    fn test_from_messages() {
321        let llm = OpenAIChat::new(create_test_config());
322        let messages = vec![Message::human("Hello"), Message::ai("Hello!")];
323        let memory: ConversationSummaryMemory<OpenAIChat> =
324            ConversationSummaryMemory::from_messages(llm, messages);
325
326        assert_eq!(memory.chat_memory().len(), 2);
327    }
328
329    #[test]
330    fn test_format_new_lines() {
331        let llm = OpenAIChat::new(create_test_config());
332        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
333
334        let new_lines = memory.format_new_lines("Hello", "Hello!");
335        assert_eq!(new_lines, "Human: Hello\nAI: Hello!");
336    }
337
338    #[tokio::test]
339    async fn test_buffer_initial_empty() {
340        let llm = OpenAIChat::new(create_test_config());
341        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
342
343        let buffer = memory.buffer().await;
344        assert!(buffer.is_empty());
345    }
346
347    #[tokio::test]
348    async fn test_load_memory_variables_empty() {
349        let llm = OpenAIChat::new(create_test_config());
350        let memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
351
352        let vars = memory.load_memory_variables(&HashMap::new()).await.unwrap();
353        let history = vars.get("history").unwrap().as_str().unwrap();
354
355        assert!(history.is_empty());
356    }
357
358    #[tokio::test]
359    async fn test_clear() {
360        let llm = OpenAIChat::new(create_test_config());
361        let mut memory: ConversationSummaryMemory<OpenAIChat> = ConversationSummaryMemory::new(llm);
362
363        memory.chat_memory.add_user_message("test");
364        memory.chat_memory.add_ai_message("reply");
365
366        memory.buffer = "Test summary".to_string();
367
368        memory.clear().await.unwrap();
369
370        assert!(memory.buffer().await.is_empty());
371        assert_eq!(memory.chat_memory().len(), 0);
372    }
373}
lc_memory/summary.rs

lc_memory/
summary.rs
