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
|
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(
|
fn image_token_budget_for_request(
|
||||||
runtime_config: &AgentRuntimeConfig,
|
runtime_config: &AgentRuntimeConfig,
|
||||||
text_only_messages: &[Message],
|
text_tokens: usize,
|
||||||
tools: Option<&Vec<crate::domain::tools::Tool>>,
|
tools_tokens: usize,
|
||||||
) -> usize {
|
) -> usize {
|
||||||
let completion_reserve = runtime_config
|
let completion_reserve = runtime_config
|
||||||
.provider
|
.provider
|
||||||
@ -170,12 +198,7 @@ fn image_token_budget_for_request(
|
|||||||
.saturating_sub(completion_reserve);
|
.saturating_sub(completion_reserve);
|
||||||
let safe_input_window = (input_window as f64 * CONTEXT_INPUT_SAFETY_RATIO) as usize;
|
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)
|
safe_input_window.saturating_sub(text_tokens + tools_tokens)
|
||||||
+ tools
|
|
||||||
.map(estimate_tokens_from_serialized_json)
|
|
||||||
.unwrap_or_default();
|
|
||||||
|
|
||||||
safe_input_window.saturating_sub(text_tokens)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn count_supported_image_media_refs(messages: &[ChatMessage]) -> usize {
|
fn count_supported_image_media_refs(messages: &[ChatMessage]) -> usize {
|
||||||
@ -700,23 +723,46 @@ impl<H: EmittedMessageHandler> PersistingEmittedMessageHandler<H> {
|
|||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl<H: EmittedMessageHandler> EmittedMessageHandler for PersistingEmittedMessageHandler<H> {
|
impl<H: EmittedMessageHandler> EmittedMessageHandler for PersistingEmittedMessageHandler<H> {
|
||||||
async fn handle(&self, message: ChatMessage) {
|
async fn handle(&self, message: ChatMessage) {
|
||||||
if let Err(e) = self.conversation_repository
|
// SQLite 是同步 I/O,直接调用会阻塞 tokio worker 线程。
|
||||||
.append_message_with_topic(&self.session_id, self.topic_id.as_deref(), &message)
|
// 用 spawn_blocking 把写操作挪到 blocking 线程池,不改变语义:
|
||||||
{
|
// 失败仅记日志不阻断主流程(与原 if let Err 一致)。
|
||||||
tracing::error!(error = %e, session_id = %self.session_id,
|
let repo = self.conversation_repository.clone();
|
||||||
"Failed to persist emitted message");
|
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;
|
self.inner.handle(message).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn handle_tool_result(&self, message: ChatMessage, duration_ms: Option<u64>) {
|
async fn handle_tool_result(&self, message: ChatMessage, duration_ms: Option<u64>) {
|
||||||
// Persist the ChatMessage first (no duration field, same as before)
|
// Persist the ChatMessage first (no duration field, same as before)
|
||||||
if let Err(e) = self.conversation_repository
|
let repo = self.conversation_repository.clone();
|
||||||
.append_message_with_topic(&self.session_id, self.topic_id.as_deref(), &message)
|
let session_id = self.session_id.clone();
|
||||||
{
|
let topic_id = self.topic_id.clone();
|
||||||
tracing::error!(error = %e, session_id = %self.session_id,
|
let msg_for_persist = message.clone();
|
||||||
"Failed to persist emitted message");
|
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;
|
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 loop_detector = LoopDetector::new(LoopDetectorConfig::default());
|
||||||
let mut emitted_messages = Vec::new();
|
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 {
|
for iteration in 0..self.max_iterations {
|
||||||
#[cfg(debug_assertions)]
|
#[cfg(debug_assertions)]
|
||||||
tracing::debug!(iteration, "Agent iteration started");
|
tracing::debug!(iteration, "Agent iteration started");
|
||||||
@ -934,14 +997,6 @@ impl AgentLoop {
|
|||||||
return Ok(cancel);
|
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
|
// Defense-in-depth: sanitize incomplete tool call sequences
|
||||||
// before EVERY LLM request, not just once at process() entry.
|
// before EVERY LLM request, not just once at process() entry.
|
||||||
// This catches edge cases where compression, persistence races,
|
// 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
|
// Set up streaming delta consumer
|
||||||
// Pre-generate the message ID so stream deltas and the final assistant
|
// Pre-generate the message ID so stream deltas and the final assistant
|
||||||
@ -1182,11 +1237,16 @@ impl AgentLoop {
|
|||||||
|
|
||||||
/// 构建 LLM 请求:过滤图片、构建系统提示、转换消息格式。
|
/// 构建 LLM 请求:过滤图片、构建系统提示、转换消息格式。
|
||||||
/// 统一了循环内调用和最终 summary 调用的请求构建逻辑。
|
/// 统一了循环内调用和最终 summary 调用的请求构建逻辑。
|
||||||
|
///
|
||||||
|
/// 优化:token 估算直接基于 ChatMessage 累加字段长度(轻量),
|
||||||
|
/// 不再构造中间 text_only_messages: Vec<Message>(避免 N 条消息 × 5 字段双克隆)。
|
||||||
|
/// 工具 token 估算在循环外预算后传入(tools_tokens)。
|
||||||
fn build_llm_request(
|
fn build_llm_request(
|
||||||
&self,
|
&self,
|
||||||
messages: &[ChatMessage],
|
messages: &[ChatMessage],
|
||||||
system_prompt_context: Option<&SystemPromptContext>,
|
system_prompt_context: Option<&SystemPromptContext>,
|
||||||
tools: Option<Vec<crate::domain::tools::Tool>>,
|
tools: Option<Vec<crate::domain::tools::Tool>>,
|
||||||
|
tools_tokens: usize,
|
||||||
) -> ChatCompletionRequest {
|
) -> ChatCompletionRequest {
|
||||||
let filtered_messages = filter_images_by_age_and_count(
|
let filtered_messages = filter_images_by_age_and_count(
|
||||||
messages,
|
messages,
|
||||||
@ -1201,18 +1261,18 @@ impl AgentLoop {
|
|||||||
.and_then(|provider| provider.build(ctx))
|
.and_then(|provider| provider.build(ctx))
|
||||||
});
|
});
|
||||||
|
|
||||||
let mut text_only_messages: Vec<Message> = Vec::with_capacity(filtered_messages.len() + 2);
|
// 轻量 token 估算:直接从 ChatMessage 累加,无需构造中间 Vec<Message>
|
||||||
if let Some(ref prompt) = system_prompt {
|
let text_tokens = estimate_text_tokens_from_chat_messages(
|
||||||
text_only_messages.push(Message::system(prompt.content.clone()));
|
&filtered_messages,
|
||||||
}
|
system_prompt.as_ref().map(|p| p.content.as_str()),
|
||||||
text_only_messages.extend(filtered_messages.iter().map(chat_message_to_text_only_llm_message));
|
);
|
||||||
|
|
||||||
let image_tokens = image_token_budget_for_request(
|
let image_tokens = image_token_budget_for_request(
|
||||||
&self.runtime_config,
|
&self.runtime_config,
|
||||||
&text_only_messages,
|
text_tokens,
|
||||||
tools.as_ref(),
|
tools_tokens,
|
||||||
);
|
);
|
||||||
let mut image_budget = ImageInlineBudget::new(image_tokens, image_count);
|
let mut image_budget = ImageInlineBudget::new(image_tokens, image_count);
|
||||||
|
|
||||||
let mut messages_for_llm: Vec<Message> = Vec::with_capacity(filtered_messages.len() + 2);
|
let mut messages_for_llm: Vec<Message> = Vec::with_capacity(filtered_messages.len() + 2);
|
||||||
if let Some(ref prompt) = system_prompt {
|
if let Some(ref prompt) = system_prompt {
|
||||||
messages_for_llm.push(Message::system(prompt.content.clone()));
|
messages_for_llm.push(Message::system(prompt.content.clone()));
|
||||||
@ -1321,7 +1381,7 @@ impl AgentLoop {
|
|||||||
);
|
);
|
||||||
messages.push(summary_request);
|
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 调用也与取消信号竞速
|
// 最终 summary 调用也与取消信号竞速
|
||||||
let final_result: Result<
|
let final_result: Result<
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user