agcore/examples/custom_tool.rs

//! custom_tool —— 自定义工具注册、单次 / 并行调用、权限检查。
//!
//! 演示：
//! 1. 实现 `BaseTool` trait（WeatherTool + DeleteFileTool）
//! 2. 注册到 `ToolRegistry`
//! 3. 单次 `invoke`（含 tool_call_id 关联）
//! 4. 并行 `invoke_all`
//! 5. 调用未注册工具 → `ToolError::NotFound`
//! 6. `PermissionChecker` 黑名单阻断 `DeleteFileTool`
//!
//! 运行：`cargo run --example custom_tool`

use std::sync::Arc;

use agcore::tools::{
    BaseTool, Permission, PermissionChecker, PermissionConfig, ToolContext, ToolError, ToolRef,
    ToolRegistry,
};
use async_trait::async_trait;
use serde_json::{json, Value};

/// 天气查询工具 —— 模拟根据城市返回天气数据。
struct WeatherTool;

#[async_trait]
impl BaseTool for WeatherTool {
    fn name(&self) -> &str {
        "get_weather"
    }
    fn description(&self) -> &str {
        "查询指定城市的天气"
    }
    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "city": { "type": "string", "description": "城市名" }
            },
            "required": ["city"]
        })
    }
    fn required_permissions(&self) -> Vec<Permission> {
        vec![Permission::Network]
    }
    async fn execute(
        &self,
        args: Value,
        _ctx: &ToolContext<'_>,
    ) -> Result<Value, ToolError> {
        let city = args["city"].as_str().unwrap_or("未知");
        // 模拟查询：根据城市名给出不同温度
        let (temperature, condition) = match city {
            "北京" => (22_i32, "晴"),
            "上海" => (25, "多云"),
            "广州" => (28, "雷阵雨"),
            _ => (20, "晴"),
        };
        Ok(json!({
            "city": city,
            "temperature": temperature,
            "condition": condition
        }))
    }
}

/// 删除文件工具 —— 用于演示权限黑名单阻断。
struct DeleteFileTool;

#[async_trait]
impl BaseTool for DeleteFileTool {
    fn name(&self) -> &str {
        "delete_file"
    }
    fn description(&self) -> &str {
        "删除指定路径的文件"
    }
    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": { "path": { "type": "string" } },
            "required": ["path"]
        })
    }
    fn required_permissions(&self) -> Vec<Permission> {
        vec![Permission::Delete]
    }
    async fn execute(
        &self,
        _args: Value,
        _ctx: &ToolContext<'_>,
    ) -> Result<Value, ToolError> {
        Ok(json!({"deleted": true}))
    }
}

#[tokio::main]
async fn main() {
    // 1. 注册工具
    let mut registry = ToolRegistry::new();
    let weather: ToolRef = Arc::new(WeatherTool);
    let deleter: ToolRef = Arc::new(DeleteFileTool);
    registry.register(weather).expect("注册 get_weather 失败");
    registry.register(deleter).expect("注册 delete_file 失败");
    println!("=== 已注册工具 ===");
    println!("{:?}", registry.list_tools());

    // 2. 单次 invoke：tool_call_id 用于回传时与原始调用关联
    println!("\n=== 单次 invoke ===");
    let result = registry
        .invoke("call_1", "get_weather", json!({"city": "北京"}))
        .await
        .expect("invoke 失败");
    println!("tool_call_id: {}", result.tool_call_id);
    println!("tool_name: {}", result.tool_name);
    println!("output: {}", result.output.unwrap());

    // 3. 并行 invoke_all：三个并行天气查询
    println!("\n=== 并行 invoke_all（30s 超时）===");
    let calls = vec![
        ("c1".into(), "get_weather".into(), json!({"city": "北京"})),
        ("c2".into(), "get_weather".into(), json!({"city": "上海"})),
        ("c3".into(), "get_weather".into(), json!({"city": "广州"})),
    ];
    let results = registry.invoke_all(calls, 30).await;
    assert_eq!(results.len(), 3);
    for r in &results {
        let output = r.output.as_ref().unwrap();
        println!("[{}] {}", r.tool_call_id, output);
    }

    // 4. 调用未注册工具
    println!("\n=== 未注册工具 ===");
    let err = registry
        .invoke("c_x", "nope_tool", json!({}))
        .await
        .unwrap_err();
    println!("错误: {err}");

    // 5. 权限检查：默认 PermissionConfig 黑名单含 Delete
    println!("\n=== 权限检查（默认 PermissionConfig，denied = [Delete, Shell]）===");
    let mut registry_with_checker = ToolRegistry::new().with_permission_checker(PermissionChecker::new(
        PermissionConfig::default(),
    ));
    registry_with_checker
        .register(Arc::new(WeatherTool) as ToolRef)
        .unwrap();
    registry_with_checker
        .register(Arc::new(DeleteFileTool) as ToolRef)
        .unwrap();

    // get_weather 声明 Network → 在 allowed 列表 → 通过
    let r = registry_with_checker
        .invoke("c1", "get_weather", json!({"city": "北京"}))
        .await
        .unwrap();
    println!(
        "get_weather 权限检查: {}",
        if r.output.is_ok() { "通过 ✓" } else { "阻断 ✗" }
    );

    // delete_file 声明 Delete → 在 denied 列表 → 阻断
    let err = registry_with_checker
        .invoke("c2", "delete_file", json!({"path": "/tmp/x"}))
        .await
        .unwrap_err();
    println!("delete_file 权限检查: 阻断 ✗ ({err})");

    println!("\n✓ custom_tool 完成");
}