徐涛
85b92ae9d4
按 PM/SA/Code Reviewer 三方审计报告修复： 🔴 阻塞修复: - CalcTool 除零 panic: `a / b` 改 `a.checked_div(b).ok_or_else(...)`, b=0 时返回 ToolError::InvalidArguments 而非 panic 🟡 警告修复: - end_to_end.rs 持久化验证注释与实际不符: 显式 drop(backend) 让注释 描述与 Arc 释放顺序一致 - quick_start EchoTool 参数验证: 用 args.get("text").and_then().ok_or_else() 替换 as_str().unwrap_or("") 静默降级, 缺失/类型错误时返回显式错误 💭 一致性修复: - end_to_end.rs EchoTool 与 quick_start 一致化 (format!("收到: {text}")) - quick_start mock 响应文本 "已通过 echo 回传" → "EchoTool 已收到并完成回传" - quick_start 断言改为检查 "收到", 与方案字面要求一致 - quick_start 末尾追加 POSIX trailing newline 验收: cargo test 200/0 + clippy 0 警告 + doc 0 warning + 10 示例 exit 0
247 lines
11 KiB
Rust
247 lines
11 KiB
Rust
//! end_to_end —— 3 工具 + 3 轮对话 + SqliteStore 持久化跨连接验证。
|
||
//!
|
||
//! 运行：`cargo run --example end_to_end`（离线，零配置）
|
||
//!
|
||
//! ## 真实 LLM Provider 切换
|
||
//!
|
||
//! 设置环境变量即可使用真实 LLM Provider：
|
||
//! - `AG_LLM_BASE_URL` —— API 端点（如 `https://api.openai.com/v1`）
|
||
//! - `AG_LLM_API_KEY` —— API key
|
||
//! - `AG_LLM_MODEL` —— 模型名（如 `gpt-4o-mini`）
|
||
//! - `AG_LLM_PROVIDER`（可选）—— Provider 类型，默认 OpenaiChat（OpenAI / DeepSeek / Qwen / Ollama）
|
||
//!
|
||
//! 未设置上述变量时自动降级为 MockProvider，零配置可运行。
|
||
|
||
use std::env;
|
||
use std::sync::Arc;
|
||
|
||
use agcore::agent::{Agent, AgentBuilder, AgentSession};
|
||
use agcore::llm::hooks::HookExecutor;
|
||
use agcore::llm::mock::MockProvider;
|
||
use agcore::llm::provider::{create_provider, LlmProvider, ProviderConfig, ProviderType};
|
||
use agcore::llm::types::{Usage, message::{ContentBlock, Message}, response_v2::{MessageResponse, StopReason}};
|
||
use agcore::memory::store::{MemoryStore, SqliteStore};
|
||
use agcore::memory::types::{MemoryFilter, MemoryItem};
|
||
use agcore::tools::{BaseTool, ToolContext, ToolError, ToolRegistry};
|
||
use async_trait::async_trait;
|
||
use serde_json::{Value, json};
|
||
use tempfile::TempDir;
|
||
use time::OffsetDateTime;
|
||
|
||
// === Agent ===
|
||
|
||
struct AssistantAgent;
|
||
impl Agent for AssistantAgent {
|
||
fn name(&self) -> &str { "end-to-end assistant" }
|
||
fn system_prompt(&self) -> Option<&str> { Some("简洁助手，必要时调用工具完成任务。") }
|
||
}
|
||
|
||
// === Tools ===
|
||
|
||
struct EchoTool;
|
||
#[async_trait]
|
||
impl BaseTool for EchoTool {
|
||
fn name(&self) -> &str { "echo" }
|
||
fn description(&self) -> &str { "回显输入文本" }
|
||
fn parameters(&self) -> Value {
|
||
json!({"type":"object","properties":{"text":{"type":"string"}},"required":["text"]})
|
||
}
|
||
async fn execute(&self, args: Value, _: &ToolContext<'_>) -> Result<Value, ToolError> {
|
||
let text = args.get("text").and_then(|v| v.as_str())
|
||
.ok_or_else(|| ToolError::InvalidArguments("text".into(), "需要 string 类型的 text 参数".into()))?;
|
||
Ok(json!({"echoed": format!("收到: {text}")}))
|
||
}
|
||
}
|
||
|
||
/// 四则运算：'a op b' 格式（ponytail: 基础 +-*/ 不引入 rhai 依赖）。
|
||
struct CalcTool;
|
||
#[async_trait]
|
||
impl BaseTool for CalcTool {
|
||
fn name(&self) -> &str { "calc" }
|
||
fn description(&self) -> &str { "四则运算：'a op b' 格式，op ∈ {+, -, *, /}" }
|
||
fn parameters(&self) -> Value {
|
||
json!({"type":"object","properties":{"expr":{"type":"string"}},"required":["expr"]})
|
||
}
|
||
async fn execute(&self, args: Value, _: &ToolContext<'_>) -> Result<Value, ToolError> {
|
||
let expr = args["expr"].as_str().unwrap_or("");
|
||
let parts: Vec<&str> = expr.split_whitespace().collect();
|
||
if parts.len() != 3 {
|
||
return Err(ToolError::InvalidArguments("expr".into(), "需要 'a op b' 三段式".into()));
|
||
}
|
||
let a: i64 = parts[0].parse().map_err(|_| ToolError::InvalidArguments("expr".into(), format!("无法解析 '{}'", parts[0])))?;
|
||
let b: i64 = parts[2].parse().map_err(|_| ToolError::InvalidArguments("expr".into(), format!("无法解析 '{}'", parts[2])))?;
|
||
let result = match parts[1] {
|
||
"+" => a + b,
|
||
"-" => a - b,
|
||
"*" => a * b,
|
||
"/" => a.checked_div(b).ok_or_else(|| {
|
||
ToolError::InvalidArguments("expr".into(), "除数不能为 0".into())
|
||
})?,
|
||
op => return Err(ToolError::InvalidArguments("expr".into(), format!("不支持的运算符: {op}"))),
|
||
};
|
||
Ok(json!({"result": result}))
|
||
}
|
||
}
|
||
|
||
/// 通过 MemoryStore trait 读写笔记：直接持有 Arc<dyn MemoryStore>，
|
||
/// 绕开 AgentSession 封装（NoteTool 在 tool.execute 中直接操作 store）。
|
||
/// 关键前缀 "note:" 用于 list 过滤。
|
||
struct NoteTool { store: Arc<dyn MemoryStore> }
|
||
impl NoteTool { const PREFIX: &'static str = "note:"; }
|
||
|
||
#[async_trait]
|
||
impl BaseTool for NoteTool {
|
||
fn name(&self) -> &str { "note" }
|
||
fn description(&self) -> &str { "笔记 save/query: save(key, content) / query()" }
|
||
fn parameters(&self) -> Value {
|
||
json!({
|
||
"type":"object",
|
||
"properties":{
|
||
"action":{"type":"string","enum":["save","query"]},
|
||
"key":{"type":"string"},
|
||
"content":{"type":"string"}
|
||
},
|
||
"required":["action"]
|
||
})
|
||
}
|
||
async fn execute(&self, args: Value, _: &ToolContext<'_>) -> Result<Value, ToolError> {
|
||
let action = args["action"].as_str().unwrap_or("");
|
||
match action {
|
||
"save" => {
|
||
let key = args["key"].as_str().unwrap_or("");
|
||
let content = args["content"].as_str().unwrap_or("");
|
||
let item = MemoryItem {
|
||
id: format!("{}{}", Self::PREFIX, key),
|
||
content: content.to_string(),
|
||
metadata: json!({}),
|
||
created_at: OffsetDateTime::now_utc(),
|
||
};
|
||
self.store.save(item).await
|
||
.map_err(|e| ToolError::ExecutionFailed("note".into(), e.to_string()))?;
|
||
Ok(json!({"saved": key}))
|
||
}
|
||
"query" => {
|
||
let filter = MemoryFilter { prefix: Some(Self::PREFIX.into()), ..Default::default() };
|
||
let items = self.store.list(&filter).await
|
||
.map_err(|e| ToolError::ExecutionFailed("note".into(), e.to_string()))?;
|
||
let notes: Vec<String> = items.into_iter().map(|i| i.content).collect();
|
||
Ok(json!({"notes": notes}))
|
||
}
|
||
_ => Err(ToolError::InvalidArguments("action".into(), format!("未知 action: {action}"))),
|
||
}
|
||
}
|
||
}
|
||
|
||
// === Mock response helper ===
|
||
|
||
fn resp(content: Vec<ContentBlock>, stop: StopReason, u: (u32, u32)) -> MessageResponse {
|
||
MessageResponse { id: String::new(), model: "mock".into(),
|
||
message: Message::Assistant { content },
|
||
usage: Usage::from_input_output(u.0, u.1),
|
||
stop_reason: stop, extra: Default::default() }
|
||
}
|
||
|
||
fn mock_responses() -> Vec<MessageResponse> {
|
||
vec![
|
||
// 第 1 轮：calc(25 * 4) → tool_result(100) → 文本回答
|
||
resp(vec![ContentBlock::ToolUse { id: "t1".into(), name: "calc".into(),
|
||
input: json!({"expr": "25 * 4"}) }], StopReason::ToolUse, (5, 8)),
|
||
resp(vec![ContentBlock::Text { text: "25 * 4 = 100".into() }], StopReason::Stop, (8, 12)),
|
||
// 第 2 轮：note(save, last_calc, "100") → tool_result(saved) → 文本回答
|
||
resp(vec![ContentBlock::ToolUse { id: "t2".into(), name: "note".into(),
|
||
input: json!({"action": "save", "key": "last_calc", "content": "100"}) }],
|
||
StopReason::ToolUse, (10, 14)),
|
||
resp(vec![ContentBlock::Text { text: "已记录：last_calc = 100".into() }], StopReason::Stop, (12, 16)),
|
||
// 第 3 轮：note(query) → tool_result([100]) → 文本回答
|
||
resp(vec![ContentBlock::ToolUse { id: "t3".into(), name: "note".into(),
|
||
input: json!({"action": "query"}) }], StopReason::ToolUse, (8, 8)),
|
||
resp(vec![ContentBlock::Text { text: "您刚才的计算结果是 100".into() }], StopReason::Stop, (10, 14)),
|
||
// 后续冗余响应（防止队列耗尽报错）
|
||
resp(vec![ContentBlock::Text { text: "done".into() }], StopReason::Stop, (1, 1)),
|
||
resp(vec![ContentBlock::Text { text: "done".into() }], StopReason::Stop, (1, 1)),
|
||
resp(vec![ContentBlock::Text { text: "done".into() }], StopReason::Stop, (1, 1)),
|
||
]
|
||
}
|
||
|
||
// === Provider selection ===
|
||
|
||
fn select_provider() -> Arc<dyn LlmProvider> {
|
||
if env::var("AG_LLM_BASE_URL").is_ok() && env::var("AG_LLM_API_KEY").is_ok() {
|
||
let cfg = ProviderConfig::from_env("AG_LLM").expect("AG_LLM_* 环境变量解析失败");
|
||
let provider_type = env::var("AG_LLM_PROVIDER").ok()
|
||
.and_then(|s| s.parse::<ProviderType>().ok())
|
||
.unwrap_or(ProviderType::OpenaiChat);
|
||
Arc::from(create_provider(provider_type, cfg).expect("Provider 创建失败"))
|
||
} else {
|
||
let found: Vec<&str> = ["AG_LLM_BASE_URL", "AG_LLM_API_KEY", "AG_LLM_MODEL"]
|
||
.iter().filter(|k| env::var(k).is_ok()).copied().collect();
|
||
eprintln!("AG_LLM_* 环境变量不完整（检测到: {:?}），回退到 MockProvider", found);
|
||
Arc::new(MockProvider::new(mock_responses()))
|
||
}
|
||
}
|
||
|
||
// === Main ===
|
||
|
||
#[tokio::main]
|
||
async fn main() {
|
||
println!("=== agcore 端到端演示 ===");
|
||
let dir = TempDir::new().expect("TempDir 创建失败");
|
||
let db_path = dir.path().join("agcore.db");
|
||
let backend: Arc<dyn MemoryStore> =
|
||
Arc::new(SqliteStore::open(&db_path).expect("SqliteStore 打开失败"));
|
||
println!("💾 SqliteStore: {}", db_path.display());
|
||
let provider_label = if env::var("AG_LLM_BASE_URL").is_ok() && env::var("AG_LLM_API_KEY").is_ok() {
|
||
"真实 LLM Provider"
|
||
} else {
|
||
"MockProvider (离线回退模式)"
|
||
};
|
||
println!("🔄 Provider: {provider_label}");
|
||
|
||
let mut registry = ToolRegistry::new();
|
||
registry.register(Arc::new(EchoTool)).unwrap();
|
||
registry.register(Arc::new(CalcTool)).unwrap();
|
||
registry.register(Arc::new(NoteTool { store: backend.clone() })).unwrap();
|
||
println!("🔧 注册工具: {:?}", registry.list_tools());
|
||
|
||
let bundle = Arc::new(AgentBuilder::new()
|
||
.provider(select_provider())
|
||
.tool_registry(Arc::new(registry))
|
||
.hook_executor(Arc::new(HookExecutor::new()))
|
||
.build().expect("RuntimeBundle 装配失败"));
|
||
|
||
let mut session = AgentSession::new(Arc::new(AssistantAgent), "e2e-1", bundle.clone());
|
||
|
||
println!("\n第 1 轮 用户: 帮我算 25 * 4");
|
||
let r1 = session.submit_turn("帮我算 25 * 4").await.expect("turn 1 失败");
|
||
println!(" → 回答: {}", r1.text());
|
||
|
||
println!("\n第 2 轮 用户: 记下来：结果是 100");
|
||
let r2 = session.submit_turn("记下来：结果是 100").await.expect("turn 2 失败");
|
||
println!(" → 回答: {}", r2.text());
|
||
|
||
println!("\n第 3 轮 用户: 我刚才算了什么？");
|
||
let r3 = session.submit_turn("我刚才算了什么？").await.expect("turn 3 失败");
|
||
println!(" → 回答: {}", r3.text());
|
||
|
||
let total = session.usage().total();
|
||
println!("\n📊 用量: prompt={}, completion={}, total={}",
|
||
total.prompt_tokens, total.completion_tokens, total.total_tokens);
|
||
|
||
println!("\n=== 持久化验证 ===");
|
||
// 显式释放所有对 backend 的 Arc 引用，确保 SqliteStore Connection 真正关闭。
|
||
// 释放顺序：session → bundle（间接持有 NoteTool → backend clone）→ backend 局部变量。
|
||
drop(session); // session.bundle Arc 计数 -1
|
||
drop(bundle); // bundle Arc 计数归零 → registry → NoteTool → backend clone Arc 计数 2→1
|
||
drop(backend); // backend 局部变量 Arc 计数 1→0 → SqliteStore::drop → Connection 自动 close
|
||
let backend2: Arc<dyn MemoryStore> =
|
||
Arc::new(SqliteStore::open(&db_path).expect("重开 SqliteStore 失败"));
|
||
let filter = MemoryFilter { prefix: Some("note:".into()), ..Default::default() };
|
||
let items = backend2.list(&filter).await.expect("list 失败");
|
||
println!("✓ 跨连接数据存活: 找到 {} 条 note", items.len());
|
||
assert!(!items.is_empty(), "持久化验证失败：重开后无数据");
|
||
for i in &items {
|
||
println!(" - {} = {}", i.id, i.content);
|
||
}
|
||
|
||
println!("\n✓ 端到端演示完成");
|
||
} |