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    fn estimate_recent_tool_message_length(
76        &self,
77        messages: &[Message],
78        recent_tool_call_ids: Option<&BTreeSet<String>>,
79    ) -> u64 {
80        let Some(recent_tool_call_ids) = recent_tool_call_ids.filter(|ids| !ids.is_empty()) else {
81            return 0;
82        };
83        let tool_messages = messages
84            .iter()
85            .filter(|message| {
86                message.role == MessageRole::Tool
87                    && message
88                        .tool_call_id
89                        .as_ref()
90                        .is_some_and(|tool_call_id| recent_tool_call_ids.contains(tool_call_id))
91            })
92            .cloned()
93            .collect::<Vec<_>>();
94        count_messages_tokens(&tool_messages, &self.config.model)
95    }
96}
vv_agent/memory/manager/limits.rs

vv_agent/memory/manager/
limits.rs
