synaptic_memory/
summary.rs1use std::{collections::HashMap, sync::Arc};
2
3use async_trait::async_trait;
4use synaptic_core::{ChatModel, ChatRequest, MemoryStore, Message, SynapseError};
5use tokio::sync::RwLock;
6
7pub 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 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 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 let messages = self.store.load(session_id).await?;
57 if messages.len() > self.buffer_size * 2 {
58 let split_point = messages.len() - self.buffer_size;
60 let older = &messages[..split_point];
61 let recent = &messages[split_point..];
62
63 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 {
82 let mut summaries = self.summary.write().await;
83 summaries.insert(session_id.to_string(), new_summary);
84 }
85
86 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}