abu_tool/
lib.rs

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

abu_tool/
lib.rs
