abu_tool/
lib.rs

1mod register;
2pub use register::ToolRegister;
3pub use abu_base::chat::ToolDefinition;
4use serde_json::Value;
5
6#[async_trait::async_trait]
7pub trait Tool: Send + Sync {
8    fn name(&self) -> &'static str;
9    fn description(&self) -> &'static str;
10    fn parameters(&self) -> Vec<ToolParameter>;
11    async fn execute(&self, args: Value) -> ToolResult<ToolCallResult>;
12
13    fn to_function_define(&self) -> ToolDefinition {
14        ToolDefinition {
15            name: self.name().to_string(),
16            description: self.description().to_string(),
17            schema: build_params_schema(&self.parameters())
18        }
19    }
20}
21
22#[derive(Debug, Clone)]
23pub struct ToolParameter {
24    pub name: String,
25    pub required: bool,
26    pub description: Option<String>,
27    pub kind: ToolParameterKind,
28}
29
30#[derive(Debug, Clone)]
31pub enum ToolParameterKind {
32    Object(Vec<ToolParameter>),
33    Array(Box<ToolParameterKind>),
34    String(Option<Vec<String>>),
35    Integer,
36    Number,
37    Boolean,
38}
39
40impl ToolParameter {
41    pub fn integer(name: impl Into<String>) -> Self {
42        Self {
43            name: name.into(),
44            required: true,
45            description: None,
46            kind: ToolParameterKind::Integer
47        }
48    }
49
50    pub fn number(name: impl Into<String>) -> Self {
51        Self {
52            name: name.into(),
53            required: true,
54            description: None,
55            kind: ToolParameterKind::Number
56        }
57    }
58
59    pub fn string(name: impl Into<String>) -> Self {
60        Self {
61            name: name.into(),
62            required: true,
63            description: None,
64            kind: ToolParameterKind::String(None)
65        }
66    }
67
68    pub fn string_with(name: impl Into<String>, enums: Vec<String>) -> Self {
69        Self {
70            name: name.into(),
71            required: true,
72            description: None,
73            kind: ToolParameterKind::String(Some(enums)),
74        }
75    }
76
77    pub fn array(name: impl Into<String>, kind: ToolParameterKind) -> Self {
78        Self {
79            name: name.into(),
80            required: true,
81            description: None,
82            kind: ToolParameterKind::Array(Box::new(kind))
83        }
84    }
85
86    pub fn object(name: impl Into<String>, field: Vec<ToolParameter>) -> Self {
87        Self {
88            name: name.into(),
89            required: true,
90            description: None,
91            kind: ToolParameterKind::Object(field),
92        }
93    }
94
95    pub fn required(self, value: bool) -> Self {
96        Self { required: value, ..self }
97    }
98
99    pub fn description(self, value: impl Into<String>) -> Self {
100        Self { description: Some(value.into()), ..self }
101    }
102
103    pub fn to_schema(&self) -> serde_json::Value {
104        let mut schema = self.kind.to_schema();
105        if let Some(desc) = self.description.as_ref() {
106            schema["description"] = serde_json::Value::String(desc.to_string());
107        }
108        schema
109    }
110
111    pub fn build_params_properties(params: &[ToolParameter]) -> serde_json::Map<String, serde_json::Value> {
112        let mut properties = serde_json::Map::new();
113        for param in params {
114            let mut param_schema = param.kind.to_schema();        
115            if let Some(desc) = &param.description {
116                param_schema["description"] = serde_json::json!(desc);
117            }
118            
119            properties.insert(param.name.clone(), param_schema);
120        }
121        properties
122    }
123    
124    pub fn build_params_required(params: &[ToolParameter]) -> Vec<String> {
125        params.iter()
126            .filter(|p| p.required)
127            .map(|p| p.name.clone())
128            .collect()
129    }
130}
131
132impl ToolParameterKind {
133    pub fn to_schema(&self) -> serde_json::Value {
134        match &self {
135            Self::Object(params) => build_params_schema(&params),
136            Self::Array(kind) => serde_json::json!({
137                "type": "array",
138                "items": kind.to_schema(),
139            }),
140            Self::String(enums) => match enums {
141                Some(enums) => serde_json::json!({ "type": "string", "enums": enums }),
142                None => serde_json::json!({ "type": "string" }),
143            }
144            Self::Boolean => serde_json::json!({ "type": "boolean" }),
145            Self::Number => serde_json::json!({ "type": "number" }),
146            Self::Integer => serde_json::json!({ "type": "integer" }),
147        }
148    }
149}
150
151pub(crate) fn build_params_schema(params: &[ToolParameter]) -> serde_json::Value {
152    let properties = ToolParameter::build_params_properties(params);
153    let required = ToolParameter::build_params_required(params);
154    serde_json::json!({
155        "type": "object",
156        "properties": properties,
157        "required": required,
158    })
159}
160
161#[derive(Debug, Clone)]
162pub struct ToolCallResult {
163    pub is_error: bool,
164    pub context: String,
165}
166
167impl ToolCallResult {
168    pub fn error(context: impl Into<String>) -> Self {
169        Self { is_error: true, context: context.into() }
170    }
171
172    pub fn success(context: impl Into<String>) -> Self {
173        Self { is_error: false, context: context.into() }
174    }
175}
176
177#[derive(Debug, thiserror::Error)]
178pub enum ToolError {
179    #[error(transparent)]
180    SerdeJson(#[from] serde_json::Error),
181
182    #[error("tool {0} not found")]
183    ToolNotFound(String),
184
185    #[error("arg {0} not found")]
186    ArgNotFound(String),
187
188    #[error("arg parse failed, expect: Expect {0}")]
189    ArgParse(&'static str),
190}
191
192pub type ToolResult<T> = std::result::Result<T, ToolError>;
abu_tool/lib.rs

abu_tool/
lib.rs
