Skip to main content

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(crate) fn should_attempt_compaction(
39        &self,
40        messages: &[Message],
41        total_tokens: Option<u64>,
42        recent_tool_call_ids: Option<&BTreeSet<String>>,
43    ) -> bool {
44        let used_tokens =
45            self.calculate_effective_length(messages, total_tokens, recent_tool_call_ids);
46        self.should_preemptive_microcompact(used_tokens)
47            || used_tokens > self.autocompact_threshold()
48    }
49
50    pub fn warning_threshold(&self) -> u64 {
51        let threshold = self.autocompact_threshold();
52        if threshold == 0 {
53            return 0;
54        }
55        (threshold * u64::from(self.config.warning_threshold_percentage)) / 100
56    }
57
58    pub(super) fn calculate_effective_length(
59        &self,
60        messages: &[Message],
61        total_tokens: Option<u64>,
62        recent_tool_call_ids: Option<&BTreeSet<String>>,
63    ) -> u64 {
64        if let Some(total_tokens) = total_tokens.filter(|tokens| *tokens > 0) {
65            return total_tokens
66                + self.estimate_recent_tool_message_length(messages, recent_tool_call_ids);
67        }
68        count_messages_tokens(messages, &self.config.model)
69    }
70
71    fn estimate_recent_tool_message_length(
72        &self,
73        messages: &[Message],
74        recent_tool_call_ids: Option<&BTreeSet<String>>,
75    ) -> u64 {
76        let Some(recent_tool_call_ids) = recent_tool_call_ids.filter(|ids| !ids.is_empty()) else {
77            return 0;
78        };
79        let tool_messages = messages
80            .iter()
81            .filter(|message| {
82                message.role == MessageRole::Tool
83                    && message
84                        .tool_call_id
85                        .as_ref()
86                        .is_some_and(|tool_call_id| recent_tool_call_ids.contains(tool_call_id))
87            })
88            .cloned()
89            .collect::<Vec<_>>();
90        count_messages_tokens(&tool_messages, &self.config.model)
91    }
92}