1use jsonschema::Validator;
4use serde::{Deserialize, Serialize};
5use serde_json::Value as JsonValue;
6use serde_json::json;
7use std::sync::Arc;
8use thiserror::Error;
9
10pub type ProgressCallback = Arc<dyn Fn(String) + Send + Sync>;
12
13pub fn progress_callback<F: Fn(String) + Send + Sync + 'static>(f: F) -> ProgressCallback {
15 Arc::new(f)
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
25#[serde(tag = "type", content = "name", rename_all = "snake_case")]
26pub enum ToolChoice {
27 #[default]
29 Auto,
30 Named(String),
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct Tool {
38 pub name: String,
40
41 pub description: String,
43
44 pub parameters: JsonValue,
46}
47
48impl Tool {
49 pub fn new(
71 name: impl Into<String>,
72 description: impl Into<String>,
73 parameters: JsonValue,
74 ) -> Self {
75 Self {
76 name: name.into(),
77 description: description.into(),
78 parameters,
79 }
80 }
81
82 pub fn with_string_param(
97 name: impl Into<String>,
98 description: impl Into<String>,
99 param_name: impl Into<String>,
100 param_description: impl Into<String>,
101 ) -> Self {
102 let param_name = param_name.into();
103 let param_description = param_description.into();
104
105 let mut properties = serde_json::Map::new();
107 properties.insert("type".to_string(), json!("object"));
108
109 let mut obj_properties = serde_json::Map::new();
110 obj_properties.insert(
111 param_name.clone(),
112 json!({
113 "type": "string",
114 "description": param_description
115 }),
116 );
117 properties.insert(
118 "properties".to_string(),
119 serde_json::Value::Object(obj_properties),
120 );
121
122 let required_arr =
123 serde_json::Value::Array(vec![serde_json::Value::String(param_name.clone())]);
124 properties.insert("required".to_string(), required_arr);
125
126 let params = serde_json::Value::Object(properties);
127 Self::new(name, description, params)
128 }
129
130 pub fn validate(&self, args: &JsonValue) -> Result<JsonValue, ToolValidationError> {
146 validate_args_internal(&self.parameters, args)
147 }
148
149 pub fn requires_parameters(&self) -> bool {
151 self.parameters
152 .get("required")
153 .and_then(|r| r.as_array())
154 .map(|arr| !arr.is_empty())
155 .unwrap_or(false)
156 }
157}
158
159#[derive(Error, Debug)]
161pub enum ToolValidationError {
162 #[error("Invalid JSON: {0}")]
163 InvalidJson(#[from] serde_json::Error),
165
166 #[error("Schema validation failed: {0}")]
167 SchemaValidation(String),
169
170 #[error("Missing required field: {0}")]
171 MissingRequiredField(String),
173}
174
175pub fn validate_args(tool: &Tool, args: &JsonValue) -> Result<JsonValue, ToolValidationError> {
177 validate_args_internal(&tool.parameters, args)
178}
179
180fn validate_args_internal(
182 schema: &JsonValue,
183 args: &JsonValue,
184) -> Result<JsonValue, ToolValidationError> {
185 let validator =
186 Validator::new(schema).map_err(|e| ToolValidationError::SchemaValidation(e.to_string()))?;
187
188 let validation_result = validator.validate(args);
189
190 match validation_result {
191 Ok(()) => Ok(args.clone()),
192 Err(errors) => {
193 let error_msg = format!("{}", errors);
195 Err(ToolValidationError::SchemaValidation(error_msg))
196 }
197 }
198}
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203
204 #[test]
205 fn test_tool_validation() {
206 let tool = Tool::with_string_param(
207 "get_weather",
208 "Get current weather for a location",
209 "location",
210 "City name or coordinates",
211 );
212
213 let valid_args = serde_json::json!({
214 "location": "London"
215 });
216
217 let result = tool.validate(&valid_args);
218 assert!(result.is_ok());
219 }
220
221 #[test]
222 fn test_tool_validation_failure() {
223 let tool = Tool::with_string_param(
224 "get_weather",
225 "Get current weather for a location",
226 "location",
227 "City name or coordinates",
228 );
229
230 let invalid_args = serde_json::json!({});
232
233 let result = tool.validate(&invalid_args);
234 assert!(result.is_err());
235 }
236}