Skip to main content

tiny_agent/tools/
mod.rs

1mod default_tools;
2
3pub use default_tools::{register_default_tools, register_skill_tools, register_spawn_agent};
4
5use super::sandbox::Sandbox;
6use crate::shared::{ToolDefinition, UserInteraction};
7use schemars::{JsonSchema, schema_for};
8use serde::de::DeserializeOwned;
9use serde_json::Value;
10use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
11
12type BoxFuture = Pin<Box<dyn Future<Output = Result<ToolOutcome, String>> + Send>>;
13type BoxedCall = Arc<dyn Fn(Arc<dyn Sandbox>, Value) -> BoxFuture + Send + Sync>;
14
15#[derive(Debug, Clone)]
16pub enum ToolOutcome {
17    Completed(Value),
18    NeedsUserInteraction(UserInteraction),
19}
20
21impl From<Value> for ToolOutcome {
22    fn from(value: Value) -> Self {
23        Self::Completed(value)
24    }
25}
26
27#[derive(Clone)]
28pub struct ToolEntry {
29    pub name: String,
30    pub description: String,
31    pub schema: Value,
32    /// 是否只读、无副作用。只读工具之间可以并发执行(见 `core::tool_execution` 的分段调度);
33    /// 任何会改动 sandbox 状态的工具都必须保持 `false`,以免与并发读发生 read-write race。
34    pub read_only: bool,
35    call: BoxedCall,
36}
37
38#[derive(Clone)]
39pub struct ToolRegistry {
40    tools: HashMap<String, ToolEntry>,
41}
42
43impl ToolRegistry {
44    pub fn new() -> Self {
45        Self {
46            tools: HashMap::new(),
47        }
48    }
49
50    /// 注册一个工具。`read_only` 标记它是否无副作用 —— 只读工具之间可被并发调度
51    /// (见 `core::tool_execution` 的分段执行),任何会改动 sandbox 的工具都必须传 `false`。
52    /// handler 统一返回 `ToolOutcome`(普通完成用 `Value.into()`,需要用户交互用 `NeedsUserInteraction`)。
53    pub fn register<A, F, Fut>(
54        &mut self,
55        name: &str,
56        description: &str,
57        read_only: bool,
58        handler: F,
59    ) where
60        A: DeserializeOwned + JsonSchema + Send + 'static,
61        F: Fn(Arc<dyn Sandbox>, A) -> Fut + Send + Sync + 'static,
62        Fut: Future<Output = Result<ToolOutcome, String>> + Send + 'static,
63    {
64        let schema = serde_json::to_value(schema_for!(A)).unwrap();
65        let handler = Arc::new(handler);
66        let call: BoxedCall = Arc::new(move |sandbox: Arc<dyn Sandbox>, v: Value| {
67            let handler = handler.clone();
68            Box::pin(async move {
69                let args: A = serde_json::from_value(v).map_err(|e| format!("参数不合法: {e}"))?;
70                handler(sandbox, args).await
71            })
72        });
73        self.tools.insert(
74            name.into(),
75            ToolEntry {
76                name: name.into(),
77                description: description.into(),
78                schema,
79                read_only,
80                call,
81            },
82        );
83    }
84
85    /// 工具是否声明为只读。未知工具按非只读处理(最保守,强制串行)。
86    pub fn is_read_only(&self, name: &str) -> bool {
87        self.tools.get(name).is_some_and(|t| t.read_only)
88    }
89
90    /// 按名字调度执行,执行时注入该会话绑定的 sandbox。
91    pub async fn call(
92        &self,
93        name: &str,
94        args: Value,
95        sandbox: Arc<dyn Sandbox>,
96    ) -> Result<ToolOutcome, String> {
97        let tool = self
98            .tools
99            .get(name)
100            .ok_or_else(|| format!("工具不存在: {name}"))?;
101        (tool.call)(sandbox, args).await
102    }
103
104    /// 将工具 struct 转成 `ToolDefinition` 列表
105    pub fn definitions(&self) -> Vec<ToolDefinition> {
106        self.tools
107            .values()
108            .map(|t| ToolDefinition {
109                name: t.name.clone(),
110                desc: t.description.clone(),
111                arguments: t.schema.clone(),
112            })
113            .collect()
114    }
115}
116
117impl Default for ToolRegistry {
118    fn default() -> Self {
119        Self::new()
120    }
121}