feat: 实现两段历史压缩的同步处理,替换会话的活动历史记录
This commit is contained in:
parent
1af0fab3ad
commit
51e06c8f73
@ -5,7 +5,6 @@ use crate::bus::{
|
|||||||
use crate::config::LLMProviderConfig;
|
use crate::config::LLMProviderConfig;
|
||||||
use crate::providers::{ChatCompletionRequest, LLMProvider, Message, create_provider};
|
use crate::providers::{ChatCompletionRequest, LLMProvider, Message, create_provider};
|
||||||
use crate::text::{char_count, take_prefix_chars};
|
use crate::text::{char_count, take_prefix_chars};
|
||||||
|
|
||||||
use crate::agent::{AgentError, AgentRuntimeConfig};
|
use crate::agent::{AgentError, AgentRuntimeConfig};
|
||||||
|
|
||||||
const TOKEN_ESTIMATE_SAFETY_MULTIPLIER: f64 = 1.2;
|
const TOKEN_ESTIMATE_SAFETY_MULTIPLIER: f64 = 1.2;
|
||||||
@ -13,6 +12,145 @@ const CJK_CHARS_PER_TOKEN: f64 = 2.0;
|
|||||||
const OTHER_CHARS_PER_TOKEN: f64 = 4.0;
|
const OTHER_CHARS_PER_TOKEN: f64 = 4.0;
|
||||||
const JSON_OVERHEAD_PER_MESSAGE: usize = 50;
|
const JSON_OVERHEAD_PER_MESSAGE: usize = 50;
|
||||||
|
|
||||||
|
/// System context marker for the heavy compression (older segment) summary.
|
||||||
|
pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER: &str = "history_compaction_older";
|
||||||
|
/// System context marker for the light compression (newer segment) summary.
|
||||||
|
pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER: &str = "history_compaction_newer";
|
||||||
|
|
||||||
|
/// Default threshold ratio: compress when estimated tokens exceed 70% of context window.
|
||||||
|
const DEFAULT_THRESHOLD_RATIO: f64 = 0.7;
|
||||||
|
|
||||||
|
/// Default budget split: older segment gets 30% of summary budget, newer gets 70%.
|
||||||
|
const OLDER_BUDGET_RATIO: f64 = 0.3;
|
||||||
|
const NEWER_BUDGET_RATIO: f64 = 0.7;
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// HistoryUnit — atomic message units for compression
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
/// An indivisible unit of conversation history. Compression operates at the
|
||||||
|
/// unit level, so tool call sequences can never be split across boundaries.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
enum HistoryUnit {
|
||||||
|
/// Always-preserved system messages (agent_prompt, scheduled_prompt).
|
||||||
|
SystemGuard(ChatMessage),
|
||||||
|
/// A user message.
|
||||||
|
UserMessage(ChatMessage),
|
||||||
|
/// An atomic tool call sequence: assistant(tool_calls) + all following
|
||||||
|
/// tool-result messages. Must be preserved or compressed as a whole.
|
||||||
|
ToolRound {
|
||||||
|
assistant: ChatMessage,
|
||||||
|
results: Vec<ChatMessage>,
|
||||||
|
},
|
||||||
|
/// A plain assistant text response (no tool_calls).
|
||||||
|
AssistantText(ChatMessage),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl HistoryUnit {
|
||||||
|
/// Estimate tokens for this unit alone.
|
||||||
|
fn estimate_tokens(&self) -> usize {
|
||||||
|
match self {
|
||||||
|
HistoryUnit::SystemGuard(msg) | HistoryUnit::UserMessage(msg) | HistoryUnit::AssistantText(msg) => {
|
||||||
|
estimate_tokens(std::slice::from_ref(msg))
|
||||||
|
}
|
||||||
|
HistoryUnit::ToolRound { assistant, results } => {
|
||||||
|
let mut all = vec![assistant.clone()];
|
||||||
|
all.extend(results.clone());
|
||||||
|
estimate_tokens(&all)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Expand this unit back into a flat list of ChatMessages.
|
||||||
|
#[allow(dead_code)]
|
||||||
|
fn into_messages(self) -> Vec<ChatMessage> {
|
||||||
|
match self {
|
||||||
|
HistoryUnit::SystemGuard(msg)
|
||||||
|
| HistoryUnit::UserMessage(msg)
|
||||||
|
| HistoryUnit::AssistantText(msg) => vec![msg],
|
||||||
|
HistoryUnit::ToolRound { assistant, results } => {
|
||||||
|
let mut msgs = vec![assistant];
|
||||||
|
msgs.extend(results);
|
||||||
|
msgs
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns true if this unit is a ToolRound.
|
||||||
|
#[allow(dead_code)]
|
||||||
|
fn is_tool_round(&self) -> bool {
|
||||||
|
matches!(self, HistoryUnit::ToolRound { .. })
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Unit parser — one forward pass, O(n)
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
fn is_system_guard(msg: &ChatMessage) -> bool {
|
||||||
|
msg.role == "system"
|
||||||
|
&& (msg.has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT)
|
||||||
|
|| msg.has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_assistant_with_tool_calls(msg: &ChatMessage) -> bool {
|
||||||
|
msg.role == "assistant"
|
||||||
|
&& msg
|
||||||
|
.tool_calls
|
||||||
|
.as_ref()
|
||||||
|
.map_or(false, |calls| !calls.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse a flat message list into atomic units. Orphaned tool results
|
||||||
|
/// (without a preceding assistant with tool_calls) are silently dropped.
|
||||||
|
fn parse_to_units(messages: &[ChatMessage]) -> Vec<HistoryUnit> {
|
||||||
|
let mut units = Vec::new();
|
||||||
|
let mut i = 0;
|
||||||
|
|
||||||
|
while i < messages.len() {
|
||||||
|
let msg = &messages[i];
|
||||||
|
|
||||||
|
if is_system_guard(msg) {
|
||||||
|
units.push(HistoryUnit::SystemGuard(msg.clone()));
|
||||||
|
i += 1;
|
||||||
|
} else if msg.role == "user" {
|
||||||
|
units.push(HistoryUnit::UserMessage(msg.clone()));
|
||||||
|
i += 1;
|
||||||
|
} else if is_assistant_with_tool_calls(msg) {
|
||||||
|
let assistant = msg.clone();
|
||||||
|
let mut results = Vec::new();
|
||||||
|
i += 1;
|
||||||
|
while i < messages.len() && messages[i].role == "tool" {
|
||||||
|
results.push(messages[i].clone());
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
units.push(HistoryUnit::ToolRound { assistant, results });
|
||||||
|
} else if msg.role == "assistant" {
|
||||||
|
units.push(HistoryUnit::AssistantText(msg.clone()));
|
||||||
|
i += 1;
|
||||||
|
} else if msg.role == "tool" {
|
||||||
|
// Orphaned tool result — drop with a warning
|
||||||
|
tracing::warn!(
|
||||||
|
tool_call_id = ?msg.tool_call_id,
|
||||||
|
message_index = i,
|
||||||
|
"Dropping orphaned tool result during unit parsing"
|
||||||
|
);
|
||||||
|
i += 1;
|
||||||
|
} else {
|
||||||
|
// system messages that aren't guards — keep as plain messages
|
||||||
|
units.push(HistoryUnit::AssistantText(msg.clone()));
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
units
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Estimate total tokens from a slice of units.
|
||||||
|
fn estimate_tokens_from_units(units: &[HistoryUnit]) -> usize {
|
||||||
|
units.iter().map(|u| u.estimate_tokens()).sum()
|
||||||
|
}
|
||||||
|
|
||||||
/// Check if a character is CJK (Chinese, Japanese, Korean)
|
/// Check if a character is CJK (Chinese, Japanese, Korean)
|
||||||
fn is_cjk_char(c: char) -> bool {
|
fn is_cjk_char(c: char) -> bool {
|
||||||
matches!(c,
|
matches!(c,
|
||||||
@ -97,7 +235,7 @@ impl Default for ContextCompressionConfig {
|
|||||||
pub struct ContextCompressor {
|
pub struct ContextCompressor {
|
||||||
config: ContextCompressionConfig,
|
config: ContextCompressionConfig,
|
||||||
context_window: usize,
|
context_window: usize,
|
||||||
/// Threshold ratio to trigger compression (50% of context window)
|
/// Threshold ratio to trigger compression (70% of context window).
|
||||||
threshold_ratio: f64,
|
threshold_ratio: f64,
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -113,6 +251,10 @@ impl ContextCompressor {
|
|||||||
.clamp(MIN_SUMMARY_CHARS, MAX_SUMMARY_CHARS)
|
.clamp(MIN_SUMMARY_CHARS, MAX_SUMMARY_CHARS)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// =========================================================================
|
||||||
|
// Transcript building (shared helpers, unchanged from original)
|
||||||
|
// =========================================================================
|
||||||
|
|
||||||
fn format_transcript_entry(message: &ChatMessage) -> String {
|
fn format_transcript_entry(message: &ChatMessage) -> String {
|
||||||
let role = match message.role.as_str() {
|
let role = match message.role.as_str() {
|
||||||
"assistant" => "Assistant",
|
"assistant" => "Assistant",
|
||||||
@ -214,6 +356,86 @@ Be concise, aim for {} characters or less.
|
|||||||
|
|
||||||
{}
|
{}
|
||||||
|
|
||||||
|
"#,
|
||||||
|
target_chars, transcript
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// =========================================================================
|
||||||
|
// Two-segment compression: heavy (older) + light (newer) prompts
|
||||||
|
// =========================================================================
|
||||||
|
|
||||||
|
/// Prompt for the older segment — aggressive compression, keep only essentials.
|
||||||
|
fn build_heavy_summary_prompt(transcript: &str, target_chars: usize) -> String {
|
||||||
|
format!(
|
||||||
|
r#"You are a conversation compaction engine. Aggressively summarize the following OLDER conversation segment. These events happened earlier in the session and will be marked as "较早的操作记录(已压缩)".
|
||||||
|
|
||||||
|
A newer portion (not shown here) follows this segment. The model can use tools to re-read files for exact data, so you do NOT need to preserve full tool outputs.
|
||||||
|
|
||||||
|
=== MUST PRESERVE (keep verbatim or near-verbatim) ===
|
||||||
|
- All file paths (absolute and relative)
|
||||||
|
- All URLs
|
||||||
|
- Key numeric results and conclusions (one sentence per result)
|
||||||
|
- Error messages (exact text if important)
|
||||||
|
- User preferences and decisions
|
||||||
|
- The original user task or goal
|
||||||
|
|
||||||
|
=== SHOULD CONDENSE ===
|
||||||
|
- Each tool call → one line summary: "tool_name: 关键参数 → 结果概要"
|
||||||
|
- Multiple similar operations → group: "对 ./data/ 目录下的 10 个 CSV 文件执行了统计分析"
|
||||||
|
- Verbose tool output → drop entirely unless it contains a critical result or error
|
||||||
|
|
||||||
|
=== SHOULD DROP ===
|
||||||
|
- Full tool output (hundreds of lines of raw data)
|
||||||
|
- Repeated attempts of the same operation
|
||||||
|
- Greetings, filler text, markdown formatting
|
||||||
|
|
||||||
|
Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
OLDER SEGMENT (events from earlier in the session):
|
||||||
|
|
||||||
|
{}
|
||||||
|
|
||||||
|
"#,
|
||||||
|
target_chars, transcript
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Prompt for the newer segment — lighter compression, keep more detail.
|
||||||
|
fn build_light_summary_prompt(transcript: &str, target_chars: usize) -> String {
|
||||||
|
format!(
|
||||||
|
r#"You are a conversation compaction engine. Lightly summarize the following RECENT conversation segment. These events happened just before the current moment and will be marked as "近期操作记录(轻度压缩,保留了更多细节)".
|
||||||
|
|
||||||
|
An older segment (already summarized separately) precedes this. Keep enough detail so the model can continue the task seamlessly without re-reading files for data that appears in these recent results.
|
||||||
|
|
||||||
|
=== MUST PRESERVE (keep with more detail than a normal summary) ===
|
||||||
|
- All file paths, URLs, and identifiers
|
||||||
|
- The exact sequence of recent operations (step by step)
|
||||||
|
- Tool parameters (especially file paths, search queries, command strings)
|
||||||
|
- Key outputs and results (shortened but keep the substance)
|
||||||
|
- Error messages (complete, not truncated)
|
||||||
|
- Current task status and what should happen next
|
||||||
|
|
||||||
|
=== SHOULD CONDENSE ===
|
||||||
|
- Tool outputs → truncate to the most meaningful parts (key data, conclusions, not raw output)
|
||||||
|
- Long text → keep the essence but not every word
|
||||||
|
- Repeated similar outputs → note the pattern but keep a representative example
|
||||||
|
|
||||||
|
=== SHOULD DROP ===
|
||||||
|
- Completely irrelevant debug output
|
||||||
|
- Boilerplate text with no informational value
|
||||||
|
- Trivial operations that have no bearing on the task
|
||||||
|
|
||||||
|
Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
RECENT SEGMENT (events from just before the current moment):
|
||||||
|
|
||||||
|
{}
|
||||||
|
|
||||||
"#,
|
"#,
|
||||||
target_chars, transcript
|
target_chars, transcript
|
||||||
)
|
)
|
||||||
@ -295,12 +517,339 @@ Be concise, aim for {} characters or less.
|
|||||||
Ok(take_prefix_chars(transcript, target_chars))
|
Ok(take_prefix_chars(transcript, target_chars))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// =========================================================================
|
||||||
|
// Two-segment compression
|
||||||
|
// =========================================================================
|
||||||
|
|
||||||
|
/// Find a safe split point in the unit list. The split ensures:
|
||||||
|
/// - Accumulated tokens up to `ratio` of total are placed in the older segment.
|
||||||
|
/// - The split never lands on a SystemGuard (they're always preserved).
|
||||||
|
/// - The split lands on a unit boundary (ToolRound, AssistantText, or UserMessage).
|
||||||
|
fn find_safe_split_point(&self, units: &[HistoryUnit], ratio: f64) -> usize {
|
||||||
|
let total_tokens = estimate_tokens_from_units(units);
|
||||||
|
let target = (total_tokens as f64 * ratio) as usize;
|
||||||
|
let mut accumulated = 0;
|
||||||
|
|
||||||
|
for (i, unit) in units.iter().enumerate() {
|
||||||
|
// Never split on a SystemGuard — they must be preserved
|
||||||
|
if matches!(unit, HistoryUnit::SystemGuard(_)) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
accumulated += unit.estimate_tokens();
|
||||||
|
if accumulated >= target {
|
||||||
|
return i + 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fallback: put everything in the older segment
|
||||||
|
units.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Summarize a transcript with a custom prompt builder function.
|
||||||
|
async fn summarize_with_prompt(
|
||||||
|
&self,
|
||||||
|
provider: &dyn LLMProvider,
|
||||||
|
transcript: &str,
|
||||||
|
target_chars: usize,
|
||||||
|
build_prompt: fn(&str, usize) -> String,
|
||||||
|
) -> Result<String, AgentError> {
|
||||||
|
let request = ChatCompletionRequest {
|
||||||
|
messages: vec![
|
||||||
|
Message::system("You are a helpful assistant."),
|
||||||
|
Message::user(build_prompt(transcript, target_chars)),
|
||||||
|
],
|
||||||
|
temperature: Some(0.3),
|
||||||
|
max_tokens: Some(1000),
|
||||||
|
tools: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let response = provider
|
||||||
|
.chat(request)
|
||||||
|
.await
|
||||||
|
.map_err(|e| AgentError::LlmError(e.to_string()))?;
|
||||||
|
Ok(response.content)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Summarize a segment of units with chunked fallback.
|
||||||
|
async fn summarize_units_segment(
|
||||||
|
&self,
|
||||||
|
provider: &dyn LLMProvider,
|
||||||
|
messages: &[ChatMessage],
|
||||||
|
transcript: &str,
|
||||||
|
target_chars: usize,
|
||||||
|
build_prompt: fn(&str, usize) -> String,
|
||||||
|
) -> Result<String, AgentError> {
|
||||||
|
if messages.is_empty() {
|
||||||
|
return Ok(String::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
if char_count(transcript) <= target_chars {
|
||||||
|
return self
|
||||||
|
.summarize_with_prompt(provider, transcript, target_chars, build_prompt)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Use chunked summarization for oversized transcripts
|
||||||
|
let target = target_chars.max(1);
|
||||||
|
let mut layer = Self::chunk_messages_for_summary(messages, target);
|
||||||
|
|
||||||
|
if layer.is_empty() {
|
||||||
|
layer.push(transcript.to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
for _ in 0..6 {
|
||||||
|
if layer.len() == 1 && char_count(&layer[0]) <= target {
|
||||||
|
return self
|
||||||
|
.summarize_with_prompt(provider, &layer[0], target, build_prompt)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let per_chunk_target = (target / layer.len().max(1)).max(500).min(target);
|
||||||
|
let mut summaries = Vec::with_capacity(layer.len());
|
||||||
|
for chunk in &layer {
|
||||||
|
summaries.push(
|
||||||
|
self.summarize_with_prompt(provider, chunk, per_chunk_target, build_prompt)
|
||||||
|
.await?,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
if summaries.len() == 1 {
|
||||||
|
let summary = summaries.pop().unwrap_or_default();
|
||||||
|
if char_count(&summary) <= target {
|
||||||
|
return Ok(summary);
|
||||||
|
}
|
||||||
|
layer = Self::split_text_chunks(&summary, target);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let merged = summaries.join("\n\n");
|
||||||
|
if char_count(&merged) <= target {
|
||||||
|
return self
|
||||||
|
.summarize_with_prompt(provider, &merged, target, build_prompt)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
layer = Self::split_text_chunks(&merged, target);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(take_prefix_chars(transcript, target))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Main entry point for two-segment compression.
|
||||||
|
///
|
||||||
|
/// Splits history into older and newer segments, then summarizes each
|
||||||
|
/// with LLM using different prompts and budget allocations. The result
|
||||||
|
/// contains no tool_calls or tool-result messages — only system summaries
|
||||||
|
/// and the original user message — eliminating any risk of API 400 errors
|
||||||
|
/// from orphaned tool call sequences.
|
||||||
|
pub async fn compress_two_segment(
|
||||||
|
&self,
|
||||||
|
history: &[ChatMessage],
|
||||||
|
provider_config: &LLMProviderConfig,
|
||||||
|
) -> Result<Vec<ChatMessage>, AgentError> {
|
||||||
|
let tokens = estimate_tokens(history);
|
||||||
|
if tokens <= self.threshold() {
|
||||||
|
tracing::info!(
|
||||||
|
tokens = tokens,
|
||||||
|
threshold = self.threshold(),
|
||||||
|
msg_count = history.len(),
|
||||||
|
"Two-segment compression not needed (under threshold)"
|
||||||
|
);
|
||||||
|
return Ok(history.to_vec());
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
tokens = tokens,
|
||||||
|
threshold = self.threshold(),
|
||||||
|
msg_count = history.len(),
|
||||||
|
"Starting two-segment compression"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Step 1: Parse into atomic units
|
||||||
|
let units = parse_to_units(history);
|
||||||
|
|
||||||
|
// Step 2: Separate system guards + user messages from compressible units
|
||||||
|
let mut system_guards: Vec<ChatMessage> = Vec::new();
|
||||||
|
let mut user_messages: Vec<ChatMessage> = Vec::new();
|
||||||
|
let mut compressible: Vec<HistoryUnit> = Vec::new();
|
||||||
|
|
||||||
|
for unit in units {
|
||||||
|
match unit {
|
||||||
|
HistoryUnit::SystemGuard(msg) => system_guards.push(msg),
|
||||||
|
HistoryUnit::UserMessage(msg) => user_messages.push(msg),
|
||||||
|
other => compressible.push(other),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Keep only the latest user message in full; older ones go into compression
|
||||||
|
let latest_user_msg = user_messages.pop();
|
||||||
|
|
||||||
|
// Step 3: Find split point in compressible units
|
||||||
|
let total_compressible = estimate_tokens_from_units(&compressible);
|
||||||
|
let split_point = if total_compressible > 0 {
|
||||||
|
self.find_safe_split_point(&compressible, 0.5)
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
let older_units: Vec<&HistoryUnit> = compressible[..split_point].iter().collect();
|
||||||
|
let newer_units: Vec<&HistoryUnit> = compressible[split_point..].iter().collect();
|
||||||
|
|
||||||
|
// Step 4: Build transcripts
|
||||||
|
let older_messages: Vec<ChatMessage> = older_units
|
||||||
|
.iter()
|
||||||
|
.flat_map(|u| match u {
|
||||||
|
HistoryUnit::ToolRound { assistant, results } => {
|
||||||
|
let mut msgs = vec![assistant.clone()];
|
||||||
|
msgs.extend(results.clone());
|
||||||
|
msgs
|
||||||
|
}
|
||||||
|
HistoryUnit::AssistantText(msg) => vec![msg.clone()],
|
||||||
|
HistoryUnit::UserMessage(msg) => {
|
||||||
|
// Older user messages go into the compressible transcript
|
||||||
|
vec![msg.clone()]
|
||||||
|
}
|
||||||
|
_ => vec![],
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let newer_messages: Vec<ChatMessage> = newer_units
|
||||||
|
.iter()
|
||||||
|
.flat_map(|u| match u {
|
||||||
|
HistoryUnit::ToolRound { assistant, results } => {
|
||||||
|
let mut msgs = vec![assistant.clone()];
|
||||||
|
msgs.extend(results.clone());
|
||||||
|
msgs
|
||||||
|
}
|
||||||
|
HistoryUnit::AssistantText(msg) => vec![msg.clone()],
|
||||||
|
_ => vec![],
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Include older user messages in the older transcript
|
||||||
|
let older_user_msgs: Vec<ChatMessage> = user_messages
|
||||||
|
.iter()
|
||||||
|
.map(|m| {
|
||||||
|
let mut msg = m.clone();
|
||||||
|
msg.role = "user".to_string();
|
||||||
|
msg
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let all_older_messages: Vec<ChatMessage> = older_user_msgs
|
||||||
|
.iter()
|
||||||
|
.chain(older_messages.iter())
|
||||||
|
.cloned()
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let older_transcript = Self::build_transcript(&all_older_messages);
|
||||||
|
let newer_transcript = Self::build_transcript(&newer_messages);
|
||||||
|
|
||||||
|
// Step 5: Budget allocation
|
||||||
|
let total_budget = self.config.summary_max_chars;
|
||||||
|
let older_budget = (total_budget as f64 * OLDER_BUDGET_RATIO) as usize;
|
||||||
|
let newer_budget = (total_budget as f64 * NEWER_BUDGET_RATIO) as usize;
|
||||||
|
|
||||||
|
// Step 6: Create provider and run both summaries (can be parallel)
|
||||||
|
let runtime_config = AgentRuntimeConfig::from(provider_config.clone());
|
||||||
|
let provider = create_provider(runtime_config.provider)
|
||||||
|
.map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
|
||||||
|
|
||||||
|
let (older_result, newer_result) = if older_units.is_empty() && newer_units.is_empty() {
|
||||||
|
(Ok(String::new()), Ok(String::new()))
|
||||||
|
} else if older_units.is_empty() {
|
||||||
|
let result = self
|
||||||
|
.summarize_units_segment(
|
||||||
|
provider.as_ref(),
|
||||||
|
&newer_messages,
|
||||||
|
&newer_transcript,
|
||||||
|
newer_budget,
|
||||||
|
Self::build_light_summary_prompt,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
(Ok(String::new()), result)
|
||||||
|
} else if newer_units.is_empty() {
|
||||||
|
let result = self
|
||||||
|
.summarize_units_segment(
|
||||||
|
provider.as_ref(),
|
||||||
|
&all_older_messages,
|
||||||
|
&older_transcript,
|
||||||
|
older_budget,
|
||||||
|
Self::build_heavy_summary_prompt,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
(result, Ok(String::new()))
|
||||||
|
} else {
|
||||||
|
let older_fut = self.summarize_units_segment(
|
||||||
|
provider.as_ref(),
|
||||||
|
&all_older_messages,
|
||||||
|
&older_transcript,
|
||||||
|
older_budget,
|
||||||
|
Self::build_heavy_summary_prompt,
|
||||||
|
);
|
||||||
|
let newer_fut = self.summarize_units_segment(
|
||||||
|
provider.as_ref(),
|
||||||
|
&newer_messages,
|
||||||
|
&newer_transcript,
|
||||||
|
newer_budget,
|
||||||
|
Self::build_light_summary_prompt,
|
||||||
|
);
|
||||||
|
tokio::join!(older_fut, newer_fut)
|
||||||
|
};
|
||||||
|
|
||||||
|
let older_summary = older_result?;
|
||||||
|
let newer_summary = newer_result?;
|
||||||
|
|
||||||
|
// Step 7: Assemble compressed history
|
||||||
|
let mut compressed: Vec<ChatMessage> = Vec::with_capacity(4);
|
||||||
|
|
||||||
|
// System guards first
|
||||||
|
compressed.extend(system_guards);
|
||||||
|
|
||||||
|
// Latest user message
|
||||||
|
if let Some(user_msg) = latest_user_msg {
|
||||||
|
compressed.push(user_msg);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Heavy compression summary (older)
|
||||||
|
if !older_summary.is_empty() {
|
||||||
|
compressed.push(ChatMessage::system_with_context(
|
||||||
|
format!("## 较早的操作记录(已压缩)\n\n{}", older_summary),
|
||||||
|
Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER.to_string()),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Light compression summary (newer)
|
||||||
|
if !newer_summary.is_empty() {
|
||||||
|
compressed.push(ChatMessage::system_with_context(
|
||||||
|
format!(
|
||||||
|
"## 近期操作记录(轻度压缩,保留了更多细节)\n\n{}\n\n---\n以上为近期操作记录。如需精确数据,可使用工具重新读取相关文件。",
|
||||||
|
newer_summary
|
||||||
|
),
|
||||||
|
Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER.to_string()),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
original_tokens = tokens,
|
||||||
|
final_tokens = estimate_tokens(&compressed),
|
||||||
|
final_msg_count = compressed.len(),
|
||||||
|
older_units = older_units.len(),
|
||||||
|
newer_units = newer_units.len(),
|
||||||
|
"Two-segment compression completed"
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(compressed)
|
||||||
|
}
|
||||||
|
|
||||||
|
// =========================================================================
|
||||||
|
// Legacy methods (kept for backward compatibility during migration)
|
||||||
|
// =========================================================================
|
||||||
|
|
||||||
/// Create a new compressor with the given context window size.
|
/// Create a new compressor with the given context window size.
|
||||||
pub fn new(context_window: usize) -> Self {
|
pub fn new(context_window: usize) -> Self {
|
||||||
Self {
|
Self {
|
||||||
config: ContextCompressionConfig::default(),
|
config: ContextCompressionConfig::default(),
|
||||||
context_window,
|
context_window,
|
||||||
threshold_ratio: 0.5,
|
threshold_ratio: DEFAULT_THRESHOLD_RATIO,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -323,11 +872,11 @@ Be concise, aim for {} characters or less.
|
|||||||
Self {
|
Self {
|
||||||
config,
|
config,
|
||||||
context_window,
|
context_window,
|
||||||
threshold_ratio: 0.5,
|
threshold_ratio: DEFAULT_THRESHOLD_RATIO,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get the compression threshold in tokens.
|
/// Get the compression threshold in tokens (70% of context window).
|
||||||
fn threshold(&self) -> usize {
|
fn threshold(&self) -> usize {
|
||||||
(self.context_window as f64 * self.threshold_ratio) as usize
|
(self.context_window as f64 * self.threshold_ratio) as usize
|
||||||
}
|
}
|
||||||
@ -707,7 +1256,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn test_threshold() {
|
fn test_threshold() {
|
||||||
let compressor = ContextCompressor::new(128_000);
|
let compressor = ContextCompressor::new(128_000);
|
||||||
assert_eq!(compressor.threshold(), 64_000);
|
assert_eq!(compressor.threshold(), 89_600); // 70% of 128_000
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@ -6,97 +6,60 @@ use crate::agent::AgentError;
|
|||||||
|
|
||||||
use super::session::Session;
|
use super::session::Session;
|
||||||
|
|
||||||
|
/// Run two-segment history compression synchronously.
|
||||||
|
///
|
||||||
|
/// Unlike the previous background approach (tokio::spawn), this holds the
|
||||||
|
/// session lock during the LLM calls (2–5 seconds). Since the agent loop
|
||||||
|
/// has already finished by this point there is no response-time impact, and
|
||||||
|
/// the synchronous guarantee means the next execution always starts with
|
||||||
|
/// freshly compacted history.
|
||||||
pub(crate) async fn schedule_background_history_compaction(
|
pub(crate) async fn schedule_background_history_compaction(
|
||||||
session: Arc<Mutex<Session>>,
|
session: Arc<Mutex<Session>>,
|
||||||
chat_id: impl Into<String>,
|
chat_id: impl Into<String>,
|
||||||
) -> Result<(), AgentError> {
|
) -> Result<(), AgentError> {
|
||||||
let chat_id = chat_id.into();
|
let chat_id = chat_id.into();
|
||||||
|
|
||||||
let snapshot = {
|
|
||||||
let mut session_guard = session.lock().await;
|
let mut session_guard = session.lock().await;
|
||||||
let session_record = session_guard.ensure_persistent_session(&chat_id)?;
|
session_guard.ensure_persistent_session(&chat_id)?;
|
||||||
session_guard.ensure_chat_loaded(&chat_id)?;
|
session_guard.ensure_chat_loaded(&chat_id)?;
|
||||||
|
|
||||||
let history = session_guard.get_or_create_history(&chat_id).clone();
|
let history = session_guard.get_or_create_history(&chat_id).clone();
|
||||||
let compressor = session_guard.compressor().clone();
|
let compressor = session_guard.compressor().clone();
|
||||||
|
|
||||||
if !compressor.should_compress(&history) {
|
if !compressor.should_compress(&history) {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
if !session_guard.try_start_background_compaction(&chat_id) {
|
let store = session_guard.store();
|
||||||
return Ok(());
|
let session_id = session_guard.persistent_session_id(&chat_id);
|
||||||
}
|
let provider_config = session_guard.provider_config().clone();
|
||||||
|
|
||||||
(
|
|
||||||
session_guard.store(),
|
|
||||||
session_guard.persistent_session_id(&chat_id),
|
|
||||||
session_record.message_count,
|
|
||||||
history,
|
|
||||||
compressor,
|
|
||||||
session_guard.provider_config().clone(),
|
|
||||||
)
|
|
||||||
};
|
|
||||||
|
|
||||||
let (
|
|
||||||
store,
|
|
||||||
session_id,
|
|
||||||
snapshot_end_seq,
|
|
||||||
history,
|
|
||||||
compressor,
|
|
||||||
provider_config,
|
|
||||||
) = snapshot;
|
|
||||||
let session_for_task = session.clone();
|
|
||||||
let chat_id_for_task = chat_id.clone();
|
|
||||||
|
|
||||||
tokio::spawn(async move {
|
|
||||||
tracing::info!(chat_id = %chat_id_for_task, snapshot_end_seq, "Starting background history compaction");
|
|
||||||
|
|
||||||
let compaction_result = compressor
|
|
||||||
.build_compaction_plan(&history, &provider_config)
|
|
||||||
.await;
|
|
||||||
let mut committed = false;
|
|
||||||
|
|
||||||
match compaction_result {
|
|
||||||
Ok(Some(plan)) => match store.compact_active_history(
|
|
||||||
&session_id,
|
|
||||||
snapshot_end_seq,
|
|
||||||
&plan.preserved_system_messages,
|
|
||||||
&plan.summary_message,
|
|
||||||
&plan.preserved_messages,
|
|
||||||
) {
|
|
||||||
Ok(true) => {
|
|
||||||
committed = true;
|
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
chat_id = %chat_id_for_task,
|
chat_id = %chat_id,
|
||||||
snapshot_end_seq,
|
msg_count = history.len(),
|
||||||
compressed_turns = plan.compressed_turns,
|
"Starting synchronous two-segment compression"
|
||||||
preserved_turns = plan.preserved_turns,
|
|
||||||
"Background history compaction committed"
|
|
||||||
);
|
);
|
||||||
}
|
|
||||||
Ok(false) => {
|
|
||||||
tracing::info!(chat_id = %chat_id_for_task, snapshot_end_seq, "Background history compaction skipped due to stale snapshot");
|
|
||||||
}
|
|
||||||
Err(error) => {
|
|
||||||
tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Background history compaction commit failed");
|
|
||||||
}
|
|
||||||
},
|
|
||||||
Ok(None) => {
|
|
||||||
tracing::debug!(chat_id = %chat_id_for_task, "Background history compaction not needed after snapshot analysis");
|
|
||||||
}
|
|
||||||
Err(error) => {
|
|
||||||
tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Background history compaction build failed");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut session_guard = session_for_task.lock().await;
|
// Synchronous compression — holds lock during LLM calls.
|
||||||
if committed {
|
// compress_two_segment guarantees the result contains no tool_calls,
|
||||||
if let Err(error) = session_guard.reload_chat_history(&chat_id_for_task) {
|
// so there is no risk of orphaned tool call sequences.
|
||||||
tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Failed to reload history after background compaction");
|
let compressed = compressor
|
||||||
}
|
.compress_two_segment(&history, &provider_config)
|
||||||
}
|
.await?;
|
||||||
session_guard.finish_background_compaction(&chat_id_for_task);
|
|
||||||
});
|
// Replace the entire history with the compressed result.
|
||||||
|
// Since we hold the lock, no concurrent modifications can occur.
|
||||||
|
store
|
||||||
|
.replace_active_history(&session_id, &compressed)
|
||||||
|
.map_err(|e| AgentError::Other(format!("replace_active_history error: {}", e)))?;
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
chat_id = %chat_id,
|
||||||
|
compressed_msg_count = compressed.len(),
|
||||||
|
"Two-segment compression committed"
|
||||||
|
);
|
||||||
|
|
||||||
|
session_guard.reload_chat_history(&chat_id)?;
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@ -526,10 +526,12 @@ impl Session {
|
|||||||
&self.compressor
|
&self.compressor
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
pub(crate) fn try_start_background_compaction(&mut self, chat_id: &str) -> bool {
|
pub(crate) fn try_start_background_compaction(&mut self, chat_id: &str) -> bool {
|
||||||
self.history.try_start_background_compaction(chat_id)
|
self.history.try_start_background_compaction(chat_id)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
pub(crate) fn finish_background_compaction(&mut self, chat_id: &str) {
|
pub(crate) fn finish_background_compaction(&mut self, chat_id: &str) {
|
||||||
self.history.finish_background_compaction(chat_id);
|
self.history.finish_background_compaction(chat_id);
|
||||||
}
|
}
|
||||||
|
|||||||
@ -235,10 +235,12 @@ impl SessionHistory {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
pub(crate) fn try_start_background_compaction(&mut self, chat_id: &str) -> bool {
|
pub(crate) fn try_start_background_compaction(&mut self, chat_id: &str) -> bool {
|
||||||
self.compression_in_flight.insert(chat_id.to_string())
|
self.compression_in_flight.insert(chat_id.to_string())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
pub(crate) fn finish_background_compaction(&mut self, chat_id: &str) {
|
pub(crate) fn finish_background_compaction(&mut self, chat_id: &str) {
|
||||||
self.compression_in_flight.remove(chat_id);
|
self.compression_in_flight.remove(chat_id);
|
||||||
}
|
}
|
||||||
|
|||||||
@ -814,6 +814,59 @@ impl SessionStore {
|
|||||||
Ok(true)
|
Ok(true)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Replace the entire active history for a session.
|
||||||
|
///
|
||||||
|
/// This is a simpler alternative to `compact_active_history` for when the
|
||||||
|
/// compressor has already produced a complete, validated message list
|
||||||
|
/// (e.g. two-segment compression). It replaces all existing messages
|
||||||
|
/// with the new list in a single transaction.
|
||||||
|
pub fn replace_active_history(
|
||||||
|
&self,
|
||||||
|
session_id: &str,
|
||||||
|
messages: &[ChatMessage],
|
||||||
|
) -> Result<(), StorageError> {
|
||||||
|
let conn = self.pool.get()?;
|
||||||
|
let tx = conn.unchecked_transaction()?;
|
||||||
|
let now = current_timestamp();
|
||||||
|
|
||||||
|
// Delete all existing messages for this session
|
||||||
|
tx.execute(
|
||||||
|
"DELETE FROM messages WHERE session_id = ?1",
|
||||||
|
params![session_id],
|
||||||
|
)?;
|
||||||
|
|
||||||
|
// Insert new messages with sequential seq numbers
|
||||||
|
let mut active_user_turn_count = 0_i64;
|
||||||
|
for (i, message) in messages.iter().enumerate() {
|
||||||
|
let seq = (i + 1) as i64;
|
||||||
|
if message.role == "user" {
|
||||||
|
active_user_turn_count += 1;
|
||||||
|
}
|
||||||
|
insert_message_with_seq(&tx, session_id, seq, message)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
tx.execute(
|
||||||
|
"
|
||||||
|
UPDATE sessions
|
||||||
|
SET message_count = ?2,
|
||||||
|
user_turn_count = ?3,
|
||||||
|
updated_at = ?4,
|
||||||
|
last_active_at = ?4,
|
||||||
|
archived_at = NULL
|
||||||
|
WHERE id = ?1 AND deleted_at IS NULL
|
||||||
|
",
|
||||||
|
params![
|
||||||
|
session_id,
|
||||||
|
messages.len() as i64,
|
||||||
|
active_user_turn_count,
|
||||||
|
now,
|
||||||
|
],
|
||||||
|
)?;
|
||||||
|
|
||||||
|
tx.commit()?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
pub fn mark_agent_prompt_reinjected(&self, session_id: &str) -> Result<(), StorageError> {
|
pub fn mark_agent_prompt_reinjected(&self, session_id: &str) -> Result<(), StorageError> {
|
||||||
let now = current_timestamp();
|
let now = current_timestamp();
|
||||||
let conn = self.pool.get()?;
|
let conn = self.pool.get()?;
|
||||||
|
|||||||
@ -54,6 +54,15 @@ pub trait ConversationRepository: Send + Sync + 'static {
|
|||||||
summary_message: &ChatMessage,
|
summary_message: &ChatMessage,
|
||||||
preserved_messages: &[ChatMessage],
|
preserved_messages: &[ChatMessage],
|
||||||
) -> Result<bool, StorageError>;
|
) -> Result<bool, StorageError>;
|
||||||
|
|
||||||
|
/// Replace the entire active history for a session with new messages.
|
||||||
|
/// Used when the compressor has already produced a complete, validated
|
||||||
|
/// message list (e.g. two-segment compression).
|
||||||
|
fn replace_active_history(
|
||||||
|
&self,
|
||||||
|
session_id: &str,
|
||||||
|
messages: &[ChatMessage],
|
||||||
|
) -> Result<(), StorageError>;
|
||||||
}
|
}
|
||||||
|
|
||||||
pub trait PromptInjectionRepository: Send + Sync + 'static {
|
pub trait PromptInjectionRepository: Send + Sync + 'static {
|
||||||
@ -235,6 +244,14 @@ impl ConversationRepository for super::SessionStore {
|
|||||||
preserved_messages,
|
preserved_messages,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn replace_active_history(
|
||||||
|
&self,
|
||||||
|
session_id: &str,
|
||||||
|
messages: &[ChatMessage],
|
||||||
|
) -> Result<(), StorageError> {
|
||||||
|
super::SessionStore::replace_active_history(self, session_id, messages)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl PromptInjectionRepository for super::SessionStore {
|
impl PromptInjectionRepository for super::SessionStore {
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user