Skip to main content

synaptic_memory/
summary.rs

1use std::{collections::HashMap, sync::Arc};
2
3use async_trait::async_trait;
4use synaptic_core::{ChatModel, ChatRequest, MemoryStore, Message, SynapseError};
5use tokio::sync::RwLock;
6
7/// A memory strategy that summarizes older messages using a ChatModel.
8///
9/// Keeps the most recent `buffer_size` messages verbatim. When the total
10/// number of stored messages exceeds `buffer_size * 2`, the older messages
11/// are summarized into a single system message that is prepended on `load`.
12pub struct ConversationSummaryMemory {
13    store: Arc<dyn MemoryStore>,
14    model: Arc<dyn ChatModel>,
15    summary: Arc<RwLock<HashMap<String, String>>>,
16    buffer_size: usize,
17}
18
19impl ConversationSummaryMemory {
20    /// Create a new summary memory.
21    ///
22    /// - `store` — the underlying message store
23    /// - `model` — the ChatModel used to generate summaries
24    /// - `buffer_size` — number of recent messages to keep verbatim
25    pub fn new(store: Arc<dyn MemoryStore>, model: Arc<dyn ChatModel>, buffer_size: usize) -> Self {
26        Self {
27            store,
28            model,
29            summary: Arc::new(RwLock::new(HashMap::new())),
30            buffer_size,
31        }
32    }
33
34    /// Generate a summary of the given messages using the ChatModel.
35    async fn summarize(&self, messages: &[Message]) -> Result<String, SynapseError> {
36        let conversation = messages
37            .iter()
38            .map(|m| format!("{}: {}", m.role(), m.content()))
39            .collect::<Vec<_>>()
40            .join("\n");
41
42        let prompt = format!("Summarize the following conversation concisely:\n\n{conversation}");
43
44        let request = ChatRequest::new(vec![Message::human(prompt)]);
45        let response = self.model.chat(request).await?;
46        Ok(response.message.content().to_string())
47    }
48}
49
50#[async_trait]
51impl MemoryStore for ConversationSummaryMemory {
52    async fn append(&self, session_id: &str, message: Message) -> Result<(), SynapseError> {
53        self.store.append(session_id, message).await?;
54
55        // Check if we need to trigger summarization
56        let messages = self.store.load(session_id).await?;
57        if messages.len() > self.buffer_size * 2 {
58            // Summarize the older messages (everything except the most recent buffer_size)
59            let split_point = messages.len() - self.buffer_size;
60            let older = &messages[..split_point];
61            let recent = &messages[split_point..];
62
63            // Build new summary incorporating any existing summary
64            let existing_summary = {
65                let summaries = self.summary.read().await;
66                summaries.get(session_id).cloned()
67            };
68
69            let to_summarize = if let Some(ref existing) = existing_summary {
70                let mut with_context =
71                    vec![Message::system(format!("Previous summary: {existing}"))];
72                with_context.extend_from_slice(older);
73                with_context
74            } else {
75                older.to_vec()
76            };
77
78            let new_summary = self.summarize(&to_summarize).await?;
79
80            // Update the summary
81            {
82                let mut summaries = self.summary.write().await;
83                summaries.insert(session_id.to_string(), new_summary);
84            }
85
86            // Replace the store contents with just the recent messages
87            self.store.clear(session_id).await?;
88            for msg in recent {
89                self.store.append(session_id, msg.clone()).await?;
90            }
91        }
92
93        Ok(())
94    }
95
96    async fn load(&self, session_id: &str) -> Result<Vec<Message>, SynapseError> {
97        let messages = self.store.load(session_id).await?;
98        let summaries = self.summary.read().await;
99
100        if let Some(summary_text) = summaries.get(session_id) {
101            let mut result = vec![Message::system(format!(
102                "Summary of earlier conversation: {summary_text}"
103            ))];
104            result.extend(messages);
105            Ok(result)
106        } else {
107            Ok(messages)
108        }
109    }
110
111    async fn clear(&self, session_id: &str) -> Result<(), SynapseError> {
112        self.store.clear(session_id).await?;
113        let mut summaries = self.summary.write().await;
114        summaries.remove(session_id);
115        Ok(())
116    }
117}