Skip to main content

anthropic_tools/common/
tool.rs

1//! Tool definitions for function calling with the Anthropic API.
2//!
3//! This module provides types for defining tools that Claude can use:
4//!
5//! - [`Tool`] - Main tool definition with name, description, and input schema
6//! - [`JsonSchema`] - JSON Schema for tool input parameters
7//! - [`PropertyDef`] - Property definitions within a schema
8//! - [`CacheControl`] - Cache control for prompt caching
9//!
10//! # Example
11//!
12//! ```rust
13//! use anthropic_tools::common::tool::{Tool, PropertyDef};
14//!
15//! // Create a weather tool
16//! let mut tool = Tool::new("get_weather");
17//! tool.description("Get the current weather for a location")
18//!     .add_string_property("location", Some("City and state, e.g., San Francisco, CA"), true)
19//!     .add_enum_property("unit", Some("Temperature unit"), vec!["celsius", "fahrenheit"], false);
20//!
21//! // Convert to JSON for API request
22//! let json = tool.to_value();
23//! ```
24//!
25//! # With Prompt Caching
26//!
27//! ```rust
28//! use anthropic_tools::common::tool::Tool;
29//!
30//! let mut tool = Tool::new("expensive_tool");
31//! tool.description("A tool with many parameters")
32//!     .with_cache();  // Enable prompt caching
33//! ```
34
35use serde::{Deserialize, Serialize};
36use std::collections::HashMap;
37
38/// Tool definition for the Anthropic API
39#[derive(Serialize, Deserialize, Debug, Clone)]
40pub struct Tool {
41    /// Name of the tool
42    pub name: String,
43
44    /// Description of what the tool does
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub description: Option<String>,
47
48    /// JSON schema for the tool's input parameters
49    pub input_schema: JsonSchema,
50
51    /// Cache control for prompt caching
52    #[serde(skip_serializing_if = "Option::is_none")]
53    pub cache_control: Option<CacheControl>,
54}
55
56/// Cache control for prompt caching
57#[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/// JSON Schema for tool input
72#[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/// Property definition in JSON schema
88#[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    /// Create a new tool with name only
114    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    /// Set the tool description
124    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    /// Add a string property to the input schema
130    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    /// Add a number property to the input schema
144    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    /// Add a boolean property to the input schema
158    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    /// Add an enum property to the input schema
172    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    /// Add an array property to the input schema
190    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    /// Add a property with custom PropertyDef
205    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    /// Enable cache control for this tool
227    pub fn with_cache(&mut self) -> &mut Self {
228        self.cache_control = Some(CacheControl::ephemeral());
229        self
230    }
231
232    /// Build the tool and return ownership
233    pub fn build(self) -> Self {
234        self
235    }
236
237    /// Convert to serde_json::Value
238    pub fn to_value(&self) -> serde_json::Value {
239        serde_json::to_value(self).unwrap()
240    }
241}
242
243impl JsonSchema {
244    /// Create an object schema
245    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    /// Create an empty object schema (no properties)
255    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    /// Create a string property
267    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    /// Create a number property
280    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    /// Create an integer property
293    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    /// Create a boolean property
306    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    /// Create an enum (string with allowed values) property
319    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    /// Create an array property
332    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    /// Create an object property
345    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    /// Set a default value
358    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}