abu_agent/memory/
augmented.rs

1use abu_base::chat::{ChatMessage, ChatRequestBuilder};
2use abu_provider::ChatProvide;
3use tracing::debug;
4use crate::{AgentCtxError, AgentError};
5use super::{Memory, SliceWindowMemory};
6
7pub struct AugmentedMemory<P> {
8    llm: P,
9    model: String,
10    recent_memory: SliceWindowMemory,
11    memory_tokens: Vec<String>,
12}
13
14impl<P: ChatProvide> AugmentedMemory<P> {
15    pub fn new(llm: P, model: impl Into<String>, window_size: usize) -> Self {
16        Self { 
17            llm,
18            model: model.into(),
19            recent_memory: SliceWindowMemory::new(window_size),
20            memory_tokens: vec![] 
21        }
22    }
23}
24
25impl<P: ChatProvide> Memory for AugmentedMemory<P> {
26    type Error = AgentCtxError;
27
28    async fn add(&mut self, user_input: &str, ai_response: &str) -> Result<(), Self::Error> {
29        self.recent_memory.add(user_input, ai_response).await.expect("recent_memory error");
30        
31
32        let fact_extraction_prompt = format!(
33            "Analyze the following conversation turn. Does it contain a core fact, preference, or decision that should be remembered long-term? \
34             Examples include user preferences ('I hate flying'), key decisions ('The budget is $1000'), or important facts ('My user ID is 12345').\n\n\
35             Conversation Turn:\nUser: {user_input}\nAI: {ai_response}\n\n\
36             If it contains such a fact, state the fact concisely in one sentence. Otherwise, respond with 'No important fact.'"
37        );
38
39        let request = ChatRequestBuilder::default()
40            .model(&self.model)
41            .messages(vec![
42                ChatMessage::system("You are a fact-extraction expert."),
43                ChatMessage::user(fact_extraction_prompt),
44            ])
45            .build()?;
46        
47        let response = self.llm
48            .chat(&request).await
49            .map_err(|e| AgentError::ChatProvider(Box::new(e)))?
50            .message;
51
52        // lowcase?
53        if !response.content.contains("No important fact") {
54            let extracted_fact = response.content;
55            debug!("--- [Memory Augmentation: New memory token created: '{}'] ---", extracted_fact);
56            self.memory_tokens.push(extracted_fact);
57        }
58
59        Ok(())
60    }
61    
62    async fn search(&self, query: &str) -> Result<Vec<ChatMessage>, Self::Error> {
63        let recent_context = self.recent_memory.search(query).await.expect("recent_memory error");
64        let mut context = vec![
65            ChatMessage::user(format!("### Key Memory Tokens (Long-Term Facts):\n{}\n\n", self.memory_tokens.join("\n"))),
66            ChatMessage::user("### Recent Conversation:\n")
67        ];
68        context.extend(recent_context);
69        Ok(context)
70    }
71
72    async fn clear(&mut self) -> Result<(), Self::Error> {
73        self.recent_memory.clear().await.expect("rec mem err");
74        self.memory_tokens.clear();
75        Ok(())
76    }
77}
abu_agent/memory/augmented.rs

abu_agent/memory/
augmented.rs
