Skip to main content

lellm_agent/runtime/context/
budget.rs

1//! 上下文预算配置 — 控制 Agent Loop 中 messages 的 Token 总量。
2
3use lellm_core::ContentBlock;
4
5/// 上下文预算配置。
6///
7/// 控制 Agent Loop 中消息历史的 Token 上限与压缩行为。
8#[derive(Debug, Clone)]
9pub struct ContextBudget {
10    /// 消息历史的最大 Token 数(默认 128k)。
11    ///
12    /// **v0.1**: 固定默认值 128k,适用于大多数模型
13    /// **v0.2**: 从 `ResolvedModel.context_window` 自动推导(window * 0.8)
14    pub max_tokens: usize,
15    /// 达到此占比时触发压缩(默认 0.8 = 80%)
16    pub warning_ratio: f32,
17    /// 压缩时保留最近多少个 Turn(默认 5)
18    pub keep_recent_turns: usize,
19    /// 单条工具结果的最大字符数(默认 4096)
20    pub max_tool_result_chars: usize,
21}
22
23impl Default for ContextBudget {
24    fn default() -> Self {
25        Self {
26            max_tokens: 128_000,
27            warning_ratio: 0.8,
28            keep_recent_turns: 5,
29            max_tool_result_chars: 4096,
30        }
31    }
32}
33
34impl ContextBudget {
35    /// 判断是否需要压缩。
36    pub fn should_compact(&self, current_tokens: usize) -> bool {
37        let threshold = (self.max_tokens as f32 * self.warning_ratio) as usize;
38        current_tokens > threshold
39    }
40
41    /// 截断单条工具结果,防止单条响应撑爆上下文。
42    pub fn truncate_tool_result(&self, text: String) -> String {
43        if text.chars().count() <= self.max_tool_result_chars {
44            return text;
45        }
46        let truncated: String = text.chars().take(self.max_tool_result_chars).collect();
47        format!(
48            "{}\n[truncated, original {} chars]",
49            truncated,
50            text.chars().count()
51        )
52    }
53
54    /// 截断 ContentBlock 列表中的文本内容(用于非流式路径)。
55    pub fn truncate_tool_result_blocks(&self, content: &[ContentBlock]) -> Vec<ContentBlock> {
56        content
57            .iter()
58            .map(|b| match b {
59                ContentBlock::Text(t) => {
60                    let truncated = self.truncate_tool_result(t.text.clone());
61                    if truncated != t.text {
62                        ContentBlock::Text(lellm_core::TextBlock {
63                            text: truncated,
64                            cache_control: None,
65                        })
66                    } else {
67                        b.clone()
68                    }
69                }
70                _ => b.clone(),
71            })
72            .collect()
73    }
74}