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