otherone_context/compact/
mod.rs

1// 作用：上下文压缩子模块
2// 关联：被 otherone-context 的 combine_context 调用
3// 预期结果：提供 token 估算、阈值检查、消息序列化、压缩 prompts、压缩执行
4
5pub mod check_threshold;
6pub mod estimate_tokens;
7pub mod messages_to_sequence;
8pub mod prompts;
9
10use crate::compact::estimate_tokens::estimate_tokens;
11use crate::compact::messages_to_sequence::messages_to_sequence;
12use crate::compact::prompts::{
13    SUMMARIZATION_SYSTEM_PROMPT, TURN_PREFIX_SUMMARIZATION_PROMPT, UPDATE_SUMMARIZATION_PROMPT,
14};
15use crate::error::ContextError;
16use otherone_ai::types::{Message, MessageContent, ProviderType};
17use otherone_storage::types::{StorageType, WriteCompactedEntryOptions};
18
19/// 压缩上下文消息
20/// 作用：保留最新的消息，压缩旧的消息为一段摘要
21/// 关联：被 combine_context 调用，当 token 使用量超过阈值时触发
22/// 预期结果：返回压缩后的 messages 数组：[compacted_message, ...messages_to_keep]
23pub async fn compact_messages(
24    messages: &[Message],
25    _context_tokens: u32,
26    context_window: u32,
27    compact_ratio: Option<f32>,
28    ai_config: Option<&serde_json::Value>,
29    has_compacted_content: bool,
30    session_id: Option<&str>,
31    storage_type: Option<&StorageType>,
32    database_config: Option<&otherone_storage::types::DatabaseConfig>,
33    original_entries: Option<&[otherone_storage::types::Entry]>,
34) -> Result<Vec<Message>, ContextError> {
35    if messages.is_empty() {
36        return Ok(Vec::new());
37    }
38
39    // 提取压缩 LLM 配置
40    let compact_llm = extract_compact_llm_config(ai_config).ok_or_else(|| {
41        ContextError::ConfigError("AI configuration is required for compaction".to_string())
42    })?;
43
44    // 设置默认保留比例为 40%
45    let keep_ratio = compact_ratio.unwrap_or(0.4);
46    let keep_token_threshold = (context_window as f32 * keep_ratio) as u32;
47
48    // 从后往前查找切割点
49    let mut accumulated_tokens: u32 = 0;
50    let mut cutoff_index: usize = 0;
51
52    for i in (0..messages.len()).rev() {
53        let message_tokens = estimate_tokens(&[messages[i].clone()]);
54        if accumulated_tokens + message_tokens <= keep_token_threshold {
55            accumulated_tokens += message_tokens;
56            cutoff_index = i;
57        } else {
58            break;
59        }
60    }
61
62    // 如果所有消息都在阈值内，不需要压缩
63    if cutoff_index == 0 {
64        return Ok(messages.to_vec());
65    }
66
67    // 根据切割点消息的 role 类型调整切割位置
68    let cutoff_message = &messages[cutoff_index];
69    let adjusted_cutoff = if cutoff_message.role == "assistant" {
70        // 从切割点往前查找第一条 user 消息
71        let mut found = cutoff_index;
72        for i in (0..cutoff_index).rev() {
73            if messages[i].role == "user" {
74                found = i + 1;
75                break;
76            }
77        }
78        found
79    } else {
80        cutoff_index
81    };
82
83    // 分割消息数组
84    let messages_to_compact = &messages[..adjusted_cutoff];
85    let messages_to_keep = &messages[adjusted_cutoff..];
86
87    if messages_to_compact.is_empty() {
88        return Ok(messages.to_vec());
89    }
90
91    // 调用 LLM 进行压缩
92    let compressed_summary =
93        call_compact_llm(messages_to_compact, &compact_llm, has_compacted_content).await?;
94
95    // 创建压缩摘要消息
96    let compacted_message = Message {
97        role: "user".to_string(),
98        content: MessageContent::Text(compressed_summary.clone()),
99        name: None,
100        tool_calls: None,
101        tool_call_id: None,
102    };
103
104    // 存储压缩记录
105    if let (Some(sid), Some(st), Some(entries)) = (session_id, storage_type, original_entries) {
106        let trigger_entry_id = if has_compacted_content && adjusted_cutoff > 0 {
107            entries
108                .get(adjusted_cutoff - 1)
109                .map(|e| e.entry_id.clone())
110                .unwrap_or_default()
111        } else {
112            entries
113                .get(adjusted_cutoff - 1)
114                .map(|e| e.entry_id.clone())
115                .unwrap_or_default()
116        };
117
118        if !trigger_entry_id.is_empty() {
119            let _ = otherone_storage::write_compacted_entry(&WriteCompactedEntryOptions {
120                storage_type: st.clone(),
121                session_id: sid.to_string(),
122                summary: compressed_summary,
123                trigger_entry_id,
124                create_at: None,
125                database_config: database_config.cloned(),
126            })
127            .await;
128        }
129    }
130
131    // 返回压缩后的消息数组
132    let mut result = vec![compacted_message];
133    result.extend_from_slice(messages_to_keep);
134    Ok(result)
135}
136
137/// 调用 LLM 进行上下文压缩
138/// 作用：调用 AI 模块将旧消息压缩为一段摘要
139/// 关联：被 compact_messages 调用
140/// 预期结果：返回压缩后的摘要文本
141async fn call_compact_llm(
142    messages_to_compact: &[Message],
143    compact_llm_config: &CompactLLMConfig,
144    has_compacted_content: bool,
145) -> Result<String, ContextError> {
146    // 将消息转换为序列文本
147    let message_sequence = messages_to_sequence(messages_to_compact);
148
149    // 根据是否已有压缩内容选择不同的 prompt
150    let user_prompt = if has_compacted_content {
151        let previous_summary = extract_message_text(&messages_to_compact[0]);
152        format!(
153            "<previous-summary>\n{}\n</previous-summary>\n\n{}\n\n{}",
154            previous_summary, UPDATE_SUMMARIZATION_PROMPT, message_sequence
155        )
156    } else {
157        format!(
158            "{}\n\n{}",
159            TURN_PREFIX_SUMMARIZATION_PROMPT, message_sequence
160        )
161    };
162
163    // 构建压缩请求的 messages
164    let compact_messages = vec![
165        Message {
166            role: "system".to_string(),
167            content: MessageContent::Text(SUMMARIZATION_SYSTEM_PROMPT.to_string()),
168            name: None,
169            tool_calls: None,
170            tool_call_id: None,
171        },
172        Message {
173            role: "user".to_string(),
174            content: MessageContent::Text(user_prompt),
175            name: None,
176            tool_calls: None,
177            tool_call_id: None,
178        },
179    ];
180
181    // 构建 AI 配置并调用
182    let ai_config = serde_json::json!({
183        "model": compact_llm_config.model,
184        "messages": compact_messages,
185    });
186
187    let response = otherone_ai::invoke_model(
188        compact_llm_config.provider.clone(),
189        &compact_llm_config.api_key,
190        &compact_llm_config.base_url,
191        ai_config,
192    )
193    .await
194    .map_err(|e| ContextError::CompactionError(e.to_string()))?;
195
196    // 提取压缩内容
197    let compressed = match compact_llm_config.provider {
198        ProviderType::OpenAI
199        | ProviderType::Fetch
200        | ProviderType::OpenRouter
201        | ProviderType::Local => response
202            .choices
203            .first()
204            .and_then(|c| c.message.as_ref())
205            .and_then(|m| m.content.as_deref())
206            .unwrap_or("")
207            .to_string(),
208        ProviderType::Anthropic => response
209            .choices
210            .first()
211            .and_then(|c| c.message.as_ref())
212            .and_then(|m| m.content.as_deref())
213            .unwrap_or("")
214            .to_string(),
215    };
216
217    if compressed.is_empty() {
218        return Err(ContextError::CompactionError(
219            "Unable to extract content from compaction response".to_string(),
220        ));
221    }
222
223    Ok(compressed)
224}
225
226/// 压缩 LLM 配置
227struct CompactLLMConfig {
228    provider: ProviderType,
229    api_key: String,
230    base_url: String,
231    model: String,
232}
233
234/// 从 AI 配置中提取压缩 LLM 配置
235fn extract_compact_llm_config(ai_config: Option<&serde_json::Value>) -> Option<CompactLLMConfig> {
236    let config = ai_config?;
237    let obj = config.as_object()?;
238
239    let provider_str = obj
240        .get("compact_llm_provider")
241        .or_else(|| obj.get("provider"))
242        .and_then(|v| v.as_str())
243        .unwrap_or("openai");
244
245    let provider = match provider_str {
246        "anthropic" => ProviderType::Anthropic,
247        "fetch" => ProviderType::Fetch,
248        _ => ProviderType::OpenAI,
249    };
250
251    let api_key: String = obj
252        .get("compact_llm_apiKey")
253        .or_else(|| obj.get("apiKey"))
254        .or_else(|| obj.get("api_key"))
255        .and_then(|v| v.as_str())
256        .filter(|s| !s.is_empty())
257        .map(|s| s.to_string())?;
258
259    let base_url: String = obj
260        .get("compact_llm_baseUrl")
261        .or_else(|| obj.get("baseUrl"))
262        .or_else(|| obj.get("base_url"))
263        .and_then(|v| v.as_str())
264        .filter(|s| !s.is_empty())
265        .map(|s| s.to_string())?;
266
267    let model: String = obj
268        .get("compact_llm_model")
269        .or_else(|| obj.get("model"))
270        .and_then(|v| v.as_str())
271        .filter(|s| !s.is_empty())
272        .map(|s| s.to_string())?;
273
274    Some(CompactLLMConfig {
275        provider,
276        api_key,
277        base_url,
278        model,
279    })
280}
281
282/// 提取 MessageContent 中的文本
283fn extract_message_text(message: &Message) -> String {
284    match &message.content {
285        MessageContent::Text(t) => t.clone(),
286        _ => "[非文本内容]".to_string(),
287    }
288}
otherone_context/compact/mod.rs

otherone_context/compact/
mod.rs
