abu_agent/memory/
summary.rs

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

abu_agent/memory/
summary.rs
