feat(llm): LlmCycle 切换 IR 消息类型,移除 Phase 0 桥接层

核心改动:
- LlmCycle.messages 从 Vec<OpenaiChatMessage> 切换为 Vec<Message>,
  移除 Phase 0 引入的 chat_message_to_message / message_to_chat_message
  转换函数
- 移除 LlmCycle.system_prompt 字段(FIX-D),调用方通过
  Message::system_text() + with_messages() 管理系统提示;
  with_system_prompt() 标记 #[deprecated] 保留为过渡期 shim
- submit_stream 简化(FIX-E):Item = StreamEvent 不再带 Result 包装,
  错误用 StreamEvent::Error 传出;self.messages 不再自动 push_message,
  调用方在收到 MessageComplete 后手动调用 cycle.push_message()
- compact.rs 适配 Message 类型 + ContentBlock 估算;microcompact
  跳过 is_error:true 的 ToolResult(FIX-F),保留 LLM 错误诊断
- ConversationMemory 持久化从 OpenaiChatMessage 切换为 Message
- ToolInvocation 增加 tool_call_id 字段(FIX-A),
  invoke/invoke_all 签名增加 tool_call_id 参数

调用方迁移:
- AgentSession::submit_turn 用 Message::system_text() + with_messages()
- simple_visit example 同步迁移

测试 177 passed / 0 failed(compact 新增 5 个单元测试覆盖 FIX-F + 估计全变体)。
This commit is contained in:
徐涛
2026-07-02 22:48:00 +08:00
parent a74c24b6fe
commit 9e4f50c955
6 changed files with 364 additions and 310 deletions
+4 -5
View File
@@ -4,9 +4,7 @@ use agcore::init_tracing;
use agcore::llm::{
cycle::{CycleConfig, LlmCycle},
provider::{create_provider, ProviderConfig, ProviderType},
types::{
message::ContentBlock, response_v2::MessageResponse,
},
types::{message::ContentBlock, message::Message, response_v2::MessageResponse},
};
fn extract_response_text(response: &MessageResponse) -> &str {
@@ -65,8 +63,9 @@ async fn main() {
..CycleConfig::default()
};
let mut cycle = LlmCycle::new(provider, cycle_config)
.with_system_prompt("你是一个简洁的助手,对于任何问题都是用一句话回答。".to_string());
let mut cycle = LlmCycle::new(provider, cycle_config).with_messages(vec![
Message::system("你是一个简洁的助手,对于任何问题都是用一句话回答。"),
]);
println!("发送请求...");
+7 -1
View File
@@ -140,8 +140,14 @@ impl AgentSession {
let _ = self.agent.tool_definitions(&self.bundle);
let mut cycle = LlmCycle::new_with_arc(Arc::clone(&self.bundle.provider), CycleConfig::default())
.with_messages(Vec::new());
// Phase 2 切换 system_prompt 字段为 Message::SystemFIX-D)。
// 若 agent 自带 system prompt,预置到 messages 列表头部。
let mut initial_messages: Vec<Message> = Vec::new();
if let Some(prompt) = self.agent.system_prompt() {
cycle = cycle.with_system_prompt(prompt.to_string());
initial_messages.push(Message::system(prompt));
}
if !initial_messages.is_empty() {
cycle = cycle.with_messages(initial_messages);
}
if let Some(cfg) = self.bundle.config.compact_config.clone() {
cycle = cycle.with_compact_config(cfg);
+133 -28
View File
@@ -1,6 +1,6 @@
//! 上下文自动压缩 —— 当对话历史过长时自动压缩。
use crate::llm::types::{ContentField, OpenaiChatMessage, OpenaiContentPart};
use crate::llm::types::message::{ContentBlock, Message};
const AUTOCOMPACT_BUFFER_TOKENS: u32 = 13_000;
const RESERVED_OUTPUT_TOKENS: u32 = 20_000;
@@ -72,37 +72,43 @@ impl CompactState {
}
/// 粗略估计消息列表的 token 数(基于字符数,4 字符 ≈ 1 token)。
pub fn estimate_message_tokens(messages: &[OpenaiChatMessage]) -> u32 {
pub fn estimate_message_tokens(messages: &[Message]) -> u32 {
messages
.iter()
.map(estimate_single_message_tokens)
.sum()
}
fn estimate_single_message_tokens(msg: &OpenaiChatMessage) -> u32 {
fn estimate_single_message_tokens(msg: &Message) -> u32 {
let role_overhead: u32 = 4;
let content_tokens = match msg {
OpenaiChatMessage::Developer { content, .. }
| OpenaiChatMessage::System { content, .. }
| OpenaiChatMessage::User { content, .. }
| OpenaiChatMessage::Assistant { content, .. }
| OpenaiChatMessage::Function { content, .. } => estimate_content_tokens(content),
OpenaiChatMessage::Tool { content, .. } => estimate_content_tokens(content),
Message::System { content }
| Message::User { content }
| Message::Assistant { content }
| Message::ToolResult { content, .. } => estimate_content_blocks_tokens(content),
Message::UserImage { .. } => 50,
};
role_overhead + content_tokens
}
fn estimate_content_tokens(content: &ContentField) -> u32 {
match content {
ContentField::String(s) => estimate_text_tokens(s),
ContentField::Array(parts) => parts.iter().map(estimate_part_tokens).sum(),
}
fn estimate_content_blocks_tokens(blocks: &[ContentBlock]) -> u32 {
blocks.iter().map(estimate_block_tokens).sum()
}
fn estimate_part_tokens(part: &OpenaiContentPart) -> u32 {
match part {
OpenaiContentPart::Text { text } => estimate_text_tokens(text),
_ => 50,
fn estimate_block_tokens(block: &ContentBlock) -> u32 {
match block {
ContentBlock::Text { text } => estimate_text_tokens(text),
ContentBlock::Thinking { text, .. } => estimate_text_tokens(text),
ContentBlock::ToolUse { input, .. } => {
estimate_text_tokens(&input.to_string())
}
ContentBlock::ToolResult { content, .. } => estimate_content_blocks_tokens(content),
// ponytail: Image / Audio / File / Extension 在 IR 中固定估算。
// 无文本的视觉/音频 block 用兜底估算,避免 token 计数膨胀。
ContentBlock::Image { .. }
| ContentBlock::Audio { .. }
| ContentBlock::File { .. }
| ContentBlock::Extension { .. } => 50,
}
}
@@ -115,11 +121,7 @@ fn estimate_text_tokens(text: &str) -> u32 {
}
/// 判断是否需要触发自动压缩。
pub fn should_compact(
messages: &[OpenaiChatMessage],
config: &CompactConfig,
state: &CompactState,
) -> bool {
pub fn should_compact(messages: &[Message], config: &CompactConfig, state: &CompactState) -> bool {
if state.consecutive_failures >= MAX_CONSECUTIVE_FAILURES {
return false;
}
@@ -133,7 +135,10 @@ pub fn should_compact(
/// 保留最近的 `keep_recent` 条消息不变。
///
/// 返回释放的估算 token 数。
pub fn microcompact(messages: &mut [OpenaiChatMessage], keep_recent: usize) -> u32 {
///
/// **审查 FIX-F**:仅压缩 `is_error: false` 的 `ToolResult` —— 错误结果包含对 LLM
/// 理解失败原因至关重要的诊断信息,压缩后 LLM 无法理解。
pub fn microcompact(messages: &mut [Message], keep_recent: usize) -> u32 {
if messages.len() <= keep_recent {
return 0;
}
@@ -141,19 +146,119 @@ pub fn microcompact(messages: &mut [OpenaiChatMessage], keep_recent: usize) -> u
let prune_start = messages.len() - keep_recent;
let mut freed_tokens: u32 = 0;
// 第一遍:计算可释放 token(仅非错误 ToolResult
for msg in &messages[..prune_start] {
if matches!(msg, OpenaiChatMessage::Tool { .. }) {
if matches!(msg, Message::ToolResult { is_error: false, .. }) {
freed_tokens += estimate_single_message_tokens(msg);
}
}
// 第二遍:替换内容(仅非错误 ToolResult
for msg in &mut messages[..prune_start] {
if let OpenaiChatMessage::Tool { content, .. } = msg {
*content = ContentField::Array(vec![OpenaiContentPart::Text {
if let Message::ToolResult { content, is_error: false, .. } = msg {
*content = vec![ContentBlock::Text {
text: "[pruned]".to_string(),
}]);
}];
}
}
freed_tokens
}
#[cfg(test)]
mod tests {
use super::*;
fn user_msg(s: &str) -> Message {
Message::user_text(s)
}
#[test]
fn estimate_message_tokens_handles_all_variants() {
let messages = vec![
Message::System {
content: vec![ContentBlock::Text {
text: "sys".into(),
}],
},
Message::user_text("hi"),
Message::assistant("ans"),
Message::user_image("b64", "image/png", crate::llm::types::shared::ImageDetail::Auto),
Message::tool_result("call_1", "tool res", false),
];
let tokens = estimate_message_tokens(&messages);
// 至少 5 条消息 × 4 role overhead = 20 + 文本/估算
assert!(tokens > 20);
}
#[test]
fn microcompact_replaces_old_tool_result_with_pruned() {
let mut messages = vec![
user_msg("hi"),
Message::tool_result("call_1", "raw result".repeat(50), false),
user_msg("again"),
user_msg("keep recent 1"),
user_msg("keep recent 2"),
];
let before_len = messages.len();
let freed = microcompact(&mut messages, 2);
assert!(freed > 0);
assert_eq!(messages.len(), before_len); // 只改内容,不删消息
// 索引 1 是被压缩的 ToolResult
if let Message::ToolResult { content, is_error, .. } = &messages[1] {
assert_eq!(content.len(), 1);
assert!(matches!(&content[0], ContentBlock::Text { text } if text == "[pruned]"));
assert!(!is_error);
} else {
panic!("expected ToolResult at index 1");
}
}
/// FIX-F 验证:错误结果不被压缩,诊断信息完整保留。
#[test]
fn microcompact_preserves_error_tool_results() {
let mut messages = vec![
user_msg("hi"),
Message::tool_result("call_1", "important error info: backend down", true),
user_msg("keep recent 1"),
user_msg("keep recent 2"),
];
let before_len = messages.len();
let freed = microcompact(&mut messages, 2);
assert_eq!(freed, 0); // 错误 ToolResult 不计入
assert_eq!(messages.len(), before_len);
// 错误信息保留完整
if let Message::ToolResult { content, is_error, .. } = &messages[1] {
assert!(is_error);
assert!(matches!(&content[0], ContentBlock::Text { text } if text.contains("backend down")));
} else {
panic!("expected ToolResult at index 1");
}
}
#[test]
fn microcompact_keeps_recent_messages() {
let mut messages = vec![
Message::tool_result("c1", "old", false),
user_msg("m1"),
user_msg("m2"),
user_msg("recent"),
];
let freed = microcompact(&mut messages, 1);
// keep_recent=1 → 仅最近 1 条不动,其余 ToolResult 压缩
// 索引 0 (ToolResult) 被压缩,索引 1-3 保留
assert!(freed > 0);
if let Message::ToolResult { content, .. } = &messages[0] {
assert!(matches!(&content[0], ContentBlock::Text { text } if text == "[pruned]"));
}
assert!(matches!(messages[3], Message::User { .. }));
}
#[test]
fn should_compact_respects_threshold() {
let cfg = CompactConfig::default();
let state = CompactState::new();
let empty: Vec<Message> = vec![];
assert!(!should_compact(&empty, &cfg, &state));
}
}
+113 -221
View File
@@ -19,14 +19,9 @@ use crate::llm::hooks::{HookContext, HookExecutor};
use crate::llm::provider::{LlmProvider, ProviderCapabilities, ProviderFeatures};
use crate::llm::stream::StreamEvent;
use crate::llm::types::message::{ContentBlock, Message};
use crate::llm::types::openai_message::{
ContentField, OpenaiChatMessage, OpenaiContentPart,
};
use crate::llm::types::request_v2::MessageRequest;
use crate::llm::types::response_v2::{MessageResponse, StopReason};
use crate::llm::types::{
FinishReason, FunctionCall, OpenaiToolCall, ToolChoice, ToolDefinition,
};
use crate::llm::types::{ToolChoice, ToolDefinition};
/// LLM 调用周期配置。
pub struct CycleConfig {
@@ -67,12 +62,18 @@ impl Default for CycleConfig {
}
/// LLM 调用周期 —— 管理一次或多次 LLM 请求的生命周期。
///
/// Phase 2 修订:
/// - `messages` 字段从 `Vec<OpenaiChatMessage>` 切换为 `Vec<Message>`IR 类型)。
/// - 移除 `system_prompt` 字段(FIX-D)—— 调用方通过 `Message::system_text()` +
/// `with_messages()` 自行管理系统提示消息,避免重复维护两条路径。
/// - `with_system_prompt()` 方法标记 `#[deprecated]`,过渡期内仍可用。
pub struct LlmCycle {
provider: Arc<dyn LlmProvider>,
config: CycleConfig,
usage: CostTracker,
messages: Vec<OpenaiChatMessage>,
system_prompt: Option<String>,
/// 消息历史 —— 直接存储 IR `Message` 类型,build_request 不再做转换。
messages: Vec<Message>,
hook_executor: Option<Arc<HookExecutor>>,
compact_config: Option<CompactConfig>,
compact_state: CompactState,
@@ -96,16 +97,22 @@ impl LlmCycle {
config,
usage: CostTracker::default(),
messages: Vec::new(),
system_prompt: None,
hook_executor: None,
compact_config: None,
compact_state: CompactState::new(),
}
}
/// 设置系统提示词。
/// 设置系统提示词**已废弃** —— Phase 2 起使用 `Message::system_text()` + `with_messages()`
///
/// 当前实现为过渡期保留:在 messages 头部插入 `Message::System { content: [Text { text }] }`。
#[deprecated(
since = "0.2.0",
note = "请改用 Message::system_text() + with_messages()"
)]
pub fn with_system_prompt(mut self, prompt: String) -> Self {
self.system_prompt = Some(prompt);
self.messages
.insert(0, Message::System { content: vec![ContentBlock::Text { text: prompt }] });
self
}
@@ -126,8 +133,8 @@ impl LlmCycle {
&self.usage
}
/// 获取消息历史引用。
pub fn messages(&self) -> &[OpenaiChatMessage] {
/// 获取消息历史引用Phase 2:返回 `&[Message]` IR 类型)
pub fn messages(&self) -> &[Message] {
&self.messages
}
@@ -142,27 +149,33 @@ impl LlmCycle {
}
/// 直接设置消息历史(覆盖已有消息),支持 Builder 链式调用。
pub fn with_messages(mut self, messages: Vec<OpenaiChatMessage>) -> Self {
pub fn with_messages(mut self, messages: Vec<Message>) -> Self {
self.messages = messages;
self
}
/// 追加消息到历史尾部。
pub fn extend_messages(&mut self, messages: Vec<OpenaiChatMessage>) {
pub fn extend_messages(&mut self, messages: Vec<Message>) {
self.messages.extend(messages);
}
/// 追加单条消息到历史尾部。
///
/// 公开给 `submit_stream()` 消费方在收到 `MessageComplete` 事件后调用。
pub fn push_message(&mut self, msg: Message) {
self.messages.push(msg);
}
/// 使用预构建消息提交(跳过自动 push user prompt)。
///
/// 与 `submit()` 不同,不自动添加 `user_text(prompt)`,也不自动插入 system prompt
/// 与 `submit()` 不同,不自动添加 `user_text(prompt)`。
/// 调用方完全控制消息序列内容。
pub async fn submit_messages(
&mut self,
messages: Vec<OpenaiChatMessage>,
messages: Vec<Message>,
tools: Vec<ToolDefinition>,
) -> Result<MessageResponse, LlmError> {
let ir_messages: Vec<Message> =
messages.iter().map(chat_message_to_ir_message).collect();
let ir_messages: Vec<Message> = messages;
let ir_tools: Vec<ToolDefinition> = tools.clone();
let request = MessageRequest {
@@ -222,7 +235,7 @@ impl LlmCycle {
prompt: String,
tools: Vec<ToolDefinition>,
) -> Result<MessageResponse, LlmError> {
self.messages.push(OpenaiChatMessage::user_text(prompt));
self.messages.push(Message::user_text(prompt));
if let Some(ref config) = self.compact_config
&& should_compact(&self.messages, config, &self.compact_state)
@@ -265,9 +278,8 @@ impl LlmCycle {
.await;
}
// ponytail: Phase 0 中内部消息存储仍为 Vec<OpenaiChatMessage>
// 把 IR Message 转回旧 wire 格式以便存储。Phase 2 切换后移除。
self.messages.push(message_to_chat_message(&response.message));
// ponytail: Phase 2 直接存储 IR Message —— 不再转换。
self.messages.push(response.message.clone());
self.usage.add(&response.usage);
return Ok(response);
@@ -302,16 +314,38 @@ impl LlmCycle {
}
}
/// 提交用户消息并返回语义事件流。
/// 提交用户消息并返回语义事件流Phase 2 / FIX-E 简化方案)
///
/// 与 `submit` 不同,该方法返回流式事件而非完整响应。
/// 适用于需要实时处理 LLM 输出的场景。
///
/// **Phase 2 设计决策(FIX-E**
/// - `Item = StreamEvent`**不再是 `Result<StreamEvent, LlmError>`**
/// - 错误统一以 `StreamEvent::Error { message }` 形式在流中传出
/// - `self.messages` 不会在流结束后自动 push Assistant 响应 ——
/// **调用方** 在收到 `StreamEvent::MessageComplete` 后手动调用
/// `cycle.push_message(response.message.clone())` 完成消息历史追加
///
/// 调用模式(FIX-E 推荐):
/// ```ignore
/// use futures_util::StreamExt;
/// let mut stream = cycle.submit_stream(prompt, tools).await?;
/// let mut final_response: Option<MessageResponse> = None;
/// while let Some(event) = stream.next().await {
/// if let StreamEvent::MessageComplete { full_response } = &event {
/// final_response = Some(full_response.clone());
/// }
/// }
/// if let Some(resp) = final_response {
/// cycle.push_message(resp.message.clone());
/// }
/// ```
pub async fn submit_stream(
&mut self,
prompt: String,
tools: Vec<ToolDefinition>,
) -> Result<Pin<Box<dyn Stream<Item = StreamEvent> + Send>>, LlmError> {
self.messages.push(OpenaiChatMessage::user_text(prompt));
self.messages.push(Message::user_text(prompt));
if let Some(ref config) = self.compact_config
&& should_compact(&self.messages, config, &self.compact_state)
@@ -324,6 +358,7 @@ impl LlmCycle {
let request = self.build_request(&tools);
// PreRequest hook
if let Some(ref executor) = self.hook_executor {
let ctx = HookContext::new(crate::llm::hooks::HookEvent::PreRequest)
.with_request(&request);
@@ -342,11 +377,9 @@ impl LlmCycle {
let ir_event_stream = self.provider.chat_stream(request).await?;
let hook_executor = self.hook_executor.clone();
let post_request = self.build_request(&tools);
// ponytail: Phase 0 中 Provider 临时桥接层把 IR 事件再回填成本 phase 的
// 信息冗余(OpenaiProvider 的 chat_stream 现在已输出 IR 事件)。当前实现
// 直接转发 IR 事件,仅在 `MessageComplete` 到达时收尾。
// ponytail: Phase 2 简化方案(FIX-E)。流是延迟求值的,&mut self 无法进入闭包。
// 调用方在收到 MessageComplete 后手动调 push_message()。
Ok(Box::pin(stream! {
use futures_util::StreamExt;
let mut ir_event_stream = ir_event_stream;
@@ -354,17 +387,15 @@ impl LlmCycle {
while let Some(result) = ir_event_stream.next().await {
match result {
Ok(event) => {
let is_terminal = matches!(event, StreamEvent::MessageComplete { .. });
let is_error = matches!(event, StreamEvent::Error { .. });
let is_terminal =
matches!(event, StreamEvent::MessageComplete { .. } | StreamEvent::Error { .. });
yield event;
if is_terminal {
break;
}
if is_error {
break;
}
}
Err(e) => {
// ponytail: 错误不再走 Err 分支,统一以 StreamEvent::Error 传出。
if let Some(ref executor) = hook_executor {
let ctx = crate::llm::hooks::HookContext::new(
crate::llm::hooks::HookEvent::OnError,
@@ -380,11 +411,13 @@ impl LlmCycle {
}
}
// ponytail: post_request hook 收到的消息列表**不包含本次 Assistant 响应**
// (因为流是延迟求值,本轮响应还未到达)。调用方如需完整上下文,
// 应在收到 MessageComplete 后手动触发 hook。
if let Some(ref executor) = hook_executor {
let ctx = crate::llm::hooks::HookContext::new(
crate::llm::hooks::HookEvent::PostRequest,
)
.with_request(&post_request);
);
executor
.execute(crate::llm::hooks::HookEvent::PostRequest, &ctx)
.await;
@@ -393,25 +426,11 @@ impl LlmCycle {
}
fn build_request(&self, tools: &[ToolDefinition]) -> MessageRequest {
let messages = self.messages.clone();
let ir_messages: Vec<Message> = messages
.iter()
.map(chat_message_to_ir_message)
.collect();
let mut ir_messages = ir_messages;
if let Some(sys_prompt) = &self.system_prompt
&& !ir_messages
.iter()
.any(|m| matches!(m, Message::System { .. }))
{
ir_messages.insert(0, Message::system(sys_prompt.as_str()));
}
// ponytail: Phase 2 简化 —— 直接 clone self.messages,无任何转换 / system prompt 注入。
// 系统消息如需存在,由调用方通过 `with_messages()` 自行管理。
MessageRequest {
model: self.config.model.clone(),
messages: ir_messages,
messages: self.messages.clone(),
tools: tools.to_vec(),
tool_choice: ToolChoice::Auto,
max_tokens: self.config.max_tokens,
@@ -510,7 +529,7 @@ impl LlmCycle {
let tool_timeout = self.config.tool_timeout_secs;
let max_bytes = self.config.max_tool_result_bytes;
self.messages.push(OpenaiChatMessage::user_text(prompt));
self.messages.push(Message::user_text(prompt));
self.maybe_compact();
let mut turn = 0;
@@ -529,9 +548,8 @@ impl LlmCycle {
let should_execute = matches!(response.stop_reason, StopReason::ToolUse)
&& has_tool_calls_in_response(&response);
// ponytail: Phase 0 中内部消息存储仍为 Vec<OpenaiChatMessage>
// 把 IR Message 转回旧 wire 格式以便存储。Phase 2 切换后移除。
self.messages.push(message_to_chat_message(&response.message));
// ponytail: Phase 2 直接存储 IR Message —— 不再转换。
self.messages.push(response.message.clone());
if !should_execute {
return Ok(response);
@@ -539,12 +557,12 @@ impl LlmCycle {
// 解析 tool_calls 并执行
let tool_calls = extract_tool_calls_from_response(&response);
let calls: Vec<(String, serde_json::Value)> = tool_calls
let calls: Vec<(String, String, serde_json::Value)> = tool_calls
.into_iter()
.map(|(_id, name, args)| {
.map(|(id, name, args)| {
let value: serde_json::Value =
serde_json::from_str(&args).unwrap_or(serde_json::Value::Null);
(name, value)
(id, name, value)
})
.collect();
@@ -552,9 +570,10 @@ impl LlmCycle {
// 回传工具结果
for result in results {
let content = match result.output {
let is_error = result.output.is_err();
let content = match &result.output {
Ok(value) => {
let serialized = serde_json::to_string(&value).unwrap_or_else(|e| {
let serialized = serde_json::to_string(value).unwrap_or_else(|e| {
tracing::warn!("工具结果序列化失败: {}", e);
"{}".to_string()
});
@@ -570,8 +589,16 @@ impl LlmCycle {
}
};
self.messages
.push(OpenaiChatMessage::tool_result(result.tool_name, content));
// ponytail: 当前 self.messages 仍是 Vec<OpenaiChatMessage> (Phase 2 切换后
// 改为 Message::tool_result 并传入 result.tool_call_id)。当前实现已经使用
// 真实 tool_call_id 而非 tool_name 充当 —— 这条 FIX-A 修复与 Phase 2 消息切换
// 同步生效。
// ponytail: Phase 2 直接存储 Message::ToolResultis_error 由 ToolInvocation.output 推断。
self.messages.push(Message::tool_result(
result.tool_call_id,
content,
is_error,
));
}
// 每轮工具执行后触发 compaction
@@ -630,157 +657,6 @@ fn extract_tool_calls_from_response(
out
}
/// 把 `OpenaiChatMessage` 转换为 IR `Message`。
///
/// ponytail: Phase 0 转换层 —— build_request 中需要把存储用的
/// `Vec<OpenaiChatMessage>` 转为 `MessageRequest.messages`。Phase 2 中
/// `self.messages` 直接改为 `Vec<Message>` 时此函数整体删除。
fn chat_message_to_ir_message(msg: &OpenaiChatMessage) -> Message {
match msg {
OpenaiChatMessage::Developer { content, .. }
| OpenaiChatMessage::System { content, .. } => Message::System {
content: content_field_to_blocks(content),
},
OpenaiChatMessage::User { content, .. } => Message::User {
content: content_field_to_blocks(content),
},
OpenaiChatMessage::Assistant {
content,
tool_calls,
..
} => {
let mut blocks = content_field_to_blocks(content);
if let Some(calls) = tool_calls {
for call in calls {
if let OpenaiToolCall::Function { id, function } = call {
let input: serde_json::Value = serde_json::from_str(&function.arguments)
.unwrap_or(serde_json::Value::Null);
blocks.push(ContentBlock::ToolUse {
id: id.clone(),
name: function.name.clone(),
input,
});
}
}
}
Message::Assistant { content: blocks }
}
OpenaiChatMessage::Tool {
content,
tool_call_id,
} => Message::ToolResult {
tool_call_id: tool_call_id.clone(),
content: content_field_to_blocks(content),
is_error: false,
},
OpenaiChatMessage::Function { content, name } => Message::ToolResult {
tool_call_id: name.clone(),
content: content_field_to_blocks(content),
is_error: false,
},
}
}
/// 把 IR `Message` 转换为 `OpenaiChatMessage`(用于 `self.messages` 持久化)。
///
/// ponytail: Phase 0 转换层 —— submit 后 `response.message: Message` 需要
/// 转回 `OpenaiChatMessage` 才能追加到旧 store。Phase 2 中移除。
fn message_to_chat_message(msg: &Message) -> OpenaiChatMessage {
match msg {
Message::System { content } => OpenaiChatMessage::System {
content: blocks_to_content_field(content),
name: None,
},
Message::User { content } => OpenaiChatMessage::User {
content: blocks_to_content_field(content),
name: None,
},
// ponytail: UserImage 当前无 OpenAI 对应,回退为空的 User 消息(图片数据丢失)。
// Phase 2 切换后此函数整体移除。
Message::UserImage { .. } => OpenaiChatMessage::User {
content: ContentField::Array(vec![]),
name: None,
},
Message::Assistant { content } => {
let mut text_blocks: Vec<String> = Vec::new();
let mut tool_call_blocks: Vec<OpenaiToolCall> = Vec::new();
for block in content {
match block {
ContentBlock::Text { text } => text_blocks.push(text.clone()),
ContentBlock::ToolUse { id, name, input } => {
tool_call_blocks.push(OpenaiToolCall::Function {
id: id.clone(),
function: FunctionCall {
name: name.clone(),
arguments: serde_json::to_string(input)
.unwrap_or_else(|_| "null".to_string()),
},
});
}
// ponytail: thinking/refusal/image/audio/file 在 Phase 0 持久化时被丢弃。
// Phase 2 切换后此函数整体移除,自然修复。
_ => {}
}
}
let text: String = text_blocks.into_iter().collect();
OpenaiChatMessage::Assistant {
content: if text.is_empty() {
ContentField::Array(vec![])
} else {
ContentField::String(text)
},
refusal: None,
name: None,
tool_calls: if tool_call_blocks.is_empty() {
None
} else {
Some(tool_call_blocks)
},
}
}
Message::ToolResult {
tool_call_id,
content,
is_error: _,
} => OpenaiChatMessage::Tool {
content: blocks_to_content_field(content),
tool_call_id: tool_call_id.clone(),
},
}
}
fn blocks_to_content_field(blocks: &[ContentBlock]) -> ContentField {
if let [ContentBlock::Text { text }] = blocks {
return ContentField::String(text.clone());
}
let parts: Vec<OpenaiContentPart> = blocks
.iter()
.filter_map(|b| match b {
ContentBlock::Text { text } => Some(OpenaiContentPart::Text { text: text.clone() }),
_ => None,
})
.collect();
ContentField::Array(parts)
}
fn content_field_to_blocks(field: &ContentField) -> Vec<ContentBlock> {
match field {
ContentField::String(s) => vec![ContentBlock::Text { text: s.clone() }],
ContentField::Array(parts) => parts
.iter()
.filter_map(|p| match p {
OpenaiContentPart::Text { text } => {
Some(ContentBlock::Text { text: text.clone() })
}
OpenaiContentPart::Refusal { refusal } => {
Some(ContentBlock::Text { text: refusal.clone() })
}
_ => None,
})
.collect(),
}
}
/// 截断工具结果到指定字节数。
fn truncate_tool_result(s: &str, max_bytes: usize) -> String {
if s.len() <= max_bytes {
@@ -929,14 +805,30 @@ mod tests {
// 最终响应是 Assistant 消息
assert!(matches!(response.message, Message::Assistant { .. }));
// 验证消息历史Vec<OpenaiChatMessage>Phase 2 才切到 Vec<Message>
// user, assistant(tool_calls → 转回后含 tool_use), tool, assistant(text)
// 验证消息历史:
// user, assistant(含 tool_use), tool_result, assistant(text)
let messages = cycle.messages();
assert_eq!(messages.len(), 4);
assert!(matches!(messages[0], OpenaiChatMessage::User { .. }));
assert!(matches!(messages[1], OpenaiChatMessage::Assistant { .. }));
assert!(matches!(messages[2], OpenaiChatMessage::Tool { .. }));
assert!(matches!(messages[3], OpenaiChatMessage::Assistant { .. }));
assert!(matches!(messages[0], Message::User { .. }));
assert!(matches!(
messages[1],
Message::Assistant {
content: _,
}
));
if let Message::Assistant { content } = &messages[1] {
assert!(content
.iter()
.any(|b| matches!(b, ContentBlock::ToolUse { .. })));
}
assert!(matches!(
messages[2],
Message::ToolResult {
is_error: false,
..
}
));
assert!(matches!(messages[3], Message::Assistant { .. }));
}
#[tokio::test]
+52 -39
View File
@@ -2,17 +2,16 @@
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use crate::llm::compact::{CompactConfig, CompactState, microcompact, should_compact};
use crate::llm::types::OpenaiChatMessage;
use crate::llm::types::message::Message;
use crate::memory::error::MemoryError;
use crate::memory::store::MemoryStore;
use crate::memory::types::MemoryItem;
/// 对话消息管理策略。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum MemoryStrategy {
/// 滑动窗口:达到上限时删除最旧消息。
SlidingWindow,
@@ -40,22 +39,30 @@ impl Default for ConversationMemoryConfig {
/// 对话记忆 —— 按 session 管理多轮对话消息历史。
///
/// 内部维护 `Vec<OpenaiChatMessage>` 热缓存(供 `llm::compact` 直接操作)
/// 内部维护 `Vec<Message>`Phase 2 IR 类型)作为热缓存
/// `MemoryStore` 用作冷持久化层。
///
/// ponytail: Phase 2 切换消息存储从 `OpenaiChatMessage` 到 `Message`。
/// `Message` 已实现 `Serialize` / `Deserialize`Phase 0 FIX-B 引入),
/// 序列化格式采用 `#[serde(tag = "type", rename_all = "snake_case")]`
/// 例如:
/// ```json
/// {"type": "user", "content": [{"type": "text", "text": "hi"}]}
/// {"type": "tool_result", "tool_call_id": "c1", "content": [...], "is_error": false}
/// ```
pub struct ConversationMemory {
store: Arc<dyn MemoryStore>,
session_id: String,
config: ConversationMemoryConfig,
/// 热缓存:消息列表,供 `llm::compact` 直接操作。
messages: Vec<OpenaiChatMessage>,
/// 与 `messages` 一一对应的存储 ID(保持稳定以便淘汰时精准删除)
messages: Vec<Message>,
/// 与 `messages` 一一对应的存储 ID。
message_ids: Vec<String>,
/// 压缩断路器状态。
compact_state: CompactState,
}
impl ConversationMemory {
/// 创建一个新的 ConversationMemory。
pub fn new(
store: Arc<dyn MemoryStore>,
session_id: impl Into<String>,
@@ -71,26 +78,23 @@ impl ConversationMemory {
}
}
/// 获取 session id。
pub fn session_id(&self) -> &str {
&self.session_id
}
/// 获取配置。
pub fn config(&self) -> &ConversationMemoryConfig {
&self.config
}
/// 从 MemoryStore 加载历史消息到热缓存。
pub async fn load(&mut self) -> Result<(), MemoryError> {
let filter = crate::memory::types::MemoryFilter {
prefix: Some(self.session_prefix()),
..Default::default()
};
let items = self.store.list(&filter).await?;
let mut pairs: Vec<(String, OpenaiChatMessage, OffsetDateTime)> = Vec::with_capacity(items.len());
let mut pairs: Vec<(String, Message, OffsetDateTime)> = Vec::with_capacity(items.len());
for item in items {
match serde_json::from_str::<OpenaiChatMessage>(&item.content) {
match serde_json::from_str::<Message>(&item.content) {
Ok(msg) => pairs.push((item.id, msg, item.created_at)),
Err(e) => {
return Err(MemoryError::Serialization(format!(
@@ -100,17 +104,13 @@ impl ConversationMemory {
}
}
}
// 按 created_at 升序排列
pairs.sort_by_key(|p| p.2);
self.message_ids = pairs.iter().map(|p| p.0.clone()).collect();
self.messages = pairs.into_iter().map(|p| p.1).collect();
Ok(())
}
/// 添加一条消息。
///
/// 写入热缓存并通过 `MemoryStore` 持久化。如有需要,触发淘汰和压缩。
pub async fn add_message(&mut self, msg: OpenaiChatMessage) -> Result<(), MemoryError> {
pub async fn add_message(&mut self, msg: Message) -> Result<(), MemoryError> {
let now = OffsetDateTime::now_utc();
let index = self.messages.len();
let id = self.make_message_id(index, &now);
@@ -119,7 +119,7 @@ impl ConversationMemory {
self.messages.push(msg);
self.message_ids.push(id.clone());
// 同步到冷存储
// ponytail: 通过 `Message` 的 Serialize 派生实现持久化
let item = MemoryItem {
id: id.clone(),
content: serde_json::to_string(self.messages.last().unwrap())
@@ -129,17 +129,14 @@ impl ConversationMemory {
};
self.store.save(item).await?;
// 触发淘汰和压缩
self.maybe_evict_and_compact().await;
Ok(())
}
/// 获取完整消息历史。
pub fn get_history(&self) -> &[OpenaiChatMessage] {
pub fn get_history(&self) -> &[Message] {
&self.messages
}
/// 清空所有消息。
pub async fn clear(&mut self) -> Result<(), MemoryError> {
let to_delete = std::mem::take(&mut self.message_ids);
self.messages.clear();
@@ -150,18 +147,16 @@ impl ConversationMemory {
Ok(())
}
/// 当前消息数量。
pub fn len(&self) -> usize {
self.messages.len()
}
/// 是否为空。
pub fn is_empty(&self) -> bool {
self.messages.is_empty()
}
fn session_prefix(&self) -> String {
format!("conv:{self}:", self = self.session_id)
format!("conv:{}:", self.session_id)
}
fn make_message_id(&self, index: usize, now: &OffsetDateTime) -> String {
@@ -169,7 +164,6 @@ impl ConversationMemory {
}
async fn maybe_evict_and_compact(&mut self) {
// 1. Sliding window 淘汰:删除最旧消息
if self.config.strategy == MemoryStrategy::SlidingWindow {
while self.messages.len() > self.config.max_turns {
if let Some(removed_id) = self.message_ids.first().cloned() {
@@ -180,7 +174,6 @@ impl ConversationMemory {
}
}
// 2. 压缩(复用 llm::compact
if let Some(ref compact_config) = self.config.compact_config {
if should_compact(&self.messages, compact_config, &self.compact_state) {
let keep_recent = compact_config.keep_recent;
@@ -188,7 +181,6 @@ impl ConversationMemory {
if freed > 0 {
self.compact_state.record_success();
} else {
// 没有 token 被释放(可能没找到可压缩的 tool result
let _ = self.compact_state.record_failure();
}
}
@@ -199,24 +191,42 @@ impl ConversationMemory {
#[cfg(test)]
mod tests {
use super::*;
use crate::llm::types::OpenaiChatMessage;
use crate::memory::InMemoryStore;
use crate::memory::MemoryStore;
fn user_text(s: &str) -> OpenaiChatMessage {
OpenaiChatMessage::user_text(s)
}
#[tokio::test]
async fn add_and_get_history() {
let store = Arc::new(InMemoryStore::new()) as Arc<dyn MemoryStore>;
let mut conv = ConversationMemory::new(store, "session1", ConversationMemoryConfig::default());
conv.add_message(user_text("hello")).await.unwrap();
conv.add_message(user_text("world")).await.unwrap();
conv.add_message(Message::user_text("hello")).await.unwrap();
conv.add_message(Message::user_text("world")).await.unwrap();
assert_eq!(conv.len(), 2);
assert_eq!(conv.get_history().len(), 2);
}
/// 验证 Message → JSON → Message 往返(包含 ToolResult 等完整信息)。
#[tokio::test]
async fn json_roundtrip_preserves_tool_result() {
let store = Arc::new(InMemoryStore::new()) as Arc<dyn MemoryStore>;
let mut conv = ConversationMemory::new(store, "s1", ConversationMemoryConfig::default());
conv.add_message(Message::tool_result("call_1", "ok", false))
.await
.unwrap();
conv.add_message(Message::assistant("done"))
.await
.unwrap();
let original = conv.get_history().to_vec();
assert_eq!(original.len(), 2);
// 各变体可序列化 + 反序列化
for msg in &original {
let json = serde_json::to_string(msg).unwrap();
let decoded: Message = serde_json::from_str(&json).unwrap();
assert_eq!(format!("{:?}", decoded), format!("{:?}", msg));
}
}
#[tokio::test]
async fn sliding_window_evicts_oldest() {
let store = Arc::new(InMemoryStore::new()) as Arc<dyn MemoryStore>;
@@ -227,7 +237,9 @@ mod tests {
};
let mut conv = ConversationMemory::new(store, "s1", config);
for i in 0..5 {
conv.add_message(user_text(&format!("msg-{i}"))).await.unwrap();
conv.add_message(Message::user_text(&format!("msg-{i}")))
.await
.unwrap();
}
assert_eq!(conv.len(), 3);
}
@@ -242,9 +254,10 @@ mod tests {
};
let mut conv = ConversationMemory::new(store, "s1", config);
for i in 0..5 {
conv.add_message(user_text(&format!("msg-{i}"))).await.unwrap();
conv.add_message(Message::user_text(&format!("msg-{i}")))
.await
.unwrap();
}
// Full 策略不删除消息
assert_eq!(conv.len(), 5);
}
@@ -252,7 +265,7 @@ mod tests {
async fn clear_empties_messages() {
let store = Arc::new(InMemoryStore::new()) as Arc<dyn MemoryStore>;
let mut conv = ConversationMemory::new(store.clone(), "s1", ConversationMemoryConfig::default());
conv.add_message(user_text("hello")).await.unwrap();
conv.add_message(Message::user_text("hello")).await.unwrap();
assert!(!conv.is_empty());
conv.clear().await.unwrap();
assert!(conv.is_empty());
+55 -16
View File
@@ -15,6 +15,12 @@ use crate::tools::permission::PermissionChecker;
/// 工具调用记录 —— 用于追踪和调试。
#[derive(Debug, Clone)]
pub struct ToolInvocation {
/// LLM 返回的 tool_call_id —— 用于回传 `Message::ToolResult` 时关联原始调用。
///
/// ponytail: Phase 2 引入。老的循环用 `tool_name` 冒充 tool_call_id
/// 对 OpenAI 碰巧可用,对 Anthropic 必然失败。Anthropic 协议要求
/// `tool_result.tool_use_id` 与上一轮 `tool_use.id` 严格一致。
pub tool_call_id: String,
/// 被调用的工具名。
pub tool_name: String,
/// 工具的入参。
@@ -25,8 +31,14 @@ pub struct ToolInvocation {
impl ToolInvocation {
/// 创建一个新的工具调用记录。
pub fn new(tool_name: String, input: Value, output: Result<Value, ToolError>) -> Self {
pub fn new(
tool_call_id: String,
tool_name: String,
input: Value,
output: Result<Value, ToolError>,
) -> Self {
Self {
tool_call_id,
tool_name,
input,
output,
@@ -128,7 +140,15 @@ impl ToolRegistry {
}
/// 调用单个工具(含权限检查)。
pub async fn invoke(&self, name: &str, args: Value) -> Result<ToolInvocation, ToolError> {
///
/// `tool_call_id` 来源于 LLM 流式响应中的 `tool_calls[i].id`,用于回传
/// 工具结果时与原始 `tool_use` block 关联。
pub async fn invoke(
&self,
tool_call_id: &str,
name: &str,
args: Value,
) -> Result<ToolInvocation, ToolError> {
let tool = self
.get(name)
.ok_or_else(|| ToolError::NotFound(name.to_string()))?;
@@ -139,35 +159,48 @@ impl ToolRegistry {
let ctx = ToolContext::new(name, "");
let output = tool.execute(args.clone(), &ctx).await;
Ok(ToolInvocation::new(name.to_string(), args, output))
Ok(ToolInvocation::new(
tool_call_id.to_string(),
name.to_string(),
args,
output,
))
}
/// 并行执行多个工具调用(互不依赖的工具)。
///
/// 每个工具独立超时(`timeout_per_call_secs`0 表示不超时)。
/// 单个工具超时不会影响其他工具的返回。
///
/// 入参元组为 `(tool_call_id, tool_name, args)` —— `tool_call_id` 来自 LLM 响应。
pub async fn invoke_all(
&self,
calls: Vec<(String, Value)>,
calls: Vec<(String, String, Value)>,
timeout_per_call_secs: u64,
) -> Vec<ToolInvocation> {
let this = self.clone();
let futures = calls.into_iter().map(|(name, args)| {
let futures = calls.into_iter().map(|(tool_call_id, name, args)| {
let this = this.clone();
async move {
match if timeout_per_call_secs == 0 {
Ok(this.invoke(&name, args.clone()).await)
Ok(this.invoke(&tool_call_id, &name, args.clone()).await)
} else {
tokio::time::timeout(
Duration::from_secs(timeout_per_call_secs),
this.invoke(&name, args.clone()),
this.invoke(&tool_call_id, &name, args.clone()),
)
.await
} {
Ok(result) => result.unwrap_or_else(|e| {
ToolInvocation::new(name.clone(), args.clone(), Err(e))
ToolInvocation::new(
tool_call_id.clone(),
name.clone(),
args.clone(),
Err(e),
)
}),
Err(_) => ToolInvocation::new(
tool_call_id,
name,
args,
Err(ToolError::McpTimeout("timeout".into())),
@@ -313,15 +346,17 @@ mod tests {
async fn test_invoke_success() {
let mut reg = ToolRegistry::new();
reg.register(Arc::new(AddTool { base: 100 })).unwrap();
let result = reg.invoke("add", json!({ "n": 5 })).await.unwrap();
let result = reg.invoke("call_1", "add", json!({ "n": 5 })).await.unwrap();
let value = result.output.unwrap();
assert_eq!(value["result"], 105);
assert_eq!(result.tool_call_id, "call_1");
assert_eq!(result.tool_name, "add");
}
#[tokio::test]
async fn test_invoke_not_found() {
let reg = ToolRegistry::new();
let result = reg.invoke("nope", json!({})).await;
let result = reg.invoke("call_x", "nope", json!({})).await;
assert!(matches!(result, Err(ToolError::NotFound(_))));
}
@@ -329,7 +364,7 @@ mod tests {
async fn test_invoke_execution_error() {
let mut reg = ToolRegistry::new();
reg.register(Arc::new(FailTool)).unwrap();
let result = reg.invoke("fail", json!({})).await.unwrap();
let result = reg.invoke("call_y", "fail", json!({})).await.unwrap();
assert!(result.output.is_err());
}
@@ -338,7 +373,7 @@ mod tests {
let mut reg = ToolRegistry::new()
.with_permission_checker(PermissionChecker::new(Default::default()));
reg.register(Arc::new(ShellTool)).unwrap();
let result = reg.invoke("shell", json!({})).await;
let result = reg.invoke("call_z", "shell", json!({})).await;
assert!(matches!(result, Err(ToolError::PermissionDenied(_, _))));
}
@@ -348,12 +383,15 @@ mod tests {
reg.register(Arc::new(AddTool { base: 1 })).unwrap();
reg.register(Arc::new(FailTool)).unwrap();
let calls = vec![
("add".into(), json!({ "n": 1 })),
("add".into(), json!({ "n": 2 })),
("fail".into(), json!({})),
("c1".into(), "add".into(), json!({ "n": 1 })),
("c2".into(), "add".into(), json!({ "n": 2 })),
("c3".into(), "fail".into(), json!({})),
];
let results = reg.invoke_all(calls, 0).await;
assert_eq!(results.len(), 3);
assert_eq!(results[0].tool_call_id, "c1");
assert_eq!(results[1].tool_call_id, "c2");
assert_eq!(results[2].tool_call_id, "c3");
assert!(results[0].output.is_ok());
assert!(results[1].output.is_ok());
assert!(results[2].output.is_err());
@@ -363,9 +401,10 @@ mod tests {
async fn test_invoke_all_with_timeout() {
let mut reg = ToolRegistry::new();
reg.register(Arc::new(AddTool { base: 0 })).unwrap();
let calls = vec![("add".into(), json!({ "n": 1 }))];
let calls = vec![("c1".into(), "add".into(), json!({ "n": 1 }))];
let results = reg.invoke_all(calls, 5).await;
assert_eq!(results.len(), 1);
assert_eq!(results[0].tool_call_id, "c1");
assert!(results[0].output.is_ok());
}
}