style(tools, llm): 统一导入顺序与代码格式

This commit is contained in:
徐涛
2026-07-05 08:19:13 +08:00
parent 98dfe6c1ed
commit 5648b1d217
39 changed files with 390 additions and 397 deletions
+6 -7
View File
@@ -15,9 +15,9 @@ use std::sync::Arc;
use agcore::agent::{Agent, AgentBuilder, AgentSession}; use agcore::agent::{Agent, AgentBuilder, AgentSession};
use agcore::llm::hooks::HookExecutor; use agcore::llm::hooks::HookExecutor;
use agcore::llm::mock::MockProvider; use agcore::llm::mock::MockProvider;
use agcore::llm::types::Usage;
use agcore::llm::types::message::{ContentBlock, Message}; use agcore::llm::types::message::{ContentBlock, Message};
use agcore::llm::types::response_v2::{MessageResponse, StopReason}; use agcore::llm::types::response_v2::{MessageResponse, StopReason};
use agcore::llm::types::Usage;
use agcore::tools::ToolRegistry; use agcore::tools::ToolRegistry;
/// 计算器角色 Agent。 /// 计算器角色 Agent。
@@ -72,7 +72,10 @@ async fn main() {
// 4. 提交第一轮 // 4. 提交第一轮
println!("=== 提交第 1 轮 ==="); println!("=== 提交第 1 轮 ===");
let resp = session.submit_turn("1+1=?").await.expect("submit_turn 失败"); let resp = session
.submit_turn("1+1=?")
.await
.expect("submit_turn 失败");
println!("LLM: {}", resp.text()); println!("LLM: {}", resp.text());
session session
.set_session_data("last_q", "1+1=?") .set_session_data("last_q", "1+1=?")
@@ -107,11 +110,7 @@ async fn main() {
// 8. 跨 session 数据隔离验证 // 8. 跨 session 数据隔离验证
println!("=== 数据隔离验证 ==="); println!("=== 数据隔离验证 ===");
let other = AgentSession::new( let other = AgentSession::new(Arc::new(CalculatorAgent), "other-session", bundle);
Arc::new(CalculatorAgent),
"other-session",
bundle,
);
assert!( assert!(
other.get_session_data("last_q").await.unwrap().is_none(), other.get_session_data("last_q").await.unwrap().is_none(),
"新会话不应看到旧 session 的 last_q" "新会话不应看到旧 session 的 last_q"
+7 -16
View File
@@ -80,11 +80,8 @@ async fn main() {
// 3. 多角色混合 + clear // 3. 多角色混合 + clear
println!("\n=== 多角色写入 + clear ==="); println!("\n=== 多角色写入 + clear ===");
let store3 = Arc::new(InMemoryStore::new()); let store3 = Arc::new(InMemoryStore::new());
let mut memory3 = ConversationMemory::new( let mut memory3 =
store3, ConversationMemory::new(store3, "session-3", ConversationMemoryConfig::default());
"session-3",
ConversationMemoryConfig::default(),
);
memory3 memory3
.add_message(Message::user_text("你好")) .add_message(Message::user_text("你好"))
.await .await
@@ -98,7 +95,9 @@ async fn main() {
.await .await
.unwrap(); .unwrap();
memory3 memory3
.add_message(Message::assistant("我无法查询实时天气,但你可以查看天气应用。")) .add_message(Message::assistant(
"我无法查询实时天气,但你可以查看天气应用。",
))
.await .await
.unwrap(); .unwrap();
println!( println!(
@@ -119,16 +118,8 @@ async fn main() {
// 4. Session 隔离 // 4. Session 隔离
println!("\n=== Session 隔离(共用 InMemoryStore==="); println!("\n=== Session 隔离(共用 InMemoryStore===");
let store4 = Arc::new(InMemoryStore::new()); let store4 = Arc::new(InMemoryStore::new());
let mut a = ConversationMemory::new( let mut a = ConversationMemory::new(store4.clone(), "s-a", ConversationMemoryConfig::default());
store4.clone(), let mut b = ConversationMemory::new(store4.clone(), "s-b", ConversationMemoryConfig::default());
"s-a",
ConversationMemoryConfig::default(),
);
let mut b = ConversationMemory::new(
store4.clone(),
"s-b",
ConversationMemoryConfig::default(),
);
a.add_message(Message::user_text("A 的消息")).await.unwrap(); a.add_message(Message::user_text("A 的消息")).await.unwrap();
b.add_message(Message::user_text("B 的消息")).await.unwrap(); b.add_message(Message::user_text("B 的消息")).await.unwrap();
println!( println!(
+10 -15
View File
@@ -17,7 +17,7 @@ use agcore::tools::{
ToolRegistry, ToolRegistry,
}; };
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::{json, Value}; use serde_json::{Value, json};
/// 天气查询工具 —— 模拟根据城市返回天气数据。 /// 天气查询工具 —— 模拟根据城市返回天气数据。
struct WeatherTool; struct WeatherTool;
@@ -42,11 +42,7 @@ impl BaseTool for WeatherTool {
fn required_permissions(&self) -> Vec<Permission> { fn required_permissions(&self) -> Vec<Permission> {
vec![Permission::Network] vec![Permission::Network]
} }
async fn execute( async fn execute(&self, args: Value, _ctx: &ToolContext<'_>) -> Result<Value, ToolError> {
&self,
args: Value,
_ctx: &ToolContext<'_>,
) -> Result<Value, ToolError> {
let city = args["city"].as_str().unwrap_or("未知"); let city = args["city"].as_str().unwrap_or("未知");
// 模拟查询:根据城市名给出不同温度 // 模拟查询:根据城市名给出不同温度
let (temperature, condition) = match city { let (temperature, condition) = match city {
@@ -84,11 +80,7 @@ impl BaseTool for DeleteFileTool {
fn required_permissions(&self) -> Vec<Permission> { fn required_permissions(&self) -> Vec<Permission> {
vec![Permission::Delete] vec![Permission::Delete]
} }
async fn execute( async fn execute(&self, _args: Value, _ctx: &ToolContext<'_>) -> Result<Value, ToolError> {
&self,
_args: Value,
_ctx: &ToolContext<'_>,
) -> Result<Value, ToolError> {
Ok(json!({"deleted": true})) Ok(json!({"deleted": true}))
} }
} }
@@ -138,9 +130,8 @@ async fn main() {
// 5. 权限检查:默认 PermissionConfig 黑名单含 Delete // 5. 权限检查:默认 PermissionConfig 黑名单含 Delete
println!("\n=== 权限检查(默认 PermissionConfigdenied = [Delete, Shell]==="); println!("\n=== 权限检查(默认 PermissionConfigdenied = [Delete, Shell]===");
let mut registry_with_checker = ToolRegistry::new().with_permission_checker(PermissionChecker::new( let mut registry_with_checker = ToolRegistry::new()
PermissionConfig::default(), .with_permission_checker(PermissionChecker::new(PermissionConfig::default()));
));
registry_with_checker registry_with_checker
.register(Arc::new(WeatherTool) as ToolRef) .register(Arc::new(WeatherTool) as ToolRef)
.unwrap(); .unwrap();
@@ -155,7 +146,11 @@ async fn main() {
.unwrap(); .unwrap();
println!( println!(
"get_weather 权限检查: {}", "get_weather 权限检查: {}",
if r.output.is_ok() { "通过 ✓" } else { "阻断 ✗" } if r.output.is_ok() {
"通过 ✓"
} else {
"阻断 ✗"
}
); );
// delete_file 声明 Delete → 在 denied 列表 → 阻断 // delete_file 声明 Delete → 在 denied 列表 → 阻断
+19 -19
View File
@@ -38,14 +38,26 @@ async fn main() {
let ks = KnowledgeStore::new(store); let ks = KnowledgeStore::new(store);
let pages = vec![ let pages = vec![
make_page("rust-1", "Rust 入门", "Rust 是一门系统级编程语言,注重安全性与并发。"), make_page(
make_page("python-1", "Python 简介", "Python 是一门动态类型的高级编程语言。"), "rust-1",
"Rust 入门",
"Rust 是一门系统级编程语言,注重安全性与并发。",
),
make_page(
"python-1",
"Python 简介",
"Python 是一门动态类型的高级编程语言。",
),
make_page( make_page(
"langgraph-1", "langgraph-1",
"LangGraph 框架", "LangGraph 框架",
"LangGraph 是 LangChain 的状态图扩展,用于构建多步 Agent。", "LangGraph 是 LangChain 的状态图扩展,用于构建多步 Agent。",
), ),
make_page("rust-async", "Rust 异步编程", "Rust 异步基于 tokio 与 futures 抽象。"), make_page(
"rust-async",
"Rust 异步编程",
"Rust 异步基于 tokio 与 futures 抽象。",
),
]; ];
for p in &pages { for p in &pages {
ks.add_page(p.clone()).await.expect("保存页面失败"); ks.add_page(p.clone()).await.expect("保存页面失败");
@@ -62,14 +74,8 @@ async fn main() {
let result = retriever.retrieve("Rust 异步").await.unwrap(); let result = retriever.retrieve("Rust 异步").await.unwrap();
println!("query: {}", result.query); println!("query: {}", result.query);
for item in &result.items { for item in &result.items {
println!( println!(" 命中: {} (score={:.3})", item.page.title, item.score);
" 命中: {} (score={:.3})", assert!((0.0..=1.0).contains(&item.score), "score 应在 [0, 1] 区间");
item.page.title, item.score
);
assert!(
(0.0..=1.0).contains(&item.score),
"score 应在 [0, 1] 区间"
);
} }
assert!(!result.items.is_empty(), "应至少命中一个页面"); assert!(!result.items.is_empty(), "应至少命中一个页面");
@@ -85,14 +91,8 @@ async fn main() {
min_score: 0.5, min_score: 0.5,
}; };
let retriever2 = MemoryRetriever::new(ks2, cfg); let retriever2 = MemoryRetriever::new(ks2, cfg);
let result = retriever2 let result = retriever2.retrieve("完全不相关的火锅配方").await.unwrap();
.retrieve("完全不相关的火锅配方") println!("无关 query → items.len = {} (期望 0)", result.items.len());
.await
.unwrap();
println!(
"无关 query → items.len = {} (期望 0)",
result.items.len()
);
assert!(result.items.is_empty()); assert!(result.items.is_empty());
// 4. max_results 截断 // 4. max_results 截断
+7 -5
View File
@@ -11,7 +11,7 @@
use agcore::llm::types::message::{ContentBlock, Message}; use agcore::llm::types::message::{ContentBlock, Message};
use agcore::prompt::{ use agcore::prompt::{
validate_messages, PromptComposer, PromptTemplate, PromptTemplateRegistry, TemplateContext, PromptComposer, PromptTemplate, PromptTemplateRegistry, TemplateContext, validate_messages,
}; };
fn message_text(msg: &Message) -> String { fn message_text(msg: &Message) -> String {
@@ -33,9 +33,8 @@ fn message_text(msg: &Message) -> String {
fn main() { fn main() {
// 1. PromptTemplate::compile + render —— 直接构造模板 // 1. PromptTemplate::compile + render —— 直接构造模板
println!("=== PromptTemplate::compile + render ==="); println!("=== PromptTemplate::compile + render ===");
let tpl = PromptTemplate::compile( let tpl =
"今日 {{location}} 天气:{{condition}},温度 {{temperature}}", PromptTemplate::compile("今日 {{location}} 天气:{{condition}},温度 {{temperature}}")
)
.expect("编译失败"); .expect("编译失败");
let mut ctx = TemplateContext::new(); let mut ctx = TemplateContext::new();
ctx.insert("location", "北京"); ctx.insert("location", "北京");
@@ -58,7 +57,10 @@ fn main() {
.register("weather", "今日 {{location}}{{condition}}") .register("weather", "今日 {{location}}{{condition}}")
.expect("注册失败"); .expect("注册失败");
registry registry
.register("greet", "你好 {{name}}{{#if formal}} 见到您很荣幸。{{/if}}") .register(
"greet",
"你好 {{name}}{{#if formal}} 见到您很荣幸。{{/if}}",
)
.expect("注册失败"); .expect("注册失败");
let mut ctx = TemplateContext::new(); let mut ctx = TemplateContext::new();
+2 -4
View File
@@ -17,9 +17,9 @@ use std::sync::Arc;
use agcore::llm::cycle::{CycleConfig, LlmCycle}; use agcore::llm::cycle::{CycleConfig, LlmCycle};
use agcore::llm::mock::MockProvider; use agcore::llm::mock::MockProvider;
use agcore::llm::provider::LlmProvider; use agcore::llm::provider::LlmProvider;
use agcore::llm::types::Usage;
use agcore::llm::types::message::{ContentBlock, Message}; use agcore::llm::types::message::{ContentBlock, Message};
use agcore::llm::types::response_v2::{MessageResponse, StopReason, StreamEvent}; use agcore::llm::types::response_v2::{MessageResponse, StopReason, StreamEvent};
use agcore::llm::types::Usage;
use futures_util::StreamExt; use futures_util::StreamExt;
/// 构造预设的纯文本响应。 /// 构造预设的纯文本响应。
@@ -99,9 +99,7 @@ async fn main() {
// 上层 Agent 通过 `match` 或 `?` 处理 `AgentError::Llm(_)`。 // 上层 Agent 通过 `match` 或 `?` 处理 `AgentError::Llm(_)`。
println!("\n=== 阶段 2:错误路径(队列耗尽)==="); println!("\n=== 阶段 2:错误路径(队列耗尽)===");
let mut cycle = LlmCycle::new_with_arc(dyn_provider, CycleConfig::default()); let mut cycle = LlmCycle::new_with_arc(dyn_provider, CycleConfig::default());
let result = cycle let result = cycle.submit_stream("第二次提问".to_string(), vec![]).await;
.submit_stream("第二次提问".to_string(), vec![])
.await;
match result { match result {
Ok(_) => panic!("阶段 2 必须失败(队列耗尽)"), Ok(_) => panic!("阶段 2 必须失败(队列耗尽)"),
Err(e) => { Err(e) => {
+15 -3
View File
@@ -67,21 +67,33 @@ async fn main() {
assert!(step.status.is_pending()); assert!(step.status.is_pending());
step.status = StepStatus::Running; step.status = StepStatus::Running;
println!("Running: pending={}, terminal={}", step.status.is_pending(), step.status.is_terminal()); println!(
"Running: pending={}, terminal={}",
step.status.is_pending(),
step.status.is_terminal()
);
step.status = StepStatus::Completed(ChatResponse { step.status = StepStatus::Completed(ChatResponse {
message: OpenaiChatMessage::assistant_text("天气:晴,22°C"), message: OpenaiChatMessage::assistant_text("天气:晴,22°C"),
usage: Usage::from_input_output(5, 10), usage: Usage::from_input_output(5, 10),
stop_reason: Some(FinishReason::Stop), stop_reason: Some(FinishReason::Stop),
}); });
println!("Completed: pending={}, terminal={}", step.status.is_pending(), step.status.is_terminal()); println!(
"Completed: pending={}, terminal={}",
step.status.is_pending(),
step.status.is_terminal()
);
assert!(step.status.is_terminal()); assert!(step.status.is_terminal());
// 3. 失败路径 // 3. 失败路径
println!("\n=== Step 状态机:失败路径 ==="); println!("\n=== Step 状态机:失败路径 ===");
let mut fail_step = Step::new(0, "调用天气 API"); let mut fail_step = Step::new(0, "调用天气 API");
fail_step.status = StepStatus::Failed(AgentError::Other("API 不可用".into())); fail_step.status = StepStatus::Failed(AgentError::Other("API 不可用".into()));
println!("Failed: pending={}, terminal={}", fail_step.status.is_pending(), fail_step.status.is_terminal()); println!(
"Failed: pending={}, terminal={}",
fail_step.status.is_pending(),
fail_step.status.is_terminal()
);
assert!(fail_step.status.is_terminal()); assert!(fail_step.status.is_terminal());
// 4. 跳过路径 // 4. 跳过路径
+1 -1
View File
@@ -24,5 +24,5 @@ pub use error::AgentError;
pub use runtime::{AgentConfig, RuntimeBundle}; pub use runtime::{AgentConfig, RuntimeBundle};
pub use session::AgentSession; pub use session::AgentSession;
pub use session_memory::SessionMemory; pub use session_memory::SessionMemory;
pub use task::{Plan, PlanParser, Step, StepStatus, TaskAgent};
pub use task::JsonPlanParser; pub use task::JsonPlanParser;
pub use task::{Plan, PlanParser, Step, StepStatus, TaskAgent};
+8 -6
View File
@@ -92,15 +92,17 @@ impl AgentBuilder {
/// `AgentError::Config(...)`,提示调用 `.provider(...)` / `.tool_registry(...)` / /// `AgentError::Config(...)`,提示调用 `.provider(...)` / `.tool_registry(...)` /
/// `.hook_executor(...)` 补齐。不 panic。 /// `.hook_executor(...)` 补齐。不 panic。
pub fn build(self) -> Result<RuntimeBundle, AgentError> { pub fn build(self) -> Result<RuntimeBundle, AgentError> {
let provider = self let provider = self.provider.ok_or_else(|| {
.provider AgentError::Config("缺少 LLM provider,请先调用 .provider(...)".into())
.ok_or_else(|| AgentError::Config("缺少 LLM provider,请先调用 .provider(...)".into()))?; })?;
let tool_registry = self let tool_registry = self
.tool_registry .tool_registry
.ok_or_else(|| AgentError::Config("缺少 tool_registry,请先调用 .tool_registry(...)(即使是空 ToolRegistry 也需要传入)".into()))?; .ok_or_else(|| AgentError::Config("缺少 tool_registry,请先调用 .tool_registry(...)(即使是空 ToolRegistry 也需要传入)".into()))?;
let hook_executor = self let hook_executor = self.hook_executor.ok_or_else(|| {
.hook_executor AgentError::Config(
.ok_or_else(|| AgentError::Config("缺少 hook_executor,请先调用 .hook_executor(...)(空 HookExecutor 也可)".into()))?; "缺少 hook_executor,请先调用 .hook_executor(...)(空 HookExecutor 也可)".into(),
)
})?;
let config = self.config.unwrap_or_default(); let config = self.config.unwrap_or_default();
+1 -1
View File
@@ -16,8 +16,8 @@ use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use crate::llm::compact::CompactConfig; use crate::llm::compact::CompactConfig;
use crate::llm::provider::LlmProvider;
use crate::llm::hooks::HookExecutor; use crate::llm::hooks::HookExecutor;
use crate::llm::provider::LlmProvider;
use crate::memory::retriever::MemoryRetriever; use crate::memory::retriever::MemoryRetriever;
use crate::memory::store::MemoryStore; use crate::memory::store::MemoryStore;
use crate::tools::ToolRegistry; use crate::tools::ToolRegistry;
+8 -7
View File
@@ -126,8 +126,7 @@ impl AgentSession {
let hook_executor = Arc::clone(&self.bundle.hook_executor); let hook_executor = Arc::clone(&self.bundle.hook_executor);
// 1. 触发 OnTurnStart hook // 1. 触发 OnTurnStart hook
let start_ctx = let start_ctx = HookContext::new(HookEvent::OnTurnStart).with_turn_index(turn_index);
HookContext::new(HookEvent::OnTurnStart).with_turn_index(turn_index);
hook_executor hook_executor
.execute(HookEvent::OnTurnStart, &start_ctx) .execute(HookEvent::OnTurnStart, &start_ctx)
.await; .await;
@@ -137,7 +136,8 @@ impl AgentSession {
// submit_with_tools 内部从 registry 自行取 definitions,此处仅消费以触发 // submit_with_tools 内部从 registry 自行取 definitions,此处仅消费以触发
// 子 trait 覆盖(白名单/过滤)的副作用。 // 子 trait 覆盖(白名单/过滤)的副作用。
let _ = self.agent.tool_definitions(&self.bundle); let _ = self.agent.tool_definitions(&self.bundle);
let mut cycle = LlmCycle::new_with_arc(Arc::clone(&self.bundle.provider), CycleConfig::default()) let mut cycle =
LlmCycle::new_with_arc(Arc::clone(&self.bundle.provider), CycleConfig::default())
.with_messages(Vec::new()); .with_messages(Vec::new());
// Phase 2 切换 system_prompt 字段为 Message::SystemFIX-D)。 // Phase 2 切换 system_prompt 字段为 Message::SystemFIX-D)。
// 若 agent 自带 system prompt,预置到 messages 列表头部。 // 若 agent 自带 system prompt,预置到 messages 列表头部。
@@ -223,9 +223,7 @@ mod tests {
id: String::new(), id: String::new(),
model: String::new(), model: String::new(),
message: Message::Assistant { message: Message::Assistant {
content: vec![ContentBlock::Text { content: vec![ContentBlock::Text { text: text.into() }],
text: text.into(),
}],
}, },
usage: crate::llm::types::Usage::from_input_output(10, 5), usage: crate::llm::types::Usage::from_input_output(10, 5),
stop_reason: StopReason::Stop, stop_reason: StopReason::Stop,
@@ -280,7 +278,10 @@ mod tests {
assert!(session.get_session_data("k").await.unwrap().is_none()); assert!(session.get_session_data("k").await.unwrap().is_none());
session.set_session_data("k", "v").await.unwrap(); session.set_session_data("k", "v").await.unwrap();
assert_eq!(session.get_session_data("k").await.unwrap(), Some("v".into())); assert_eq!(
session.get_session_data("k").await.unwrap(),
Some("v".into())
);
// 覆盖写 // 覆盖写
session.set_session_data("k", "v2").await.unwrap(); session.set_session_data("k", "v2").await.unwrap();
assert_eq!( assert_eq!(
+2 -10
View File
@@ -78,11 +78,7 @@ impl SessionMemory {
prefix: Some(format!("{}:", self.namespace)), prefix: Some(format!("{}:", self.namespace)),
..Default::default() ..Default::default()
}; };
let items = self let items = self.store.list(&filter).await.map_err(AgentError::Memory)?;
.store
.list(&filter)
.await
.map_err(AgentError::Memory)?;
let mut lines = Vec::with_capacity(items.len() + 2); let mut lines = Vec::with_capacity(items.len() + 2);
lines.push("<session-context>".to_string()); lines.push("<session-context>".to_string());
@@ -113,11 +109,7 @@ impl SessionMemory {
prefix: Some(format!("{}:", self.namespace)), prefix: Some(format!("{}:", self.namespace)),
..Default::default() ..Default::default()
}; };
let items = self let items = self.store.list(&filter).await.map_err(AgentError::Memory)?;
.store
.list(&filter)
.await
.map_err(AgentError::Memory)?;
for item in items { for item in items {
self.store self.store
+2 -7
View File
@@ -130,9 +130,7 @@ impl PlanParser for JsonPlanParser {
.collect::<Result<Vec<_>, AgentError>>()?; .collect::<Result<Vec<_>, AgentError>>()?;
if steps.is_empty() { if steps.is_empty() {
return Err(AgentError::PlanParse( return Err(AgentError::PlanParse("Plan 至少需要一个步骤".into()));
"Plan 至少需要一个步骤".into(),
));
} }
Ok(Plan { Ok(Plan {
@@ -203,10 +201,7 @@ mod tests {
let plan = Plan { let plan = Plan {
id: "p1".into(), id: "p1".into(),
goal: "test goal".into(), goal: "test goal".into(),
steps: vec![ steps: vec![Step::new(0, "first"), Step::new(1, "second")],
Step::new(0, "first"),
Step::new(1, "second"),
],
}; };
assert_eq!(plan.steps.len(), 2); assert_eq!(plan.steps.len(), 2);
assert_eq!(plan.steps[0].index, 0); assert_eq!(plan.steps[0].index, 0);
+32 -16
View File
@@ -73,10 +73,7 @@ impl CompactState {
/// 粗略估计消息列表的 token 数(基于字符数,4 字符 ≈ 1 token)。 /// 粗略估计消息列表的 token 数(基于字符数,4 字符 ≈ 1 token)。
pub fn estimate_message_tokens(messages: &[Message]) -> u32 { pub fn estimate_message_tokens(messages: &[Message]) -> u32 {
messages messages.iter().map(estimate_single_message_tokens).sum()
.iter()
.map(estimate_single_message_tokens)
.sum()
} }
fn estimate_single_message_tokens(msg: &Message) -> u32 { fn estimate_single_message_tokens(msg: &Message) -> u32 {
@@ -99,9 +96,7 @@ fn estimate_block_tokens(block: &ContentBlock) -> u32 {
match block { match block {
ContentBlock::Text { text } => estimate_text_tokens(text), ContentBlock::Text { text } => estimate_text_tokens(text),
ContentBlock::Thinking { text, .. } => estimate_text_tokens(text), ContentBlock::Thinking { text, .. } => estimate_text_tokens(text),
ContentBlock::ToolUse { input, .. } => { ContentBlock::ToolUse { input, .. } => estimate_text_tokens(&input.to_string()),
estimate_text_tokens(&input.to_string())
}
ContentBlock::ToolResult { content, .. } => estimate_content_blocks_tokens(content), ContentBlock::ToolResult { content, .. } => estimate_content_blocks_tokens(content),
// ponytail: Image / Audio / File / Extension 在 IR 中固定估算。 // ponytail: Image / Audio / File / Extension 在 IR 中固定估算。
// 无文本的视觉/音频 block 用兜底估算,避免 token 计数膨胀。 // 无文本的视觉/音频 block 用兜底估算,避免 token 计数膨胀。
@@ -148,14 +143,25 @@ pub fn microcompact(messages: &mut [Message], keep_recent: usize) -> u32 {
// 第一遍:计算可释放 token(仅非错误 ToolResult // 第一遍:计算可释放 token(仅非错误 ToolResult
for msg in &messages[..prune_start] { for msg in &messages[..prune_start] {
if matches!(msg, Message::ToolResult { is_error: false, .. }) { if matches!(
msg,
Message::ToolResult {
is_error: false,
..
}
) {
freed_tokens += estimate_single_message_tokens(msg); freed_tokens += estimate_single_message_tokens(msg);
} }
} }
// 第二遍:替换内容(仅非错误 ToolResult // 第二遍:替换内容(仅非错误 ToolResult
for msg in &mut messages[..prune_start] { for msg in &mut messages[..prune_start] {
if let Message::ToolResult { content, is_error: false, .. } = msg { if let Message::ToolResult {
content,
is_error: false,
..
} = msg
{
*content = vec![ContentBlock::Text { *content = vec![ContentBlock::Text {
text: "[pruned]".to_string(), text: "[pruned]".to_string(),
}]; }];
@@ -177,13 +183,15 @@ mod tests {
fn estimate_message_tokens_handles_all_variants() { fn estimate_message_tokens_handles_all_variants() {
let messages = vec![ let messages = vec![
Message::System { Message::System {
content: vec![ContentBlock::Text { content: vec![ContentBlock::Text { text: "sys".into() }],
text: "sys".into(),
}],
}, },
Message::user_text("hi"), Message::user_text("hi"),
Message::assistant("ans"), Message::assistant("ans"),
Message::user_image("b64", "image/png", crate::llm::types::shared::ImageDetail::Auto), Message::user_image(
"b64",
"image/png",
crate::llm::types::shared::ImageDetail::Auto,
),
Message::tool_result("call_1", "tool res", false), Message::tool_result("call_1", "tool res", false),
]; ];
let tokens = estimate_message_tokens(&messages); let tokens = estimate_message_tokens(&messages);
@@ -205,7 +213,10 @@ mod tests {
assert!(freed > 0); assert!(freed > 0);
assert_eq!(messages.len(), before_len); // 只改内容,不删消息 assert_eq!(messages.len(), before_len); // 只改内容,不删消息
// 索引 1 是被压缩的 ToolResult // 索引 1 是被压缩的 ToolResult
if let Message::ToolResult { content, is_error, .. } = &messages[1] { if let Message::ToolResult {
content, is_error, ..
} = &messages[1]
{
assert_eq!(content.len(), 1); assert_eq!(content.len(), 1);
assert!(matches!(&content[0], ContentBlock::Text { text } if text == "[pruned]")); assert!(matches!(&content[0], ContentBlock::Text { text } if text == "[pruned]"));
assert!(!is_error); assert!(!is_error);
@@ -228,9 +239,14 @@ mod tests {
assert_eq!(freed, 0); // 错误 ToolResult 不计入 assert_eq!(freed, 0); // 错误 ToolResult 不计入
assert_eq!(messages.len(), before_len); assert_eq!(messages.len(), before_len);
// 错误信息保留完整 // 错误信息保留完整
if let Message::ToolResult { content, is_error, .. } = &messages[1] { if let Message::ToolResult {
content, is_error, ..
} = &messages[1]
{
assert!(is_error); assert!(is_error);
assert!(matches!(&content[0], ContentBlock::Text { text } if text.contains("backend down"))); assert!(
matches!(&content[0], ContentBlock::Text { text } if text.contains("backend down"))
);
} else { } else {
panic!("expected ToolResult at index 1"); panic!("expected ToolResult at index 1");
} }
+20 -21
View File
@@ -8,11 +8,9 @@
use serde_json::Value; use serde_json::Value;
use crate::llm::types::message::{ContentBlock, Message};
use crate::llm::types::openai_message::{
ContentField, OpenaiChatMessage, OpenaiContentPart,
};
use crate::llm::types::OpenaiToolCall; use crate::llm::types::OpenaiToolCall;
use crate::llm::types::message::{ContentBlock, Message};
use crate::llm::types::openai_message::{ContentField, OpenaiChatMessage, OpenaiContentPart};
/// `OpenaiChatMessage` → IR `Message`。 /// `OpenaiChatMessage` → IR `Message`。
/// ///
@@ -24,11 +22,10 @@ use crate::llm::types::OpenaiToolCall;
/// - `Function`(已废弃)→ `Message::ToolResult``name` 作为 `tool_call_id` 兜底) /// - `Function`(已废弃)→ `Message::ToolResult``name` 作为 `tool_call_id` 兜底)
pub fn from_openai(msg: &OpenaiChatMessage) -> Message { pub fn from_openai(msg: &OpenaiChatMessage) -> Message {
match msg { match msg {
OpenaiChatMessage::Developer { content, .. } | OpenaiChatMessage::System { content, .. } => { OpenaiChatMessage::Developer { content, .. }
Message::System { | OpenaiChatMessage::System { content, .. } => Message::System {
content: content_to_blocks(content), content: content_to_blocks(content),
} },
}
OpenaiChatMessage::User { content, .. } => Message::User { OpenaiChatMessage::User { content, .. } => Message::User {
content: content_to_blocks(content), content: content_to_blocks(content),
}, },
@@ -86,7 +83,11 @@ pub fn to_openai(msg: &Message) -> OpenaiChatMessage {
content: blocks_to_content(content), content: blocks_to_content(content),
name: None, name: None,
}, },
Message::UserImage { data, mime_type, detail } => { Message::UserImage {
data,
mime_type,
detail,
} => {
// ponytail: 构造为单 image part 的 User 消息(OpenAI 多模态格式)。 // ponytail: 构造为单 image part 的 User 消息(OpenAI 多模态格式)。
let mime = mime_type.clone(); let mime = mime_type.clone();
let is_url = data.starts_with("http://") || data.starts_with("https://"); let is_url = data.starts_with("http://") || data.starts_with("https://");
@@ -167,18 +168,17 @@ pub fn content_to_blocks(field: &ContentField) -> Vec<ContentBlock> {
ContentField::Array(parts) => parts ContentField::Array(parts) => parts
.iter() .iter()
.filter_map(|p| match p { .filter_map(|p| match p {
OpenaiContentPart::Text { text } => { OpenaiContentPart::Text { text } => Some(ContentBlock::Text { text: text.clone() }),
Some(ContentBlock::Text { text: text.clone() }) OpenaiContentPart::Refusal { refusal } => Some(ContentBlock::Text {
} text: refusal.clone(),
OpenaiContentPart::Refusal { refusal } => { }),
Some(ContentBlock::Text { text: refusal.clone() })
}
OpenaiContentPart::Image { image_url, .. } => { OpenaiContentPart::Image { image_url, .. } => {
// ponytail: 简化处理 —— URL 直接通过,data URI 拆出 // ponytail: 简化处理 —— URL 直接通过,data URI 拆出
// data:<mime>;base64,<b64> → ImageSource { data: b64, mime, is_url: false }。 // data:<mime>;base64,<b64> → ImageSource { data: b64, mime, is_url: false }。
let url = &image_url.url; let url = &image_url.url;
if let Some(rest) = url.strip_prefix("data:") if let Some(rest) = url.strip_prefix("data:")
&& let Some((mime, b64)) = rest.split_once(";base64,") { && let Some((mime, b64)) = rest.split_once(";base64,")
{
return Some(ContentBlock::Image { return Some(ContentBlock::Image {
source: crate::llm::types::message::ImageSource { source: crate::llm::types::message::ImageSource {
data: b64.to_string(), data: b64.to_string(),
@@ -263,7 +263,9 @@ mod tests {
match ir { match ir {
Message::System { content } => { Message::System { content } => {
assert_eq!(content.len(), 1); assert_eq!(content.len(), 1);
assert!(matches!(&content[0], ContentBlock::Text { text } if text == "you are helpful")); assert!(
matches!(&content[0], ContentBlock::Text { text } if text == "you are helpful")
);
} }
_ => panic!("expected System variant"), _ => panic!("expected System variant"),
} }
@@ -385,10 +387,7 @@ mod tests {
assert_eq!(parts.len(), 1); assert_eq!(parts.len(), 1);
match &parts[0] { match &parts[0] {
OpenaiContentPart::Image { image_url, .. } => { OpenaiContentPart::Image { image_url, .. } => {
assert_eq!( assert_eq!(image_url.url, "data:image/png;base64,BASE64DATA");
image_url.url,
"data:image/png;base64,BASE64DATA"
);
} }
_ => panic!("expected Image part"), _ => panic!("expected Image part"),
} }
+37 -34
View File
@@ -9,10 +9,10 @@ pub use usage::{CostTracker, Usage};
use std::pin::Pin; use std::pin::Pin;
use std::sync::Arc; use std::sync::Arc;
use futures_core::stream::Stream;
use async_stream::stream; use async_stream::stream;
use futures_core::stream::Stream;
use crate::llm::compact::{should_compact, microcompact, CompactConfig, CompactState}; use crate::llm::compact::{CompactConfig, CompactState, microcompact, should_compact};
use crate::llm::cycle::retry::should_retry; use crate::llm::cycle::retry::should_retry;
use crate::llm::error::LlmError; use crate::llm::error::LlmError;
use crate::llm::hooks::{HookContext, HookExecutor}; use crate::llm::hooks::{HookContext, HookExecutor};
@@ -113,8 +113,12 @@ impl LlmCycle {
note = "请改用 Message::system_text() + with_messages()" note = "请改用 Message::system_text() + with_messages()"
)] )]
pub fn with_system_prompt(mut self, prompt: String) -> Self { pub fn with_system_prompt(mut self, prompt: String) -> Self {
self.messages self.messages.insert(
.insert(0, Message::System { content: vec![ContentBlock::Text { text: prompt }] }); 0,
Message::System {
content: vec![ContentBlock::Text { text: prompt }],
},
);
self self
} }
@@ -188,8 +192,8 @@ impl LlmCycle {
}; };
if let Some(ref executor) = self.hook_executor { if let Some(ref executor) = self.hook_executor {
let ctx = HookContext::new(crate::llm::hooks::HookEvent::PreRequest) let ctx =
.with_request(&request); HookContext::new(crate::llm::hooks::HookEvent::PreRequest).with_request(&request);
let results = executor let results = executor
.execute(crate::llm::hooks::HookEvent::PreRequest, &ctx) .execute(crate::llm::hooks::HookEvent::PreRequest, &ctx)
.await; .await;
@@ -218,7 +222,8 @@ impl LlmCycle {
} }
Err(e) => { Err(e) => {
if let Some(ref executor) = self.hook_executor { if let Some(ref executor) = self.hook_executor {
let ctx = HookContext::new(crate::llm::hooks::HookEvent::OnError).with_error(&e); let ctx =
HookContext::new(crate::llm::hooks::HookEvent::OnError).with_error(&e);
executor executor
.execute(crate::llm::hooks::HookEvent::OnError, &ctx) .execute(crate::llm::hooks::HookEvent::OnError, &ctx)
.await; .await;
@@ -359,8 +364,8 @@ impl LlmCycle {
// PreRequest hook // PreRequest hook
if let Some(ref executor) = self.hook_executor { if let Some(ref executor) = self.hook_executor {
let ctx = HookContext::new(crate::llm::hooks::HookEvent::PreRequest) let ctx =
.with_request(&request); HookContext::new(crate::llm::hooks::HookEvent::PreRequest).with_request(&request);
let results = executor let results = executor
.execute(crate::llm::hooks::HookEvent::PreRequest, &ctx) .execute(crate::llm::hooks::HookEvent::PreRequest, &ctx)
.await; .await;
@@ -496,8 +501,8 @@ impl LlmCycle {
} }
Err(e) => { Err(e) => {
if let Some(ref executor) = self.hook_executor { if let Some(ref executor) = self.hook_executor {
let ctx = HookContext::new(crate::llm::hooks::HookEvent::OnError) let ctx =
.with_error(&e); HookContext::new(crate::llm::hooks::HookEvent::OnError).with_error(&e);
executor executor
.execute(crate::llm::hooks::HookEvent::OnError, &ctx) .execute(crate::llm::hooks::HookEvent::OnError, &ctx)
.await; .await;
@@ -593,11 +598,8 @@ impl LlmCycle {
// 真实 tool_call_id 而非 tool_name 充当 —— 这条 FIX-A 修复与 Phase 2 消息切换 // 真实 tool_call_id 而非 tool_name 充当 —— 这条 FIX-A 修复与 Phase 2 消息切换
// 同步生效。 // 同步生效。
// ponytail: Phase 2 直接存储 Message::ToolResultis_error 由 ToolInvocation.output 推断。 // ponytail: Phase 2 直接存储 Message::ToolResultis_error 由 ToolInvocation.output 推断。
self.messages.push(Message::tool_result( self.messages
result.tool_call_id, .push(Message::tool_result(result.tool_call_id, content, is_error));
content,
is_error,
));
} }
// 每轮工具执行后触发 compaction // 每轮工具执行后触发 compaction
@@ -641,9 +643,7 @@ fn has_tool_calls_in_response(response: &MessageResponse) -> bool {
/// ponytail: 当前 Phase 0 实现,`arguments_json_string` 内含 JSON 序列化的 input。 /// ponytail: 当前 Phase 0 实现,`arguments_json_string` 内含 JSON 序列化的 input。
/// 消费方在调用 `registry.invoke_all()` 时反序列化一次。该小段冗余序列化 /// 消费方在调用 `registry.invoke_all()` 时反序列化一次。该小段冗余序列化
/// 在 Phase 2 切换为 `Vec<Message>` 后可整体消除。 /// 在 Phase 2 切换为 `Vec<Message>` 后可整体消除。
fn extract_tool_calls_from_response( fn extract_tool_calls_from_response(response: &MessageResponse) -> Vec<(String, String, String)> {
response: &MessageResponse,
) -> Vec<(String, String, String)> {
let mut out = Vec::new(); let mut out = Vec::new();
if let Message::Assistant { content } = &response.message { if let Message::Assistant { content } = &response.message {
for block in content { for block in content {
@@ -665,7 +665,11 @@ fn truncate_tool_result(s: &str, max_bytes: usize) -> String {
while end > 0 && !s.is_char_boundary(end) { while end > 0 && !s.is_char_boundary(end) {
end -= 1; end -= 1;
} }
format!("{}\n\n[... truncated, original size: {} bytes ...]", &s[..end], s.len()) format!(
"{}\n\n[... truncated, original size: {} bytes ...]",
&s[..end],
s.len()
)
} }
#[cfg(test)] #[cfg(test)]
@@ -675,7 +679,7 @@ mod tests {
use crate::tools::{BaseTool, ToolRegistry}; use crate::tools::{BaseTool, ToolRegistry};
use async_trait::async_trait; use async_trait::async_trait;
use futures_core::Stream; use futures_core::Stream;
use serde_json::{json, Value}; use serde_json::{Value, json};
use std::pin::Pin; use std::pin::Pin;
/// 模拟 Provider —— 预定义响应序列,按调用顺序返回。 /// 模拟 Provider —— 预定义响应序列,按调用顺序返回。
@@ -729,9 +733,7 @@ mod tests {
id: String::new(), id: String::new(),
model: String::new(), model: String::new(),
message: Message::Assistant { message: Message::Assistant {
content: vec![ContentBlock::Text { content: vec![ContentBlock::Text { text: text.into() }],
text: text.into(),
}],
}, },
usage: empty_usage(), usage: empty_usage(),
stop_reason: StopReason::Stop, stop_reason: StopReason::Stop,
@@ -755,7 +757,9 @@ mod tests {
MessageResponse { MessageResponse {
id: String::new(), id: String::new(),
model: String::new(), model: String::new(),
message: Message::Assistant { content: tool_blocks }, message: Message::Assistant {
content: tool_blocks,
},
usage: empty_usage(), usage: empty_usage(),
stop_reason: StopReason::ToolUse, stop_reason: StopReason::ToolUse,
extra: std::collections::HashMap::new(), extra: std::collections::HashMap::new(),
@@ -810,16 +814,13 @@ mod tests {
let messages = cycle.messages(); let messages = cycle.messages();
assert_eq!(messages.len(), 4); assert_eq!(messages.len(), 4);
assert!(matches!(messages[0], Message::User { .. })); assert!(matches!(messages[0], Message::User { .. }));
assert!(matches!( assert!(matches!(messages[1], Message::Assistant { content: _ }));
messages[1],
Message::Assistant {
content: _,
}
));
if let Message::Assistant { content } = &messages[1] { if let Message::Assistant { content } = &messages[1] {
assert!(content assert!(
content
.iter() .iter()
.any(|b| matches!(b, ContentBlock::ToolUse { .. }))); .any(|b| matches!(b, ContentBlock::ToolUse { .. }))
);
} }
assert!(matches!( assert!(matches!(
messages[2], messages[2],
@@ -875,7 +876,9 @@ mod tests {
registry.register(std::sync::Arc::new(AddTool)).unwrap(); registry.register(std::sync::Arc::new(AddTool)).unwrap();
let result = cycle.submit_with_tools("test".to_string(), &registry).await; let result = cycle.submit_with_tools("test".to_string(), &registry).await;
assert!(matches!(result, Err(LlmError::Other(msg)) if msg.contains("达到最大工具循环轮次"))); assert!(
matches!(result, Err(LlmError::Other(msg)) if msg.contains("达到最大工具循环轮次"))
);
} }
#[tokio::test] #[tokio::test]
+9 -3
View File
@@ -10,7 +10,9 @@ use std::time::Duration;
#[derive(thiserror::Error, Debug)] #[derive(thiserror::Error, Debug)]
pub enum LlmError { pub enum LlmError {
/// API 认证失败(API key 无效、过期或权限不足)。 /// API 认证失败(API key 无效、过期或权限不足)。
#[error("LLM 认证失败: {0}。请检查环境变量中的 API key(如 OPENAI_API_KEY / ANTHROPIC_API_KEY)是否正确")] #[error(
"LLM 认证失败: {0}。请检查环境变量中的 API key(如 OPENAI_API_KEY / ANTHROPIC_API_KEY)是否正确"
)]
Authentication(String), Authentication(String),
/// 请求被限流,可选地附带重试等待时间。可重试。 /// 请求被限流,可选地附带重试等待时间。可重试。
@@ -18,7 +20,9 @@ pub enum LlmError {
RateLimit { retry_after: Option<Duration> }, RateLimit { retry_after: Option<Duration> },
/// HTTP 请求失败(网络错误或非 2xx 状态码),包含状态码与响应体。 /// HTTP 请求失败(网络错误或非 2xx 状态码),包含状态码与响应体。
#[error("LLM 请求失败(HTTP {status}: {body}。请检查 Provider 端点地址(base_url)和网络连通性")] #[error(
"LLM 请求失败(HTTP {status}: {body}。请检查 Provider 端点地址(base_url)和网络连通性"
)]
Request { status: u16, body: String }, Request { status: u16, body: String },
/// 请求超时。可重试。 /// 请求超时。可重试。
@@ -30,7 +34,9 @@ pub enum LlmError {
Stream(String), Stream(String),
/// 上下文长度超出模型窗口限制。 /// 上下文长度超出模型窗口限制。
#[error("LLM 上下文超限:当前 {actual} tokens > 模型上限 {limit} tokens。请减少消息历史、缩短 prompt,或启用 auto-compactionllm::compact")] #[error(
"LLM 上下文超限:当前 {actual} tokens > 模型上限 {limit} tokens。请减少消息历史、缩短 prompt,或启用 auto-compactionllm::compact"
)]
ContextLength { actual: u32, limit: u32 }, ContextLength { actual: u32, limit: u32 },
/// 其他未分类的 LLM 调用失败。 /// 其他未分类的 LLM 调用失败。
+1 -3
View File
@@ -130,9 +130,7 @@ impl Default for HookExecutor {
impl HookExecutor { impl HookExecutor {
/// 创建一个空的执行器。 /// 创建一个空的执行器。
pub fn new() -> Self { pub fn new() -> Self {
Self { Self { hooks: Vec::new() }
hooks: Vec::new(),
}
} }
/// 注册一个钩子到指定事件点。 /// 注册一个钩子到指定事件点。
+10 -8
View File
@@ -97,8 +97,7 @@ impl LlmProvider for MockProvider {
async fn chat_stream( async fn chat_stream(
&self, &self,
_request: MessageRequest, _request: MessageRequest,
) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>, LlmError> ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>, LlmError> {
{
let response = self.pop()?; let response = self.pop()?;
// 提前 clone 出在 stream 闭包中需要的字段;最后 yield 时 move response。 // 提前 clone 出在 stream 闭包中需要的字段;最后 yield 时 move response。
let id = response.id.clone(); let id = response.id.clone();
@@ -206,10 +205,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn chat_returns_queued_response() { async fn chat_returns_queued_response() {
let provider = MockProvider::new(vec![text_response("hello")]); let provider = MockProvider::new(vec![text_response("hello")]);
let resp = provider let resp = provider.chat(MessageRequest::default()).await.unwrap();
.chat(MessageRequest::default())
.await
.unwrap();
assert_eq!(resp.text(), "hello"); assert_eq!(resp.text(), "hello");
assert_eq!(provider.remaining(), 0); assert_eq!(provider.remaining(), 0);
} }
@@ -231,7 +227,10 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn chat_stream_emits_text_delta_sequence() { async fn chat_stream_emits_text_delta_sequence() {
let provider = MockProvider::new(vec![text_response("hi")]); let provider = MockProvider::new(vec![text_response("hi")]);
let mut stream = provider.chat_stream(MessageRequest::default()).await.unwrap(); let mut stream = provider
.chat_stream(MessageRequest::default())
.await
.unwrap();
let mut seen_start = false; let mut seen_start = false;
let mut seen_block_start = false; let mut seen_block_start = false;
@@ -283,7 +282,10 @@ mod tests {
extra: Default::default(), extra: Default::default(),
}; };
let provider = MockProvider::new(vec![response]); let provider = MockProvider::new(vec![response]);
let mut stream = provider.chat_stream(MessageRequest::default()).await.unwrap(); let mut stream = provider
.chat_stream(MessageRequest::default())
.await
.unwrap();
let mut saw_tool_args = false; let mut saw_tool_args = false;
let mut saw_tool_end = false; let mut saw_tool_end = false;
+2 -4
View File
@@ -3,7 +3,7 @@
use std::collections::HashMap; use std::collections::HashMap;
use crate::llm::error::LlmError; use crate::llm::error::LlmError;
use crate::llm::provider::{create_provider, LlmProvider, ProviderConfig, ProviderType}; use crate::llm::provider::{LlmProvider, ProviderConfig, ProviderType, create_provider};
/// Provider 注册表 —— 管理多个 LLM Provider 实例。 /// Provider 注册表 —— 管理多个 LLM Provider 实例。
/// ///
@@ -61,8 +61,6 @@ impl ProviderRegistry {
/// 获取默认 Provider。 /// 获取默认 Provider。
pub fn get_default(&self) -> Option<&dyn LlmProvider> { pub fn get_default(&self) -> Option<&dyn LlmProvider> {
self.default_name self.default_name.as_ref().and_then(|name| self.get(name))
.as_ref()
.and_then(|name| self.get(name))
} }
} }
+7 -8
View File
@@ -14,8 +14,8 @@ use std::pin::Pin;
use std::task::{Context, Poll}; use std::task::{Context, Poll};
use futures_core::stream::Stream; use futures_core::stream::Stream;
use futures_util::future::poll_fn;
use futures_util::FutureExt; use futures_util::FutureExt;
use futures_util::future::poll_fn;
use serde_json::Value; use serde_json::Value;
use crate::llm::error::LlmError; use crate::llm::error::LlmError;
@@ -95,9 +95,7 @@ impl Stream for ChunkToLegacyEventStream {
} }
if let Some(usage) = &chunk.usage { if let Some(usage) = &chunk.usage {
return Poll::Ready(Some(LegacyStreamEvent::CostUpdate { return Poll::Ready(Some(LegacyStreamEvent::CostUpdate { usage: *usage }));
usage: *usage,
}));
} }
Poll::Ready(None) Poll::Ready(None)
@@ -143,9 +141,7 @@ fn empty_message_response() -> MessageResponse {
MessageResponse { MessageResponse {
id: String::new(), id: String::new(),
model: String::new(), model: String::new(),
message: Message::Assistant { message: Message::Assistant { content: vec![] },
content: vec![],
},
usage: Usage::default(), usage: Usage::default(),
stop_reason: StopReason::Stop, stop_reason: StopReason::Stop,
extra: HashMap::new(), extra: HashMap::new(),
@@ -172,7 +168,10 @@ fn map_legacy_to_ir(legacy: LegacyStreamEvent) -> StreamEvent {
LegacyStreamEvent::AssistantTextDelta { text } => StreamEvent::TextDelta { text }, LegacyStreamEvent::AssistantTextDelta { text } => StreamEvent::TextDelta { text },
LegacyStreamEvent::ToolExecutionStarted { input, .. } => { LegacyStreamEvent::ToolExecutionStarted { input, .. } => {
let arguments = serde_json::to_string(&input).unwrap_or_default(); let arguments = serde_json::to_string(&input).unwrap_or_default();
StreamEvent::ToolCallArgumentsDelta { index: 0, arguments } StreamEvent::ToolCallArgumentsDelta {
index: 0,
arguments,
}
} }
LegacyStreamEvent::ToolExecutionCompleted { .. } => { LegacyStreamEvent::ToolExecutionCompleted { .. } => {
// 旧 ToolExecutionCompleted 不在 IR 流协议中——工具执行是消费方职责。 // 旧 ToolExecutionCompleted 不在 IR 流协议中——工具执行是消费方职责。
+6 -19
View File
@@ -22,13 +22,9 @@ use crate::llm::types::shared::ImageDetail;
#[serde(rename_all = "snake_case")] #[serde(rename_all = "snake_case")]
pub enum Message { pub enum Message {
/// 系统提示(User & Assistant 之外的引导指令)。 /// 系统提示(User & Assistant 之外的引导指令)。
System { System { content: Vec<ContentBlock> },
content: Vec<ContentBlock>,
},
/// 用户输入。 /// 用户输入。
User { User { content: Vec<ContentBlock> },
content: Vec<ContentBlock>,
},
/// 用户的图片输入(快捷构造,免去构造 ContentBlock 的 boilerplate)。 /// 用户的图片输入(快捷构造,免去构造 ContentBlock 的 boilerplate)。
UserImage { UserImage {
data: String, data: String,
@@ -36,9 +32,7 @@ pub enum Message {
detail: ImageDetail, detail: ImageDetail,
}, },
/// Assistant 回复内容块(可能包含 text、thinking、tool_use 等多种 block 的混合)。 /// Assistant 回复内容块(可能包含 text、thinking、tool_use 等多种 block 的混合)。
Assistant { Assistant { content: Vec<ContentBlock> },
content: Vec<ContentBlock>,
},
/// 工具调用结果。 /// 工具调用结果。
ToolResult { ToolResult {
tool_call_id: String, tool_call_id: String,
@@ -130,10 +124,7 @@ pub enum ContentBlock {
signature: Option<String>, signature: Option<String>,
}, },
/// 逃生舱:Provider 特定 block 透传(OpenAI Response 内置工具等)。 /// 逃生舱:Provider 特定 block 透传(OpenAI Response 内置工具等)。
Extension { Extension { kind: String, data: Value },
kind: String,
data: Value,
},
} }
/// 内容块类型标签 —— 用于 `StreamEvent::ContentBlockStart.block_type`。 /// 内容块类型标签 —— 用于 `StreamEvent::ContentBlockStart.block_type`。
@@ -349,9 +340,7 @@ mod tests {
fn message_roundtrip_each_variant() { fn message_roundtrip_each_variant() {
let msgs = vec![ let msgs = vec![
Message::System { Message::System {
content: vec![ContentBlock::Text { content: vec![ContentBlock::Text { text: "sys".into() }],
text: "sys".into(),
}],
}, },
Message::User { Message::User {
content: vec![ContentBlock::Text { content: vec![ContentBlock::Text {
@@ -376,9 +365,7 @@ mod tests {
}, },
Message::ToolResult { Message::ToolResult {
tool_call_id: "call_1".into(), tool_call_id: "call_1".into(),
content: vec![ContentBlock::Text { content: vec![ContentBlock::Text { text: "ok".into() }],
text: "ok".into(),
}],
is_error: true, is_error: true,
}, },
]; ];
+4 -1
View File
@@ -98,5 +98,8 @@ impl From<ChatResponse> for OpenaiChatChunk {
} }
/// 工具定义别名(无新类型冲突,保留)。 /// 工具定义别名(无新类型冲突,保留)。
#[deprecated(since = "0.1.0", note = "ToolDefinition 仍直接对应 OpenAI wire-format;未来 v0.2 引入 IR 工具类型后会再次更新")] #[deprecated(
since = "0.1.0",
note = "ToolDefinition 仍直接对应 OpenAI wire-format;未来 v0.2 引入 IR 工具类型后会再次更新"
)]
pub type ToolDefinition = OpenaiToolDefinition; pub type ToolDefinition = OpenaiToolDefinition;
+7 -5
View File
@@ -11,18 +11,20 @@ pub struct StreamOptions {
pub include_obfuscation: Option<bool>, pub include_obfuscation: Option<bool>,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone, Default)]
#[derive(Default)]
pub enum ToolChoice { pub enum ToolChoice {
#[default] #[default]
None, None,
Auto, Auto,
Required, Required,
Named { name: String }, Named {
AllowedTools { tool_names: Vec<String> }, name: String,
},
AllowedTools {
tool_names: Vec<String>,
},
} }
impl Serialize for ToolChoice { impl Serialize for ToolChoice {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where where
+3 -11
View File
@@ -127,14 +127,9 @@ mod tests {
#[test] #[test]
fn extra_set_and_get_roundtrip() { fn extra_set_and_get_roundtrip() {
let mut req = MessageRequest::default(); let mut req = MessageRequest::default();
req.set_extra( req.set_extra("previous_response_id", "resp_abc123");
"previous_response_id",
"resp_abc123",
);
let v: Option<String> = req let v: Option<String> = req.get_extra("previous_response_id").expect("get_extra ok");
.get_extra("previous_response_id")
.expect("get_extra ok");
assert_eq!(v.as_deref(), Some("resp_abc123")); assert_eq!(v.as_deref(), Some("resp_abc123"));
let missing: Option<String> = req.get_extra("missing").expect("missing ok"); let missing: Option<String> = req.get_extra("missing").expect("missing ok");
@@ -174,10 +169,7 @@ mod tests {
} }
let opts: Options = req.get_extra_as().expect("get_extra_as ok"); let opts: Options = req.get_extra_as().expect("get_extra_as ok");
assert_eq!( assert_eq!(opts.web_search_options.search_context_size, "high");
opts.web_search_options.search_context_size,
"high"
);
assert_eq!(opts.user.as_deref(), Some("u_123")); assert_eq!(opts.user.as_deref(), Some("u_123"));
} }
+2 -6
View File
@@ -2,8 +2,8 @@ use crate::llm::types::openai_message::OpenaiChatMessage;
use crate::llm::types::shared::{FinishReason, ServiceTier}; use crate::llm::types::shared::{FinishReason, ServiceTier};
use crate::llm::types::tool::OpenaiToolCall; use crate::llm::types::tool::OpenaiToolCall;
use crate::llm::types::usage::Usage; use crate::llm::types::usage::Usage;
use serde::{Deserialize, Serialize};
use crate::llm::types::{ContentField, OpenaiContentPart}; use crate::llm::types::{ContentField, OpenaiContentPart};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenLogprob { pub struct TokenLogprob {
@@ -135,11 +135,7 @@ impl From<OpenaiChatMessage> for Delta {
text.push_str(&t); text.push_str(&t);
} }
} }
if text.is_empty() { if text.is_empty() { None } else { Some(text) }
None
} else {
Some(text)
}
} }
}, },
refusal: None, refusal: None,
+2 -2
View File
@@ -16,7 +16,7 @@ pub use store::{InMemoryStore, MemoryStore};
// 低频类型(配置/高级使用) // 低频类型(配置/高级使用)
pub use conversation::MemoryStrategy; pub use conversation::MemoryStrategy;
pub use knowledge::{PageIndexEntry, KNOWLEDGE_PREFIX}; pub use knowledge::{KNOWLEDGE_PREFIX, PageIndexEntry};
pub use retriever::{RetrieverConfig, RetrievalResult, ScoredItem}; pub use retriever::{RetrievalResult, RetrieverConfig, ScoredItem};
pub use store::{EvictionConfig, EvictionPolicy}; pub use store::{EvictionConfig, EvictionPolicy};
pub use types::{KnowledgePage, MemoryFilter, MemoryItem}; pub use types::{KnowledgePage, MemoryFilter, MemoryItem};
+13 -7
View File
@@ -160,7 +160,12 @@ impl ConversationMemory {
} }
fn make_message_id(&self, index: usize, now: &OffsetDateTime) -> String { fn make_message_id(&self, index: usize, now: &OffsetDateTime) -> String {
format!("{}{:010}_{}", self.session_prefix(), index, now.unix_timestamp_nanos()) format!(
"{}{:010}_{}",
self.session_prefix(),
index,
now.unix_timestamp_nanos()
)
} }
async fn maybe_evict_and_compact(&mut self) { async fn maybe_evict_and_compact(&mut self) {
@@ -175,7 +180,8 @@ impl ConversationMemory {
} }
if let Some(ref compact_config) = self.config.compact_config if let Some(ref compact_config) = self.config.compact_config
&& should_compact(&self.messages, compact_config, &self.compact_state) { && should_compact(&self.messages, compact_config, &self.compact_state)
{
let keep_recent = compact_config.keep_recent; let keep_recent = compact_config.keep_recent;
let freed = microcompact(&mut self.messages, keep_recent); let freed = microcompact(&mut self.messages, keep_recent);
if freed > 0 { if freed > 0 {
@@ -196,7 +202,8 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn add_and_get_history() { async fn add_and_get_history() {
let store = Arc::new(InMemoryStore::new()) as Arc<dyn MemoryStore>; let store = Arc::new(InMemoryStore::new()) as Arc<dyn MemoryStore>;
let mut conv = ConversationMemory::new(store, "session1", ConversationMemoryConfig::default()); let mut conv =
ConversationMemory::new(store, "session1", ConversationMemoryConfig::default());
conv.add_message(Message::user_text("hello")).await.unwrap(); conv.add_message(Message::user_text("hello")).await.unwrap();
conv.add_message(Message::user_text("world")).await.unwrap(); conv.add_message(Message::user_text("world")).await.unwrap();
assert_eq!(conv.len(), 2); assert_eq!(conv.len(), 2);
@@ -211,9 +218,7 @@ mod tests {
conv.add_message(Message::tool_result("call_1", "ok", false)) conv.add_message(Message::tool_result("call_1", "ok", false))
.await .await
.unwrap(); .unwrap();
conv.add_message(Message::assistant("done")) conv.add_message(Message::assistant("done")).await.unwrap();
.await
.unwrap();
let original = conv.get_history().to_vec(); let original = conv.get_history().to_vec();
assert_eq!(original.len(), 2); assert_eq!(original.len(), 2);
@@ -263,7 +268,8 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn clear_empties_messages() { async fn clear_empties_messages() {
let store = Arc::new(InMemoryStore::new()) as Arc<dyn MemoryStore>; let store = Arc::new(InMemoryStore::new()) as Arc<dyn MemoryStore>;
let mut conv = ConversationMemory::new(store.clone(), "s1", ConversationMemoryConfig::default()); let mut conv =
ConversationMemory::new(store.clone(), "s1", ConversationMemoryConfig::default());
conv.add_message(Message::user_text("hello")).await.unwrap(); conv.add_message(Message::user_text("hello")).await.unwrap();
assert!(!conv.is_empty()); assert!(!conv.is_empty());
conv.clear().await.unwrap(); conv.clear().await.unwrap();
+6 -3
View File
@@ -57,8 +57,8 @@ impl KnowledgeStore {
} }
let now = OffsetDateTime::now_utc(); let now = OffsetDateTime::now_utc();
let id = format!("{KNOWLEDGE_PREFIX}{}", page.id); let id = format!("{KNOWLEDGE_PREFIX}{}", page.id);
let content = serde_json::to_string(&page) let content =
.map_err(|e| MemoryError::Serialization(e.to_string()))?; serde_json::to_string(&page).map_err(|e| MemoryError::Serialization(e.to_string()))?;
let item = MemoryItem { let item = MemoryItem {
id, id,
content, content,
@@ -128,7 +128,10 @@ impl KnowledgeStore {
.filter(|entry| { .filter(|entry| {
entry.title.to_lowercase().contains(&needle) entry.title.to_lowercase().contains(&needle)
|| entry.summary.to_lowercase().contains(&needle) || entry.summary.to_lowercase().contains(&needle)
|| entry.tags.iter().any(|t| t.to_lowercase().contains(&needle)) || entry
.tags
.iter()
.any(|t| t.to_lowercase().contains(&needle))
}) })
.map(|entry| entry.id.clone()) .map(|entry| entry.id.clone())
.collect() .collect()
+15 -8
View File
@@ -97,7 +97,11 @@ impl MemoryRetriever {
// 4. 过滤 → 排序 → 截取 // 4. 过滤 → 排序 → 截取
items.retain(|i| i.score >= self.config.min_score); items.retain(|i| i.score >= self.config.min_score);
items.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); items.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
items.truncate(self.config.max_results); items.truncate(self.config.max_results);
Ok(RetrievalResult { Ok(RetrievalResult {
@@ -159,12 +163,12 @@ fn char_bigrams(s: &str) -> Vec<String> {
fn default_stop_words() -> HashSet<String> { fn default_stop_words() -> HashSet<String> {
[ [
"the", "a", "an", "is", "are", "was", "were", "be", "been", "being", "have", "has", "the", "a", "an", "is", "are", "was", "were", "be", "been", "being", "have", "has", "had",
"had", "do", "does", "did", "will", "would", "should", "could", "may", "might", "shall", "do", "does", "did", "will", "would", "should", "could", "may", "might", "shall", "can",
"can", "this", "that", "these", "those", "it", "its", "they", "them", "their", "what", "this", "that", "these", "those", "it", "its", "they", "them", "their", "what", "which",
"which", "who", "whom", "how", "when", "where", "and", "or", "but", "not", "no", "nor", "who", "whom", "how", "when", "where", "and", "or", "but", "not", "no", "nor", "so", "if",
"so", "if", "then", "else", "with", "without", "for", "to", "from", "in", "on", "at", "then", "else", "with", "without", "for", "to", "from", "in", "on", "at", "by", "of", "as",
"by", "of", "as", "into", "through", "during", "before", "after", "above", "below", "into", "through", "during", "before", "after", "above", "below",
] ]
.iter() .iter()
.map(|s| s.to_string()) .map(|s| s.to_string())
@@ -236,7 +240,10 @@ mod tests {
min_score: 0.99, min_score: 0.99,
}; };
let retriever = MemoryRetriever::new(ks, config); let retriever = MemoryRetriever::new(ks, config);
let result = retriever.retrieve("totally unrelated content").await.unwrap(); let result = retriever
.retrieve("totally unrelated content")
.await
.unwrap();
assert!(result.items.is_empty()); assert!(result.items.is_empty());
} }
+2 -2
View File
@@ -1,7 +1,7 @@
pub mod composer;
pub mod error; pub mod error;
pub mod template; pub mod template;
pub mod composer;
pub use composer::{PromptComposer, validate_messages};
pub use error::PromptError; pub use error::PromptError;
pub use template::{PromptTemplate, PromptTemplateRegistry, TemplateContext, TemplateValue}; pub use template::{PromptTemplate, PromptTemplateRegistry, TemplateContext, TemplateValue};
pub use composer::{validate_messages, PromptComposer};
+14 -12
View File
@@ -48,7 +48,11 @@ impl PromptComposer {
/// 添加一条 Tool 消息(工具执行结果回传)。 /// 添加一条 Tool 消息(工具执行结果回传)。
pub fn tool(mut self, tool_call_id: impl Into<String>, content: impl Into<String>) -> Self { pub fn tool(mut self, tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
self.push_message(Message::tool_result(tool_call_id.into(), content.into(), false)); self.push_message(Message::tool_result(
tool_call_id.into(),
content.into(),
false,
));
self self
} }
@@ -133,11 +137,7 @@ impl PromptComposer {
} }
/// 添加一条含指定 ContentBlock 的 Tool 消息。 /// 添加一条含指定 ContentBlock 的 Tool 消息。
pub fn tool_content( pub fn tool_content(mut self, tool_call_id: impl Into<String>, block: ContentBlock) -> Self {
mut self,
tool_call_id: impl Into<String>,
block: ContentBlock,
) -> Self {
self.push_message(Message::ToolResult { self.push_message(Message::ToolResult {
tool_call_id: tool_call_id.into(), tool_call_id: tool_call_id.into(),
content: vec![block], content: vec![block],
@@ -187,9 +187,7 @@ impl PromptComposer {
/// 验证消息序列是否符合 LLM API 要求(Tool 消息必须紧跟含 tool_calls 的 Assistant)。 /// 验证消息序列是否符合 LLM API 要求(Tool 消息必须紧跟含 tool_calls 的 Assistant)。
pub fn validate_messages(messages: &[Message]) -> Result<(), PromptError> { pub fn validate_messages(messages: &[Message]) -> Result<(), PromptError> {
if messages.is_empty() { if messages.is_empty() {
return Err(PromptError::InvalidSequence( return Err(PromptError::InvalidSequence("消息列表不能为空".to_string()));
"消息列表不能为空".to_string(),
));
} }
let mut last_tool_call_ids: Vec<String> = Vec::new(); let mut last_tool_call_ids: Vec<String> = Vec::new();
@@ -297,7 +295,8 @@ mod tests {
#[test] #[test]
fn test_template_if() { fn test_template_if() {
let tpl = PromptTemplate::compile("Hello {{#if name}}{{name}}{{else}}Guest{{/if}}").unwrap(); let tpl =
PromptTemplate::compile("Hello {{#if name}}{{name}}{{else}}Guest{{/if}}").unwrap();
let mut ctx = TemplateContext::new(); let mut ctx = TemplateContext::new();
ctx.insert("name", "Bob"); ctx.insert("name", "Bob");
@@ -312,11 +311,14 @@ mod tests {
fn test_template_each() { fn test_template_each() {
let tpl = PromptTemplate::compile("Items: {{#each items}}{{item}}, {{/each}}").unwrap(); let tpl = PromptTemplate::compile("Items: {{#each items}}{{item}}, {{/each}}").unwrap();
let mut ctx = TemplateContext::new(); let mut ctx = TemplateContext::new();
ctx.insert("items", TemplateValue::Array(vec![ ctx.insert(
"items",
TemplateValue::Array(vec![
TemplateValue::String("a".to_string()), TemplateValue::String("a".to_string()),
TemplateValue::String("b".to_string()), TemplateValue::String("b".to_string()),
TemplateValue::String("c".to_string()), TemplateValue::String("c".to_string()),
])); ]),
);
let result = tpl.render(&ctx).unwrap(); let result = tpl.render(&ctx).unwrap();
assert_eq!(result, "Items: a, b, c, "); assert_eq!(result, "Items: a, b, c, ");
+6 -2
View File
@@ -8,7 +8,9 @@ pub enum PromptError {
#[error("渲染错误: 变量 '{0}' 未找到。请在 TemplateContext 中插入该变量")] #[error("渲染错误: 变量 '{0}' 未找到。请在 TemplateContext 中插入该变量")]
VariableNotFound(String), VariableNotFound(String),
#[error("渲染错误: 引用的子模板 '{0}' 未注册。请先用 PromptTemplateRegistry::register 注册该子模板")] #[error(
"渲染错误: 引用的子模板 '{0}' 未注册。请先用 PromptTemplateRegistry::register 注册该子模板"
)]
PartialNotFound(String), PartialNotFound(String),
#[error("渲染错误: '{0}' 不是数组,无法遍历。请确认传入的是数组或先判空")] #[error("渲染错误: '{0}' 不是数组,无法遍历。请确认传入的是数组或先判空")]
@@ -20,7 +22,9 @@ pub enum PromptError {
#[error("渲染错误: {0}")] #[error("渲染错误: {0}")]
Render(String), Render(String),
#[error("消息序列校验失败: {0}。请检查消息角色顺序(例如 tool 必须在 assistant tool_call 之后)")] #[error(
"消息序列校验失败: {0}。请检查消息角色顺序(例如 tool 必须在 assistant tool_call 之后)"
)]
InvalidSequence(String), InvalidSequence(String),
#[error("文件读取错误: {0}。请检查模板文件路径与权限")] #[error("文件读取错误: {0}。请检查模板文件路径与权限")]
+16 -29
View File
@@ -1,6 +1,6 @@
use serde_json::Value;
use std::collections::HashMap; use std::collections::HashMap;
use std::fmt; use std::fmt;
use serde_json::Value;
use crate::prompt::error::PromptError; use crate::prompt::error::PromptError;
@@ -140,7 +140,9 @@ fn json_to_template_value(v: &Value) -> Result<TemplateValue, PromptError> {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
enum Fragment { enum Fragment {
Literal(String), Literal(String),
Variable { name: String }, Variable {
name: String,
},
If { If {
condition: String, condition: String,
body: Vec<Fragment>, body: Vec<Fragment>,
@@ -223,8 +225,7 @@ fn compile_fragments(template: &str) -> Result<Vec<Fragment>, PromptError> {
let tag = tag_content.trim(); let tag = tag_content.trim();
if let Some(rest) = tag.strip_prefix("#if ") { if let Some(rest) = tag.strip_prefix("#if ") {
let (body, else_body, new_i) = let (body, else_body, new_i) = parse_block(template, i, "if")?;
parse_block(template, i, "if")?;
let condition = rest.trim().to_string(); let condition = rest.trim().to_string();
fragments.push(Fragment::If { fragments.push(Fragment::If {
condition, condition,
@@ -331,10 +332,7 @@ fn parse_block(
Err(PromptError::Parse(format!("未闭合的 {{#{}}}", kind))) Err(PromptError::Parse(format!("未闭合的 {{#{}}}", kind)))
} }
fn parse_each_block( fn parse_each_block(template: &str, start: usize) -> Result<(Vec<Fragment>, usize), PromptError> {
template: &str,
start: usize,
) -> Result<(Vec<Fragment>, usize), PromptError> {
let bytes = template.as_bytes(); let bytes = template.as_bytes();
let len = bytes.len(); let len = bytes.len();
let mut depth = 1u32; let mut depth = 1u32;
@@ -368,9 +366,7 @@ fn parse_each_block(
} }
} }
Err(PromptError::Parse( Err(PromptError::Parse("未闭合的 {{#each}} 块".to_string()))
"未闭合的 {{#each}} 块".to_string(),
))
} }
fn parse_raw_block(template: &str, start: usize) -> Result<(String, usize), PromptError> { fn parse_raw_block(template: &str, start: usize) -> Result<(String, usize), PromptError> {
@@ -395,9 +391,7 @@ fn parse_raw_block(template: &str, start: usize) -> Result<(String, usize), Prom
} }
} }
Err(PromptError::Parse( Err(PromptError::Parse("未闭合的 {{#raw}} 块".to_string()))
"未闭合的 {{#raw}} 块".to_string(),
))
} }
// ===== Renderer ===== // ===== Renderer =====
@@ -418,33 +412,28 @@ fn render_fragments(
Fragment::Literal(text) => { Fragment::Literal(text) => {
output.push_str(text); output.push_str(text);
} }
Fragment::Variable { name } => { Fragment::Variable { name } => match ctx.get(name) {
match ctx.get(name) {
Some(val) => { Some(val) => {
output.push_str(&format!("{}", val)); output.push_str(&format!("{}", val));
} }
None => { None => {
return Err(PromptError::VariableNotFound(name.clone())); return Err(PromptError::VariableNotFound(name.clone()));
} }
} },
}
Fragment::If { Fragment::If {
condition, condition,
body, body,
else_body, else_body,
} => { } => {
let truthy = ctx let truthy = ctx.get(condition).map(|v| v.is_truthy()).unwrap_or(false);
.get(condition)
.map(|v| v.is_truthy())
.unwrap_or(false);
let target = if truthy { body } else { else_body }; let target = if truthy { body } else { else_body };
render_fragments(target, ctx, partials, output, depth + 1)?; render_fragments(target, ctx, partials, output, depth + 1)?;
} }
Fragment::Each { variable, body } => { Fragment::Each { variable, body } => {
let arr = match ctx.get(variable) { let arr = match ctx.get(variable) {
Some(val) => val.as_array().ok_or_else(|| { Some(val) => val
PromptError::NotAnArray(variable.clone()) .as_array()
})?, .ok_or_else(|| PromptError::NotAnArray(variable.clone()))?,
None => { None => {
return Err(PromptError::VariableNotFound(variable.clone())); return Err(PromptError::VariableNotFound(variable.clone()));
} }
@@ -504,10 +493,8 @@ impl PromptTemplateRegistry {
/// 延迟编译注册:只存储原始字符串,首次渲染时编译。 /// 延迟编译注册:只存储原始字符串,首次渲染时编译。
pub fn register_lazy(&mut self, name: &str, template: &str) { pub fn register_lazy(&mut self, name: &str, template: &str) {
self.templates.insert( self.templates
name.to_string(), .insert(name.to_string(), StoredTemplate::Raw(template.to_string()));
StoredTemplate::Raw(template.to_string()),
);
} }
/// 从文件读取并编译注册。 /// 从文件读取并编译注册。
+9 -3
View File
@@ -6,7 +6,9 @@ use std::sync::Arc;
#[derive(thiserror::Error, Debug, Clone)] #[derive(thiserror::Error, Debug, Clone)]
pub enum ToolError { pub enum ToolError {
/// 工具未注册。不可恢复——需调用方先 `registry.register(...)`。 /// 工具未注册。不可恢复——需调用方先 `registry.register(...)`。
#[error("工具 '{0}' 未注册。请先用 ToolRegistry::register(...) 注册该工具,或检查 LLM 输出的工具名拼写")] #[error(
"工具 '{0}' 未注册。请先用 ToolRegistry::register(...) 注册该工具,或检查 LLM 输出的工具名拼写"
)]
NotFound(String), NotFound(String),
/// 工具执行失败(可恢复——文本回传 LLM 由其决定重试或放弃)。 /// 工具执行失败(可恢复——文本回传 LLM 由其决定重试或放弃)。
@@ -14,11 +16,15 @@ pub enum ToolError {
ExecutionFailed(String, String), ExecutionFailed(String, String),
/// 工具参数无效(可恢复——文本回传 LLM)。 /// 工具参数无效(可恢复——文本回传 LLM)。
#[error("工具 '{0}' 参数无效: {1}。请检查 LLM 输出的参数是否符合 BaseTool::parameters() 声明的 JSON Schema")] #[error(
"工具 '{0}' 参数无效: {1}。请检查 LLM 输出的参数是否符合 BaseTool::parameters() 声明的 JSON Schema"
)]
InvalidArguments(String, String), InvalidArguments(String, String),
/// 权限被拒绝(不可恢复——终止循环)。 /// 权限被拒绝(不可恢复——终止循环)。
#[error("权限被拒绝: 工具 '{0}' 需要 {1} 权限。请在 PermissionConfig 中显式允许,或人工确认后绕过")] #[error(
"权限被拒绝: 工具 '{0}' 需要 {1} 权限。请在 PermissionConfig 中显式允许,或人工确认后绕过"
)]
PermissionDenied(String, String), PermissionDenied(String, String),
/// MCP 协议错误(不可恢复)。 /// MCP 协议错误(不可恢复)。
+14 -30
View File
@@ -9,16 +9,16 @@
use std::collections::HashMap; use std::collections::HashMap;
use std::process::Stdio; use std::process::Stdio;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc; use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration; use std::time::Duration;
use async_trait::async_trait; use async_trait::async_trait;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::{json, Value}; use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command}; use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::{oneshot, Mutex}; use tokio::sync::{Mutex, oneshot};
#[allow(deprecated)] #[allow(deprecated)]
use crate::llm::types::ToolDefinition; use crate::llm::types::ToolDefinition;
@@ -226,9 +226,7 @@ impl McpClient {
"version": env!("CARGO_PKG_VERSION") "version": env!("CARGO_PKG_VERSION")
} }
}); });
let _response = self let _response = self.send_request("initialize", Some(init_params)).await?;
.send_request("initialize", Some(init_params))
.await?;
// 发送 initialized 通知(无 id // 发送 initialized 通知(无 id
self.send_notification("notifications/initialized", Some(json!({}))) self.send_notification("notifications/initialized", Some(json!({})))
@@ -337,11 +335,7 @@ impl McpClient {
if let Some(state) = self.process.take() { if let Some(state) = self.process.take() {
let mut state = state.lock().await; let mut state = state.lock().await;
// 优雅等待 5 秒 // 优雅等待 5 秒
let graceful = tokio::time::timeout( let graceful = tokio::time::timeout(Duration::from_secs(5), state.child.wait()).await;
Duration::from_secs(5),
state.child.wait(),
)
.await;
if graceful.is_err() { if graceful.is_err() {
// 超时则强杀 // 超时则强杀
let _ = state.child.kill().await; let _ = state.child.kill().await;
@@ -372,11 +366,7 @@ impl McpClient {
tools tools
} }
async fn send_request( async fn send_request(&self, method: &str, params: Option<Value>) -> Result<Value, ToolError> {
&self,
method: &str,
params: Option<Value>,
) -> Result<Value, ToolError> {
let state_arc = self let state_arc = self
.process .process
.as_ref() .as_ref()
@@ -412,9 +402,11 @@ impl McpClient {
.write_all(b"\n") .write_all(b"\n")
.await .await
.map_err(|e| ToolError::McpError(format!("写入换行失败: {e}")))?; .map_err(|e| ToolError::McpError(format!("写入换行失败: {e}")))?;
state.stdin.flush().await.map_err(|e| { state
ToolError::McpError(format!("flush stdin 失败: {e}")) .stdin
})?; .flush()
.await
.map_err(|e| ToolError::McpError(format!("flush stdin 失败: {e}")))?;
} }
// 等待响应(带超时) // 等待响应(带超时)
@@ -471,10 +463,7 @@ impl McpClient {
} }
/// 持续读取 stdout,将响应分发到对应的 oneshot sender。 /// 持续读取 stdout,将响应分发到对应的 oneshot sender。
async fn read_loop( async fn read_loop(mut reader: BufReader<ChildStdout>, state: Arc<Mutex<ChildProcessState>>) {
mut reader: BufReader<ChildStdout>,
state: Arc<Mutex<ChildProcessState>>,
) {
let mut line = String::new(); let mut line = String::new();
loop { loop {
line.clear(); line.clear();
@@ -556,11 +545,7 @@ impl BaseTool for McpToolAdapter {
self.parameters.clone() self.parameters.clone()
} }
async fn execute( async fn execute(&self, _args: Value, _ctx: &ToolContext<'_>) -> Result<Value, ToolError> {
&self,
_args: Value,
_ctx: &ToolContext<'_>,
) -> Result<Value, ToolError> {
// 当前 Phase 2 实现的简化:McpToolAdapter 不持有活跃 MCP 连接。 // 当前 Phase 2 实现的简化:McpToolAdapter 不持有活跃 MCP 连接。
// 实际生产中应持有 Arc<McpClient> 并通过 mcp.call_tool() 执行。 // 实际生产中应持有 Arc<McpClient> 并通过 mcp.call_tool() 执行。
// 这里返回错误,提示需要通过其他方式调用 MCP 工具。 // 这里返回错误,提示需要通过其他方式调用 MCP 工具。
@@ -617,8 +602,7 @@ mod tests {
#[test] #[test]
fn test_jsonrpc_response_parse_error() { fn test_jsonrpc_response_parse_error() {
let s = let s = r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"Method not found"}}"#;
r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"Method not found"}}"#;
let resp: JsonRpcResponse = serde_json::from_str(s).unwrap(); let resp: JsonRpcResponse = serde_json::from_str(s).unwrap();
assert_eq!(resp.id, 1); assert_eq!(resp.id, 1);
assert!(resp.result.is_none()); assert!(resp.result.is_none());
+18 -15
View File
@@ -148,9 +148,7 @@ mod tests {
#[test] #[test]
fn test_default_config_denies_delete() { fn test_default_config_denies_delete() {
let checker = PermissionChecker::new(PermissionConfig::default()); let checker = PermissionChecker::new(PermissionConfig::default());
assert!(checker assert!(checker.check("rm_file", &p(Permission::Delete)).is_err());
.check("rm_file", &p(Permission::Delete))
.is_err());
} }
#[test] #[test]
@@ -246,12 +244,16 @@ mod tests {
allow_unspecified: false, allow_unspecified: false,
}; };
let checker = PermissionChecker::new(cfg); let checker = PermissionChecker::new(cfg);
assert!(checker assert!(
checker
.check("t", &[Permission::Custom("db:read".into())]) .check("t", &[Permission::Custom("db:read".into())])
.is_ok()); .is_ok()
assert!(checker );
assert!(
checker
.check("t", &[Permission::Custom("db:write".into())]) .check("t", &[Permission::Custom("db:write".into())])
.is_err()); .is_err()
);
} }
#[test] #[test]
@@ -262,12 +264,11 @@ mod tests {
allow_unspecified: false, allow_unspecified: false,
}; };
let checker = PermissionChecker::new(cfg); let checker = PermissionChecker::new(cfg);
assert!(checker assert!(
.check( checker
"t", .check("t", &[Permission::Read, Permission::Network])
&[Permission::Read, Permission::Network] .is_ok()
) );
.is_ok());
} }
#[test] #[test]
@@ -279,8 +280,10 @@ mod tests {
}; };
let checker = PermissionChecker::new(cfg); let checker = PermissionChecker::new(cfg);
// 任一权限不在白名单则拒绝 // 任一权限不在白名单则拒绝
assert!(checker assert!(
checker
.check("t", &[Permission::Read, Permission::Write]) .check("t", &[Permission::Read, Permission::Write])
.is_err()); .is_err()
);
} }
} }
+6 -3
View File
@@ -348,7 +348,10 @@ mod tests {
async fn test_invoke_success() { async fn test_invoke_success() {
let mut reg = ToolRegistry::new(); let mut reg = ToolRegistry::new();
reg.register(Arc::new(AddTool { base: 100 })).unwrap(); reg.register(Arc::new(AddTool { base: 100 })).unwrap();
let result = reg.invoke("call_1", "add", json!({ "n": 5 })).await.unwrap(); let result = reg
.invoke("call_1", "add", json!({ "n": 5 }))
.await
.unwrap();
let value = result.output.unwrap(); let value = result.output.unwrap();
assert_eq!(value["result"], 105); assert_eq!(value["result"], 105);
assert_eq!(result.tool_call_id, "call_1"); assert_eq!(result.tool_call_id, "call_1");
@@ -372,8 +375,8 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_invoke_with_permission_denied() { async fn test_invoke_with_permission_denied() {
let mut reg = ToolRegistry::new() let mut reg =
.with_permission_checker(PermissionChecker::new(Default::default())); ToolRegistry::new().with_permission_checker(PermissionChecker::new(Default::default()));
reg.register(Arc::new(ShellTool)).unwrap(); reg.register(Arc::new(ShellTool)).unwrap();
let result = reg.invoke("call_z", "shell", json!({})).await; let result = reg.invoke("call_z", "shell", json!({})).await;
assert!(matches!(result, Err(ToolError::PermissionDenied(_, _)))); assert!(matches!(result, Err(ToolError::PermissionDenied(_, _))));