vv_agent/memory/manager/
limits.rs

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

vv_agent/memory/manager/
limits.rs
