Skip to main content

vv_agent/memory/
microcompact.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use serde_json::json;
4
5use crate::types::{Message, MessageRole};
6
7pub const CLEARED_MARKER: &str = "[Old tool result content cleared by microcompact]";
8
9pub const COMPACTABLE_TOOLS: &[&str] = &[
10    "read_file",
11    "write_file",
12    "edit_file",
13    "find_files",
14    "search_files",
15    "bash",
16    "file_info",
17];
18
19#[derive(Debug, Clone, PartialEq)]
20pub struct MicrocompactConfig {
21    pub trigger_ratio: f64,
22    pub keep_recent_cycles: usize,
23    pub min_result_length: usize,
24    pub compactable_tools: Option<BTreeSet<String>>,
25}
26
27impl Default for MicrocompactConfig {
28    fn default() -> Self {
29        Self {
30            trigger_ratio: 0.75,
31            keep_recent_cycles: 3,
32            min_result_length: 500,
33            compactable_tools: None,
34        }
35    }
36}
37
38pub fn microcompact(
39    messages: &[Message],
40    current_cycle: u32,
41    config: &MicrocompactConfig,
42) -> (Vec<Message>, usize) {
43    if messages.is_empty() {
44        return (Vec::new(), 0);
45    }
46
47    let tool_call_names = build_tool_call_name_map(messages);
48    let inferred_cycles = infer_message_cycles(messages);
49    let max_inferred_cycle = inferred_cycles.last().copied().unwrap_or_default();
50    let effective_current_cycle = current_cycle.min(max_inferred_cycle.saturating_add(1));
51    let protected_cycle = effective_current_cycle.saturating_sub(config.keep_recent_cycles as u32);
52    let compactable_tools = config
53        .compactable_tools
54        .clone()
55        .unwrap_or_else(default_compactable_tools);
56
57    let mut cleared = 0;
58    let mut updated = Vec::with_capacity(messages.len());
59    for (message, inferred_cycle) in messages.iter().zip(inferred_cycles) {
60        if should_clear_message(
61            message,
62            inferred_cycle,
63            protected_cycle,
64            config.min_result_length,
65            &compactable_tools,
66            &tool_call_names,
67        ) {
68            updated.push(replace_content(message));
69            cleared += 1;
70        } else {
71            updated.push(message.clone());
72        }
73    }
74    (updated, cleared)
75}
76
77pub fn is_microcompacted_tool_content(content: &str) -> bool {
78    content.starts_with(CLEARED_MARKER)
79}
80
81fn default_compactable_tools() -> BTreeSet<String> {
82    COMPACTABLE_TOOLS
83        .iter()
84        .map(|tool| (*tool).to_string())
85        .collect()
86}
87
88fn build_tool_call_name_map(messages: &[Message]) -> BTreeMap<String, String> {
89    let mut tool_call_names = BTreeMap::new();
90    for message in messages {
91        if message.role != MessageRole::Assistant {
92            continue;
93        }
94        for tool_call in &message.tool_calls {
95            tool_call_names.insert(tool_call.id.clone(), tool_call.name.clone());
96        }
97    }
98    tool_call_names
99}
100
101fn infer_message_cycles(messages: &[Message]) -> Vec<u32> {
102    let mut current_cycle = 0;
103    let mut inferred = Vec::with_capacity(messages.len());
104    for message in messages {
105        if message.role == MessageRole::Assistant {
106            current_cycle += 1;
107        }
108        inferred.push(current_cycle);
109    }
110    inferred
111}
112
113fn should_clear_message(
114    message: &Message,
115    inferred_cycle: u32,
116    protected_cycle: u32,
117    min_result_length: usize,
118    compactable_tools: &BTreeSet<String>,
119    tool_call_names: &BTreeMap<String, String>,
120) -> bool {
121    if message.role != MessageRole::Tool {
122        return false;
123    }
124    if inferred_cycle >= protected_cycle {
125        return false;
126    }
127    if message.content.len() <= min_result_length.max(1) {
128        return false;
129    }
130    if is_microcompacted_tool_content(&message.content) {
131        return false;
132    }
133    let Some(tool_call_id) = message.tool_call_id.as_deref() else {
134        return false;
135    };
136    let Some(tool_name) = tool_call_names.get(tool_call_id) else {
137        return false;
138    };
139    compactable_tools.contains(tool_name)
140}
141
142fn replace_content(message: &Message) -> Message {
143    let mut updated = message.clone();
144    updated
145        .metadata
146        .insert("microcompacted".to_string(), json!(true));
147    updated.metadata.insert(
148        "microcompact_original_chars".to_string(),
149        json!(message.content.len()),
150    );
151    updated.content = CLEARED_MARKER.to_string();
152    updated
153}