vv_agent/memory/manager/
mod.rs

1use std::collections::BTreeSet;
2
3mod compaction;
4mod config;
5mod emergency;
6mod helpers;
7mod limits;
8mod microcompact;
9mod normalization;
10mod prompts;
11mod session_context;
12mod warnings;
13
14use crate::memory::message_sanitizer::filter_empty_assistant_messages;
15use crate::memory::session::SessionMemory;
16use crate::memory::token_utils::count_messages_tokens;
17use crate::types::{Message, MessageRole};
18
19pub use config::{MemoryManagerConfig, SummaryCallback};
20
21use helpers::compact_processed_image_messages;
22
23const MEMORY_SUMMARY_NAME: &str = "memory_summary";
24
25#[derive(Debug, Clone)]
26pub struct MemoryManager {
27    pub config: MemoryManagerConfig,
28    session_memory: Option<SessionMemory>,
29}
30
31impl MemoryManager {
32    pub fn new(mut config: MemoryManagerConfig) -> Self {
33        let session_memory = config.session_memory.take();
34        Self {
35            config,
36            session_memory,
37        }
38    }
39
40    pub fn compact(&mut self, messages: &[Message], force: bool) -> (Vec<Message>, bool) {
41        self.compact_for_cycle_with_usage_inner(messages, 0, None, force, None, None)
42    }
43
44    pub fn compact_for_cycle(
45        &mut self,
46        messages: &[Message],
47        cycle_index: u32,
48        force: bool,
49    ) -> (Vec<Message>, bool) {
50        self.compact_for_cycle_with_usage_inner(
51            messages,
52            cycle_index,
53            Some(cycle_index),
54            force,
55            None,
56            None,
57        )
58    }
59
60    pub fn compact_for_cycle_with_usage(
61        &mut self,
62        messages: &[Message],
63        cycle_index: u32,
64        force: bool,
65        total_tokens: Option<u64>,
66        recent_tool_call_ids: Option<&BTreeSet<String>>,
67    ) -> (Vec<Message>, bool) {
68        self.compact_for_cycle_with_usage_inner(
69            messages,
70            cycle_index,
71            Some(cycle_index),
72            force,
73            total_tokens,
74            recent_tool_call_ids,
75        )
76    }
77
78    fn compact_for_cycle_with_usage_inner(
79        &mut self,
80        messages: &[Message],
81        cycle_index: u32,
82        artifact_cycle_index: Option<u32>,
83        force: bool,
84        total_tokens: Option<u64>,
85        recent_tool_call_ids: Option<&BTreeSet<String>>,
86    ) -> (Vec<Message>, bool) {
87        if messages.is_empty() {
88            return (Vec::new(), false);
89        }
90
91        let cleaned = self.remove_previous_summary(messages);
92        let sanitized = filter_empty_assistant_messages(&cleaned);
93        let changed_by_sanitize = sanitized.len() != messages.len()
94            || sanitized
95                .iter()
96                .zip(messages.iter())
97                .any(|(left, right)| left != right);
98        let mut working_messages = self.apply_session_memory_context(&sanitized);
99        let mut message_length =
100            self.calculate_effective_length(&working_messages, total_tokens, recent_tool_call_ids);
101        if let Some(session_memory) = self.session_memory.as_mut() {
102            let text_messages = working_messages
103                .iter()
104                .filter(|message| {
105                    !matches!(message.role, MessageRole::System | MessageRole::Tool)
106                        && !message.content.trim().is_empty()
107                })
108                .count();
109            if session_memory.should_extract(message_length, text_messages)
110                && session_memory.extract(&working_messages, cycle_index as i32, message_length) > 0
111            {
112                working_messages = self.apply_session_memory_context(&sanitized);
113                message_length =
114                    self.calculate_effective_length(&working_messages, None, recent_tool_call_ids);
115            }
116        }
117        if !force && message_length <= self.autocompact_threshold() {
118            let (warned, warning_inserted) =
119                self.maybe_append_memory_warning(&working_messages, message_length);
120            if warning_inserted {
121                return (warned, true);
122            }
123            if self.should_preemptive_microcompact(message_length) {
124                let (microcompacted, cleared) =
125                    self.microcompact_messages(&working_messages, cycle_index);
126                if cleared > 0 {
127                    return (microcompacted, true);
128                }
129            }
130            return (working_messages, changed_by_sanitize);
131        }
132        let mut summary_source = self.strip_session_memory_context(&working_messages);
133        if !force {
134            let (microcompacted, cleared) =
135                self.microcompact_messages(&working_messages, cycle_index);
136            if cleared > 0 {
137                if self.calculate_effective_length(&microcompacted, None, None)
138                    <= self.autocompact_threshold()
139                {
140                    return (microcompacted, true);
141                }
142                summary_source = microcompacted;
143            }
144            let (image_compacted, image_changed) =
145                compact_processed_image_messages(&summary_source);
146            let (artifact_compacted, artifact_changed) =
147                self.compact_large_tool_results(&image_compacted, artifact_cycle_index);
148            if (image_changed || artifact_changed)
149                && count_messages_tokens(&artifact_compacted, &self.config.model)
150                    <= self.autocompact_threshold()
151            {
152                return (artifact_compacted, true);
153            }
154            if image_changed {
155                summary_source = image_compacted;
156            }
157        }
158        let (compacted, changed) = self.compress_memory(&summary_source, artifact_cycle_index);
159        if changed {
160            if let Some(session_memory) = self.session_memory.as_mut() {
161                session_memory
162                    .on_compaction(Some(count_messages_tokens(&compacted, &self.config.model)));
163            }
164        }
165        (compacted, changed)
166    }
167
168    fn remove_previous_summary(&self, messages: &[Message]) -> Vec<Message> {
169        messages
170            .iter()
171            .filter(|message| {
172                !(message.role == MessageRole::System
173                    && message.name.as_deref() == Some(MEMORY_SUMMARY_NAME))
174            })
175            .cloned()
176            .collect()
177    }
178}
vv_agent/memory/manager/mod.rs

vv_agent/memory/manager/
mod.rs
