vv_agent/memory/manager/
limits.rs

1use std::collections::BTreeSet;
2
3use super::MemoryManager;
4use crate::events::MemoryCompactTrigger;
5use crate::memory::token_utils::{compute_compaction_threshold, count_messages_tokens};
6use crate::types::{Message, MessageRole};
7
8impl MemoryManager {
9    pub fn autocompact_threshold(&self) -> u64 {
10        compute_compaction_threshold(
11            self.config.compact_threshold,
12            self.config.model_context_window,
13            self.config.reserved_output_tokens,
14            self.config.autocompact_buffer_tokens,
15        )
16    }
17
18    pub fn effective_context_window(&self) -> u64 {
19        self.config
20            .model_context_window
21            .saturating_sub(self.config.reserved_output_tokens)
22    }
23
24    pub fn estimate_memory_usage_percentage(
25        &self,
26        messages: &[Message],
27        total_tokens: Option<u64>,
28        recent_tool_call_ids: Option<&BTreeSet<String>>,
29    ) -> u64 {
30        let threshold = self.autocompact_threshold();
31        if threshold == 0 {
32            return 0;
33        }
34        let used_tokens =
35            self.calculate_effective_length(messages, total_tokens, recent_tool_call_ids);
36        (used_tokens.saturating_mul(100)) / threshold
37    }
38
39    pub(crate) fn compaction_trigger(
40        &self,
41        messages: &[Message],
42        total_tokens: Option<u64>,
43        recent_tool_call_ids: Option<&BTreeSet<String>>,
44    ) -> Option<MemoryCompactTrigger> {
45        let used_tokens =
46            self.calculate_effective_length(messages, total_tokens, recent_tool_call_ids);
47        if used_tokens > self.autocompact_threshold() {
48            return Some(MemoryCompactTrigger::FullThreshold);
49        }
50        self.should_preemptive_microcompact(used_tokens)
51            .then_some(MemoryCompactTrigger::MicroThreshold)
52    }
53
54    pub fn warning_threshold(&self) -> u64 {
55        let threshold = self.autocompact_threshold();
56        if threshold == 0 {
57            return 0;
58        }
59        (threshold * u64::from(self.config.warning_threshold_percentage)) / 100
60    }
61
62    pub(super) fn calculate_effective_length(
63        &self,
64        messages: &[Message],
65        total_tokens: Option<u64>,
66        recent_tool_call_ids: Option<&BTreeSet<String>>,
67    ) -> u64 {
68        if let Some(total_tokens) = total_tokens.filter(|tokens| *tokens > 0) {
69            return total_tokens
70                + self.estimate_recent_tool_message_length(messages, recent_tool_call_ids);
71        }
72        count_messages_tokens(messages, &self.config.model)
73    }
74
75    pub(crate) fn effective_length_for_cycle(
76        &self,
77        messages: &[Message],
78        total_tokens: Option<u64>,
79        recent_tool_call_ids: Option<&BTreeSet<String>>,
80    ) -> u64 {
81        self.calculate_effective_length(messages, total_tokens, recent_tool_call_ids)
82    }
83
84    fn estimate_recent_tool_message_length(
85        &self,
86        messages: &[Message],
87        recent_tool_call_ids: Option<&BTreeSet<String>>,
88    ) -> u64 {
89        let Some(recent_tool_call_ids) = recent_tool_call_ids.filter(|ids| !ids.is_empty()) else {
90            return 0;
91        };
92        let tool_messages = messages
93            .iter()
94            .filter(|message| {
95                message.role == MessageRole::Tool
96                    && message
97                        .tool_call_id
98                        .as_ref()
99                        .is_some_and(|tool_call_id| recent_tool_call_ids.contains(tool_call_id))
100            })
101            .cloned()
102            .collect::<Vec<_>>();
103        count_messages_tokens(&tool_messages, &self.config.model)
104    }
105}
vv_agent/memory/manager/limits.rs

vv_agent/memory/manager/
limits.rs
