perf(agent-loop): 缓存工具定义与轻量 token 估算,DB 写入 offload 至 blocking 池
- 工具定义在 process() 循环外构建一次(快照语义:单次 process 内 Arc<ToolRegistry> 不变) - 工具 token 估算循环外预算一次,避免每轮 serde_json::to_string 全量序列化 - 新增 estimate_text_tokens_from_chat_messages 直接基于 ChatMessage 累加字段长度, 消除中间 text_only_messages: Vec<Message> 构造(N 条消息 × 5 字段双克隆 → 0) - PersistingEmittedMessageHandler 用 spawn_blocking 包装 SQLite 同步写, 解除 tokio worker 线程阻塞,不改变失败语义(仅记日志不阻断主流程)
This commit is contained in:
parent
649ea2998d
commit
92db80dc3f
@ -155,10 +155,38 @@ fn estimate_tokens_from_serialized_json<T: serde::Serialize>(value: &T) -> usize
|
||||
as usize
|
||||
}
|
||||
|
||||
/// 轻量 token 估算:直接从 ChatMessage 累加文本字段长度,
|
||||
/// 避免对完整 Vec<Message> 做 serde_json::to_string(O(N) 序列化)。
|
||||
/// 精度略低于完整 JSON 序列化(忽略 role/tool_call_id 等 overhead),
|
||||
/// 但作为图片预算估算已足够——上下文窗口有 CONTEXT_INPUT_SAFETY_RATIO 兜底。
|
||||
/// 与原 text_only_messages 路径等价覆盖:content + reasoning_content + tool_calls。
|
||||
fn estimate_text_tokens_from_chat_messages(
|
||||
messages: &[ChatMessage],
|
||||
system_prompt: Option<&str>,
|
||||
) -> usize {
|
||||
let mut total_chars: usize = system_prompt.map(|s| s.len()).unwrap_or(0);
|
||||
for m in messages {
|
||||
total_chars += m.content.len();
|
||||
if let Some(rc) = m.reasoning_content.as_ref() {
|
||||
total_chars += rc.len();
|
||||
}
|
||||
if let Some(tcs) = m.tool_calls.as_ref() {
|
||||
for tc in tcs {
|
||||
total_chars += tc.name.len();
|
||||
total_chars += serde_json::to_string(&tc.arguments)
|
||||
.map(|s| s.len())
|
||||
.unwrap_or(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
((total_chars.div_ceil(TOKEN_ESTIMATE_CHARS_PER_TOKEN) as f64)
|
||||
* TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize
|
||||
}
|
||||
|
||||
fn image_token_budget_for_request(
|
||||
runtime_config: &AgentRuntimeConfig,
|
||||
text_only_messages: &[Message],
|
||||
tools: Option<&Vec<crate::domain::tools::Tool>>,
|
||||
text_tokens: usize,
|
||||
tools_tokens: usize,
|
||||
) -> usize {
|
||||
let completion_reserve = runtime_config
|
||||
.provider
|
||||
@ -170,12 +198,7 @@ fn image_token_budget_for_request(
|
||||
.saturating_sub(completion_reserve);
|
||||
let safe_input_window = (input_window as f64 * CONTEXT_INPUT_SAFETY_RATIO) as usize;
|
||||
|
||||
let text_tokens = estimate_tokens_from_serialized_json(&text_only_messages)
|
||||
+ tools
|
||||
.map(estimate_tokens_from_serialized_json)
|
||||
.unwrap_or_default();
|
||||
|
||||
safe_input_window.saturating_sub(text_tokens)
|
||||
safe_input_window.saturating_sub(text_tokens + tools_tokens)
|
||||
}
|
||||
|
||||
fn count_supported_image_media_refs(messages: &[ChatMessage]) -> usize {
|
||||
@ -700,23 +723,46 @@ impl<H: EmittedMessageHandler> PersistingEmittedMessageHandler<H> {
|
||||
#[async_trait]
|
||||
impl<H: EmittedMessageHandler> EmittedMessageHandler for PersistingEmittedMessageHandler<H> {
|
||||
async fn handle(&self, message: ChatMessage) {
|
||||
if let Err(e) = self.conversation_repository
|
||||
.append_message_with_topic(&self.session_id, self.topic_id.as_deref(), &message)
|
||||
{
|
||||
tracing::error!(error = %e, session_id = %self.session_id,
|
||||
"Failed to persist emitted message");
|
||||
}
|
||||
// SQLite 是同步 I/O,直接调用会阻塞 tokio worker 线程。
|
||||
// 用 spawn_blocking 把写操作挪到 blocking 线程池,不改变语义:
|
||||
// 失败仅记日志不阻断主流程(与原 if let Err 一致)。
|
||||
let repo = self.conversation_repository.clone();
|
||||
let session_id = self.session_id.clone();
|
||||
let topic_id = self.topic_id.clone();
|
||||
let msg_for_persist = message.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
if let Err(e) = repo.append_message_with_topic(
|
||||
&session_id,
|
||||
topic_id.as_deref(),
|
||||
&msg_for_persist,
|
||||
) {
|
||||
tracing::error!(error = %e, session_id = %session_id,
|
||||
"Failed to persist emitted message");
|
||||
}
|
||||
})
|
||||
.await
|
||||
.ok(); // JoinError 不影响主流程
|
||||
self.inner.handle(message).await;
|
||||
}
|
||||
|
||||
async fn handle_tool_result(&self, message: ChatMessage, duration_ms: Option<u64>) {
|
||||
// Persist the ChatMessage first (no duration field, same as before)
|
||||
if let Err(e) = self.conversation_repository
|
||||
.append_message_with_topic(&self.session_id, self.topic_id.as_deref(), &message)
|
||||
{
|
||||
tracing::error!(error = %e, session_id = %self.session_id,
|
||||
"Failed to persist emitted message");
|
||||
}
|
||||
let repo = self.conversation_repository.clone();
|
||||
let session_id = self.session_id.clone();
|
||||
let topic_id = self.topic_id.clone();
|
||||
let msg_for_persist = message.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
if let Err(e) = repo.append_message_with_topic(
|
||||
&session_id,
|
||||
topic_id.as_deref(),
|
||||
&msg_for_persist,
|
||||
) {
|
||||
tracing::error!(error = %e, session_id = %session_id,
|
||||
"Failed to persist emitted message");
|
||||
}
|
||||
})
|
||||
.await
|
||||
.ok();
|
||||
self.inner.handle_tool_result(message, duration_ms).await;
|
||||
}
|
||||
|
||||
@ -920,6 +966,23 @@ impl AgentLoop {
|
||||
let mut loop_detector = LoopDetector::new(LoopDetectorConfig::default());
|
||||
let mut emitted_messages = Vec::new();
|
||||
|
||||
// 工具集在 create_agent 时按当前专家 capability 过滤后注入(快照语义)。
|
||||
// 单次 process() 期间专家切换不影响当前 Arc<ToolRegistry> 引用,
|
||||
// 因此循环外构建一次 tool_defs,避免每轮重复读锁 + 全量 clone。
|
||||
// 注意:若未来支持 process() 中途热切换专家(当前不支持),需重新评估。
|
||||
let tool_defs = self.tools.get_definitions();
|
||||
let tools = if tool_defs.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(tool_defs)
|
||||
};
|
||||
// 工具 token 估算在循环外算一次(tool_defs 在 process() 期间不变),
|
||||
// 避免每轮 serde_json::to_string 全量序列化工具定义。
|
||||
let tools_tokens = tools
|
||||
.as_ref()
|
||||
.map(|t| estimate_tokens_from_serialized_json(t))
|
||||
.unwrap_or_default();
|
||||
|
||||
for iteration in 0..self.max_iterations {
|
||||
#[cfg(debug_assertions)]
|
||||
tracing::debug!(iteration, "Agent iteration started");
|
||||
@ -934,14 +997,6 @@ impl AgentLoop {
|
||||
return Ok(cancel);
|
||||
}
|
||||
|
||||
// Build request
|
||||
let tool_defs = self.tools.get_definitions();
|
||||
let tools = if tool_defs.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(tool_defs)
|
||||
};
|
||||
|
||||
// Defense-in-depth: sanitize incomplete tool call sequences
|
||||
// before EVERY LLM request, not just once at process() entry.
|
||||
// This catches edge cases where compression, persistence races,
|
||||
@ -956,7 +1011,7 @@ impl AgentLoop {
|
||||
);
|
||||
}
|
||||
|
||||
let request = self.build_llm_request(&messages, system_prompt_context, tools);
|
||||
let request = self.build_llm_request(&messages, system_prompt_context, tools.clone(), tools_tokens);
|
||||
|
||||
// Set up streaming delta consumer
|
||||
// Pre-generate the message ID so stream deltas and the final assistant
|
||||
@ -1182,11 +1237,16 @@ impl AgentLoop {
|
||||
|
||||
/// 构建 LLM 请求:过滤图片、构建系统提示、转换消息格式。
|
||||
/// 统一了循环内调用和最终 summary 调用的请求构建逻辑。
|
||||
///
|
||||
/// 优化:token 估算直接基于 ChatMessage 累加字段长度(轻量),
|
||||
/// 不再构造中间 text_only_messages: Vec<Message>(避免 N 条消息 × 5 字段双克隆)。
|
||||
/// 工具 token 估算在循环外预算后传入(tools_tokens)。
|
||||
fn build_llm_request(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
system_prompt_context: Option<&SystemPromptContext>,
|
||||
tools: Option<Vec<crate::domain::tools::Tool>>,
|
||||
tools_tokens: usize,
|
||||
) -> ChatCompletionRequest {
|
||||
let filtered_messages = filter_images_by_age_and_count(
|
||||
messages,
|
||||
@ -1201,18 +1261,18 @@ impl AgentLoop {
|
||||
.and_then(|provider| provider.build(ctx))
|
||||
});
|
||||
|
||||
let mut text_only_messages: Vec<Message> = Vec::with_capacity(filtered_messages.len() + 2);
|
||||
if let Some(ref prompt) = system_prompt {
|
||||
text_only_messages.push(Message::system(prompt.content.clone()));
|
||||
}
|
||||
text_only_messages.extend(filtered_messages.iter().map(chat_message_to_text_only_llm_message));
|
||||
|
||||
// 轻量 token 估算:直接从 ChatMessage 累加,无需构造中间 Vec<Message>
|
||||
let text_tokens = estimate_text_tokens_from_chat_messages(
|
||||
&filtered_messages,
|
||||
system_prompt.as_ref().map(|p| p.content.as_str()),
|
||||
);
|
||||
let image_tokens = image_token_budget_for_request(
|
||||
&self.runtime_config,
|
||||
&text_only_messages,
|
||||
tools.as_ref(),
|
||||
text_tokens,
|
||||
tools_tokens,
|
||||
);
|
||||
let mut image_budget = ImageInlineBudget::new(image_tokens, image_count);
|
||||
|
||||
let mut messages_for_llm: Vec<Message> = Vec::with_capacity(filtered_messages.len() + 2);
|
||||
if let Some(ref prompt) = system_prompt {
|
||||
messages_for_llm.push(Message::system(prompt.content.clone()));
|
||||
@ -1321,7 +1381,7 @@ impl AgentLoop {
|
||||
);
|
||||
messages.push(summary_request);
|
||||
|
||||
let request = self.build_llm_request(messages, system_prompt_context, None);
|
||||
let request = self.build_llm_request(messages, system_prompt_context, None, 0);
|
||||
|
||||
// 最终 summary 调用也与取消信号竞速
|
||||
let final_result: Result<
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user