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 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 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 pub fn is_read_only(&self, name: &str) -> bool {
87 self.tools.get(name).is_some_and(|t| t.read_only)
88 }
89
90 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 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}