ambi_memory/
extension.rs

1// ambi-memory/src/extension.rs
2
3use crate::error::{MemoryError, Result};
4use crate::provider::kv::KvMemoryProvider;
5use crate::provider::semantic::SemanticMemoryProvider;
6use crate::provider::summary::SummaryMemoryProvider;
7
8use ambi::agent::core::{Agent, AgentState};
9use ambi::Message;
10use async_trait::async_trait;
11use std::sync::Arc;
12
13/// The grand unified extension trait that empowers `AgentState` with 3D cognitive memory.
14#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
15#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
16pub trait AgentStateMemoryExt {
17    // --- 1. KV State Memory (Reflexion / Settings) ---
18    async fn remember_kv(
19        &self,
20        provider: &dyn KvMemoryProvider,
21        key: &str,
22        value: &str,
23    ) -> Result<()>;
24    async fn recall_kv_into_context(&mut self, provider: &dyn KvMemoryProvider) -> Result<()>;
25
26    // --- 2. Semantic Vector Memory (Long-term interactions) ---
27    async fn archive_semantic(
28        &self,
29        provider: &dyn SemanticMemoryProvider,
30        text: &str,
31    ) -> Result<()>;
32    async fn recall_semantic_into_context(
33        &mut self,
34        provider: &dyn SemanticMemoryProvider,
35        query: &str,
36        limit: usize,
37    ) -> Result<()>;
38
39    // --- 3. Rolling Summary Memory (Anti-amnesia on eviction) ---
40    async fn inject_summary_context(&mut self, provider: &dyn SummaryMemoryProvider) -> Result<()>;
41
42    /// **Black Magic**: Automatically digests evicted messages into a rolling summary using the Agent itself!
43    async fn summarize_evicted_messages(
44        &self,
45        agent: &Agent,
46        provider: &dyn SummaryMemoryProvider,
47        evicted: &[Arc<Message>],
48    ) -> Result<()>;
49}
50
51#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
52#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
53impl AgentStateMemoryExt for AgentState {
54    // --- KV Memory Implementations ---
55    async fn remember_kv(
56        &self,
57        provider: &dyn KvMemoryProvider,
58        key: &str,
59        value: &str,
60    ) -> Result<()> {
61        provider.store(&self.session_id, key, value).await
62    }
63
64    async fn recall_kv_into_context(&mut self, provider: &dyn KvMemoryProvider) -> Result<()> {
65        let memories = provider.retrieve_all(&self.session_id).await?;
66        if !memories.is_empty() {
67            let mut compiled = String::from("[STATE MEMORY]:\n");
68            for (k, v) in memories {
69                compiled.push_str(&format!("- {}: {}\n", k, v));
70            }
71            self.append_dynamic_context(&compiled);
72        }
73        Ok(())
74    }
75
76    // --- Semantic Vector Memory Implementations ---
77    async fn archive_semantic(
78        &self,
79        provider: &dyn SemanticMemoryProvider,
80        text: &str,
81    ) -> Result<()> {
82        provider.add_memory(&self.session_id, text).await
83    }
84
85    async fn recall_semantic_into_context(
86        &mut self,
87        provider: &dyn SemanticMemoryProvider,
88        query: &str,
89        limit: usize,
90    ) -> Result<()> {
91        let memories = provider
92            .search_memories(&self.session_id, query, limit)
93            .await?;
94        if !memories.is_empty() {
95            let mut compiled = String::from("[RECALLED PAST EXPERIENCES]:\n");
96            for (i, m) in memories.iter().enumerate() {
97                compiled.push_str(&format!("{}. {}\n", i + 1, m));
98            }
99            self.append_dynamic_context(&compiled);
100        }
101        Ok(())
102    }
103
104    // --- Rolling Summary Memory Implementations ---
105    async fn inject_summary_context(&mut self, provider: &dyn SummaryMemoryProvider) -> Result<()> {
106        if let Some(summary) = provider.get_summary(&self.session_id).await? {
107            let formatted = format!("[CONVERSATION SUMMARY]:\n{}\n", summary);
108            // We append this summary to the dynamic context so the LLM remembers the distant past
109            self.append_dynamic_context(&formatted);
110        }
111        Ok(())
112    }
113
114    async fn summarize_evicted_messages(
115        &self,
116        agent: &Agent,
117        provider: &dyn SummaryMemoryProvider,
118        evicted: &[Arc<Message>],
119    ) -> Result<()> {
120        if evicted.is_empty() {
121            return Ok(());
122        }
123
124        let old_summary = provider
125            .get_summary(&self.session_id)
126            .await?
127            .unwrap_or_else(|| "No previous summary.".to_string());
128
129        let mut evicted_text = String::new();
130        for msg in evicted {
131            evicted_text.push_str(&format!("{}\n", msg));
132        }
133
134        let instruction = format!(
135            "You are an AI assistant helping to compress conversational memory.\n\
136            [Current Summary]: {}\n\
137            [New Evicted Messages]: {}\n\
138            Task: Merge the new messages into the current summary. Keep it concise, capturing crucial facts, entities, and context. Output ONLY the new summary.",
139            old_summary, evicted_text
140        );
141
142        // We use the agent to process the summarization out-of-band!
143        let runner = ambi::ChatRunner::new(1); // Internal transient runner
144
145        // Create an ephemeral, isolated state just for this summarization task
146        // We do NOT want the summarization process to pollute the current user's chat history!
147        let temp_state = AgentState::new_shared(format!("{}_summarizer", self.session_id));
148
149        // Execute the summarization using the core LLM engine
150        let new_summary = runner
151            .chat(agent, &temp_state, &instruction)
152            .await
153            .map_err(|e| {
154                MemoryError::SummaryError(format!("LLM failed to generate summary: {}", e))
155            })?;
156
157        // Update the persistent summary provider
158        provider
159            .update_summary(&self.session_id, &new_summary)
160            .await?;
161
162        log::info!(
163            "Successfully compressed {} evicted messages into rolling summary.",
164            evicted.len()
165        );
166        Ok(())
167    }
168}
ambi_memory/extension.rs

ambi_memory/
extension.rs
