1use serde::{Deserialize, Serialize};
36use std::collections::HashMap;
37
38#[derive(Serialize, Deserialize, Debug, Clone)]
40pub struct Tool {
41 pub name: String,
43
44 #[serde(skip_serializing_if = "Option::is_none")]
46 pub description: Option<String>,
47
48 pub input_schema: JsonSchema,
50
51 #[serde(skip_serializing_if = "Option::is_none")]
53 pub cache_control: Option<CacheControl>,
54}
55
56#[derive(Serialize, Deserialize, Debug, Clone)]
58pub struct CacheControl {
59 #[serde(rename = "type")]
60 pub type_name: String,
61}
62
63impl CacheControl {
64 pub fn ephemeral() -> Self {
65 CacheControl {
66 type_name: "ephemeral".to_string(),
67 }
68 }
69}
70
71#[derive(Serialize, Deserialize, Debug, Clone)]
73pub struct JsonSchema {
74 #[serde(rename = "type")]
75 pub type_name: String,
76
77 #[serde(skip_serializing_if = "Option::is_none")]
78 pub properties: Option<HashMap<String, PropertyDef>>,
79
80 #[serde(skip_serializing_if = "Option::is_none")]
81 pub required: Option<Vec<String>>,
82
83 #[serde(skip_serializing_if = "Option::is_none")]
84 pub additional_properties: Option<bool>,
85}
86
87#[derive(Serialize, Deserialize, Debug, Clone)]
89pub struct PropertyDef {
90 #[serde(rename = "type")]
91 pub type_name: String,
92
93 #[serde(skip_serializing_if = "Option::is_none")]
94 pub description: Option<String>,
95
96 #[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
97 pub enum_values: Option<Vec<String>>,
98
99 #[serde(skip_serializing_if = "Option::is_none")]
100 pub items: Option<Box<PropertyDef>>,
101
102 #[serde(skip_serializing_if = "Option::is_none")]
103 pub properties: Option<HashMap<String, PropertyDef>>,
104
105 #[serde(skip_serializing_if = "Option::is_none")]
106 pub required: Option<Vec<String>>,
107
108 #[serde(rename = "default", skip_serializing_if = "Option::is_none")]
109 pub default_value: Option<serde_json::Value>,
110}
111
112impl Tool {
113 pub fn new<S: AsRef<str>>(name: S) -> Self {
115 Tool {
116 name: name.as_ref().to_string(),
117 description: None,
118 input_schema: JsonSchema::object(),
119 cache_control: None,
120 }
121 }
122
123 pub fn description<S: AsRef<str>>(&mut self, desc: S) -> &mut Self {
125 self.description = Some(desc.as_ref().to_string());
126 self
127 }
128
129 pub fn add_string_property<S: AsRef<str>>(
131 &mut self,
132 name: S,
133 description: Option<S>,
134 required: bool,
135 ) -> &mut Self {
136 self.add_property(
137 name.as_ref(),
138 PropertyDef::string(description.map(|s| s.as_ref().to_string())),
139 required,
140 )
141 }
142
143 pub fn add_number_property<S: AsRef<str>>(
145 &mut self,
146 name: S,
147 description: Option<S>,
148 required: bool,
149 ) -> &mut Self {
150 self.add_property(
151 name.as_ref(),
152 PropertyDef::number(description.map(|s| s.as_ref().to_string())),
153 required,
154 )
155 }
156
157 pub fn add_boolean_property<S: AsRef<str>>(
159 &mut self,
160 name: S,
161 description: Option<S>,
162 required: bool,
163 ) -> &mut Self {
164 self.add_property(
165 name.as_ref(),
166 PropertyDef::boolean(description.map(|s| s.as_ref().to_string())),
167 required,
168 )
169 }
170
171 pub fn add_enum_property<S: AsRef<str>>(
173 &mut self,
174 name: S,
175 description: Option<S>,
176 values: Vec<S>,
177 required: bool,
178 ) -> &mut Self {
179 self.add_property(
180 name.as_ref(),
181 PropertyDef::enum_type(
182 description.map(|s| s.as_ref().to_string()),
183 values.into_iter().map(|s| s.as_ref().to_string()).collect(),
184 ),
185 required,
186 )
187 }
188
189 pub fn add_array_property<S: AsRef<str>>(
191 &mut self,
192 name: S,
193 description: Option<S>,
194 items: PropertyDef,
195 required: bool,
196 ) -> &mut Self {
197 self.add_property(
198 name.as_ref(),
199 PropertyDef::array(description.map(|s| s.as_ref().to_string()), items),
200 required,
201 )
202 }
203
204 fn add_property(&mut self, name: &str, prop: PropertyDef, required: bool) -> &mut Self {
206 if self.input_schema.properties.is_none() {
207 self.input_schema.properties = Some(HashMap::new());
208 }
209
210 if let Some(props) = &mut self.input_schema.properties {
211 props.insert(name.to_string(), prop);
212 }
213
214 if required {
215 if self.input_schema.required.is_none() {
216 self.input_schema.required = Some(Vec::new());
217 }
218 if let Some(req) = &mut self.input_schema.required {
219 req.push(name.to_string());
220 }
221 }
222
223 self
224 }
225
226 pub fn with_cache(&mut self) -> &mut Self {
228 self.cache_control = Some(CacheControl::ephemeral());
229 self
230 }
231
232 pub fn build(self) -> Self {
234 self
235 }
236
237 pub fn to_value(&self) -> serde_json::Value {
239 serde_json::to_value(self).unwrap()
240 }
241}
242
243impl JsonSchema {
244 pub fn object() -> Self {
246 JsonSchema {
247 type_name: "object".to_string(),
248 properties: Some(HashMap::new()),
249 required: None,
250 additional_properties: None,
251 }
252 }
253
254 pub fn empty_object() -> Self {
256 JsonSchema {
257 type_name: "object".to_string(),
258 properties: None,
259 required: None,
260 additional_properties: None,
261 }
262 }
263}
264
265impl PropertyDef {
266 pub fn string(description: Option<String>) -> Self {
268 PropertyDef {
269 type_name: "string".to_string(),
270 description,
271 enum_values: None,
272 items: None,
273 properties: None,
274 required: None,
275 default_value: None,
276 }
277 }
278
279 pub fn number(description: Option<String>) -> Self {
281 PropertyDef {
282 type_name: "number".to_string(),
283 description,
284 enum_values: None,
285 items: None,
286 properties: None,
287 required: None,
288 default_value: None,
289 }
290 }
291
292 pub fn integer(description: Option<String>) -> Self {
294 PropertyDef {
295 type_name: "integer".to_string(),
296 description,
297 enum_values: None,
298 items: None,
299 properties: None,
300 required: None,
301 default_value: None,
302 }
303 }
304
305 pub fn boolean(description: Option<String>) -> Self {
307 PropertyDef {
308 type_name: "boolean".to_string(),
309 description,
310 enum_values: None,
311 items: None,
312 properties: None,
313 required: None,
314 default_value: None,
315 }
316 }
317
318 pub fn enum_type(description: Option<String>, values: Vec<String>) -> Self {
320 PropertyDef {
321 type_name: "string".to_string(),
322 description,
323 enum_values: Some(values),
324 items: None,
325 properties: None,
326 required: None,
327 default_value: None,
328 }
329 }
330
331 pub fn array(description: Option<String>, items: PropertyDef) -> Self {
333 PropertyDef {
334 type_name: "array".to_string(),
335 description,
336 enum_values: None,
337 items: Some(Box::new(items)),
338 properties: None,
339 required: None,
340 default_value: None,
341 }
342 }
343
344 pub fn object(description: Option<String>, properties: HashMap<String, PropertyDef>) -> Self {
346 PropertyDef {
347 type_name: "object".to_string(),
348 description,
349 enum_values: None,
350 items: None,
351 properties: Some(properties),
352 required: None,
353 default_value: None,
354 }
355 }
356
357 pub fn with_default(&mut self, value: serde_json::Value) -> &mut Self {
359 self.default_value = Some(value);
360 self
361 }
362}
363
364#[cfg(test)]
365mod tests {
366 use super::*;
367
368 #[test]
369 fn test_tool_builder() {
370 let mut tool = Tool::new("search");
371 tool.description("Search the web for information")
372 .add_string_property("query", Some("The search query"), true)
373 .add_number_property("limit", Some("Maximum results to return"), false);
374
375 assert_eq!(tool.name, "search");
376 assert!(tool.description.is_some());
377 assert!(tool.input_schema.properties.is_some());
378
379 let props = tool.input_schema.properties.as_ref().unwrap();
380 assert!(props.contains_key("query"));
381 assert!(props.contains_key("limit"));
382
383 let required = tool.input_schema.required.as_ref().unwrap();
384 assert!(required.contains(&"query".to_string()));
385 assert!(!required.contains(&"limit".to_string()));
386 }
387
388 #[test]
389 fn test_tool_serialize() {
390 let mut tool = Tool::new("get_weather");
391 tool.description("Get the current weather in a given location")
392 .add_string_property(
393 "location",
394 Some("The city and state, e.g. San Francisco, CA"),
395 true,
396 )
397 .add_enum_property(
398 "unit",
399 Some("Temperature unit"),
400 vec!["celsius", "fahrenheit"],
401 false,
402 );
403
404 let json = serde_json::to_string_pretty(&tool).unwrap();
405 assert!(json.contains("\"name\": \"get_weather\""));
406 assert!(json.contains("\"type\": \"object\""));
407 assert!(json.contains("\"location\""));
408 assert!(json.contains("\"required\""));
409 }
410
411 #[test]
412 fn test_property_def_string() {
413 let prop = PropertyDef::string(Some("A test property".to_string()));
414 assert_eq!(prop.type_name, "string");
415 assert_eq!(prop.description, Some("A test property".to_string()));
416 }
417
418 #[test]
419 fn test_property_def_enum() {
420 let prop = PropertyDef::enum_type(
421 Some("Choose a color".to_string()),
422 vec!["red".to_string(), "green".to_string(), "blue".to_string()],
423 );
424 assert_eq!(prop.type_name, "string");
425 assert!(prop.enum_values.is_some());
426 assert_eq!(prop.enum_values.unwrap().len(), 3);
427 }
428
429 #[test]
430 fn test_property_def_array() {
431 let items = PropertyDef::string(Some("Item in array".to_string()));
432 let prop = PropertyDef::array(Some("An array of strings".to_string()), items);
433 assert_eq!(prop.type_name, "array");
434 assert!(prop.items.is_some());
435 }
436
437 #[test]
438 fn test_tool_with_cache() {
439 let mut tool = Tool::new("cached_tool");
440 tool.with_cache();
441 assert!(tool.cache_control.is_some());
442 }
443
444 #[test]
445 fn test_tool_to_value() {
446 let mut tool = Tool::new("test");
447 tool.description("A test tool");
448 let value = tool.to_value();
449 assert!(value.is_object());
450 assert_eq!(value["name"], "test");
451 }
452}