abu_agent/memory/
summary.rs

1use abu_base::chat::ChatMessage;
2use abu_provider::ChatProvide;
3use tracing::debug;
4use crate::{model::ChatModel, AgentCtxError, AgentResult};
5
6use super::Memory;
7
8pub struct SummarizationMemory<P> {
9    llm: ChatModel<P>,
10    messages: Vec<ChatMessage>,
11    summary_threshold: usize,
12}
13
14impl<P: ChatProvide> SummarizationMemory<P> {
15    pub fn new(llm: ChatModel<P>, summary_threshold: usize) -> Self {
16        Self { 
17            llm,
18            messages: vec![], 
19            summary_threshold,
20        }
21    }
22
23    /// call llm to summary `messages` and reset `messages`
24    async fn consolidate_memory(&mut self) -> AgentResult<()> {
25        debug!("--- [Memory Consolidation Triggered] ---");
26
27        // collection all messages
28        let buffer_text = self.messages.iter()
29            .map(|m| Self::format_message(m))
30            .collect::<Vec<_>>()
31            .join("\n");
32
33        // send to llm
34        let summarization_prompt = format!(
35           "Summarize this conversation for continuity. Include:  \
36            1) What was accomplished, 2) Current state, 3) Key decisions made. \
37            Be concise but preserve critical details.\n\n{}",
38            buffer_text
39        );
40        let messages = vec![
41            ChatMessage::system("You are an expert summarization engine."),
42            ChatMessage::user(summarization_prompt),
43        ];
44        let response = self.llm.chat(messages).await?.message;
45        
46        // update messages
47        self.messages.clear();
48        self.messages.push(ChatMessage::user(format!("[Conversation compressed]: {}", response.content)));
49        self.messages.push(ChatMessage::assistant("Understood. I have the context from the summary. Continuing.", []));
50
51        Ok(())
52    }
53
54    fn format_message(msg: &ChatMessage) -> String {
55        format!("{}: {}", msg.role(), msg.content())
56    }
57}
58
59impl<P: ChatProvide> Memory for SummarizationMemory<P> {
60    type Error = AgentCtxError;
61
62    async fn add(&mut self, user_input: &str, ai_response: &str) -> Result<(), Self::Error> {
63        self.messages.push(ChatMessage::user(user_input));
64        self.messages.push(ChatMessage::assistant(ai_response, []));
65        if self.messages.len() / 2 >= self.summary_threshold {
66            self.consolidate_memory().await?;
67        }
68        Ok(())
69    }
70
71    async fn search(&self, _query: &str) -> Result<Vec<ChatMessage>, Self::Error> {
72        Ok(self.messages.clone())
73    }
74
75    async fn clear(&mut self) -> Result<(), Self::Error> {
76        self.messages.clear();
77        Ok(())
78    }
79}
abu_agent/memory/summary.rs

abu_agent/memory/
summary.rs
