magi_tool/
lib.rs

1use schemars::{
2    JsonSchema,
3    generate::SchemaSettings,
4    transform::{self, Transform},
5};
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::borrow::Cow;
9use thiserror::Error;
10
11// Use ToolKit for local in-process tool closures; use PegBoard for MCP service tools.
12#[cfg(feature = "tool-kit")]
13mod tool_kit;
14#[cfg(feature = "tool-kit")]
15pub use tool_kit::{ToolCallError, ToolKit};
16
17#[cfg(feature = "pegboard")]
18mod pegboard;
19#[cfg(feature = "pegboard")]
20pub use pegboard::{PegBoard, PegBoardError, prefix_tool_name};
21
22// Re-export commonly used rmcp types for convenience
23#[cfg(feature = "pegboard")]
24pub use rmcp::model::CallToolResult;
25
26/// Error type for tool creation and schema parsing
27#[derive(Debug, Error)]
28pub enum ToolError {
29    #[error("Failed to serialize JSON schema: {0}")]
30    SchemaSerialization(#[from] serde_json::Error),
31    #[error("Schema type field is required")]
32    MissingSchemaType,
33    #[error("Schema must be of type 'object', got: {0}")]
34    InvalidSchemaType(String),
35}
36
37/// Tool schema definition for function calling.
38/// Compatible with Claude's tool format.
39#[derive(Debug, Deserialize, Clone, Default, PartialEq, Serialize)]
40pub struct Tool {
41    pub name: Cow<'static, str>,
42    pub description: Option<Cow<'static, str>>,
43    pub input_schema: serde_json::Value,
44}
45
46#[derive(Debug, Clone)]
47#[non_exhaustive]
48pub struct AddNullable {
49    /// When set to `true` (the default), `"null"` will also be removed from the schemas `type`.
50    pub remove_null_type: bool,
51}
52
53impl Default for AddNullable {
54    fn default() -> Self {
55        Self {
56            remove_null_type: true,
57        }
58    }
59}
60
61impl AddNullable {
62    fn has_type(schema: &schemars::Schema, ty: &str) -> bool {
63        match schema.get("type") {
64            Some(Value::Array(values)) => values.iter().any(|v| v.as_str() == Some(ty)),
65            Some(Value::String(s)) => s == ty,
66            _ => false,
67        }
68    }
69}
70
71impl Transform for AddNullable {
72    fn transform(&mut self, schema: &mut schemars::Schema) {
73        if Self::has_type(schema, "null") {
74            // Don't add nullable property, just handle the null type removal
75            if let Some(ty) = schema.get_mut("type")
76                && self.remove_null_type
77            {
78                // Remove null from type array and clean up enum if present
79                if let Value::Array(array) = ty {
80                    array.retain(|t| t.as_str() != Some("null"));
81                    if array.len() == 1 {
82                        *ty = array[0].clone();
83                    }
84                }
85
86                // Also clean up enum arrays that contain null
87                if let Some(Value::Array(enum_array)) = schema.get_mut("enum") {
88                    enum_array.retain(|v| !v.is_null());
89                }
90            }
91        }
92
93        transform::transform_subschemas(self, schema);
94    }
95}
96
97pub fn get_tool<T: JsonSchema, S1, S2>(name: S1, desc: Option<S2>) -> Result<Tool, ToolError>
98where
99    S1: Into<Cow<'static, str>>,
100    S2: Into<Cow<'static, str>>,
101{
102    let json_value = parse_input_schema::<T>()?;
103    Ok(Tool {
104        name: name.into(),
105        description: desc.map(Into::into),
106        input_schema: json_value,
107    })
108}
109
110#[macro_export]
111macro_rules! define_tool {
112    ($tool_name:ident, $function_name:expr, $description:expr, $param_type:ty) => {
113        paste::paste! {
114            static [<$tool_name _ONCE_LOCK>]: std::sync::OnceLock<Result<$crate::Tool, $crate::ToolError>> = ::std::sync::OnceLock::new();
115
116            pub fn [<get_ $tool_name:lower>]() -> Result<&'static $crate::Tool, &'static $crate::ToolError> {
117                [<$tool_name _ONCE_LOCK>].get_or_init(|| {
118                    $crate::get_tool::<$param_type, _, _>(
119                        $function_name,
120                        Some($description),
121                    )
122                }).as_ref()
123            }
124        }
125    };
126}
127
128pub fn parse_input_schema<T: JsonSchema>() -> Result<serde_json::Value, ToolError> {
129    let settings = SchemaSettings::draft2019_09()
130        .with(|s| {
131            // s.option_nullable = false;
132            // s.option_add_null_type = false;
133            s.inline_subschemas = true;
134        })
135        .with_transform(AddNullable::default());
136    let schema = settings.into_generator().into_root_schema_for::<T>();
137    let mut json_value = serde_json::to_value(schema)?;
138    let schema_type = json_value.get("type").ok_or(ToolError::MissingSchemaType)?;
139
140    match schema_type {
141        Value::String(s) if s == "object" => {
142            // Valid object type
143        }
144        Value::String(s) => {
145            return Err(ToolError::InvalidSchemaType(s.clone()));
146        }
147        other => {
148            return Err(ToolError::InvalidSchemaType(format!("{:?}", other)));
149        }
150    }
151
152    if let Some(obj) = json_value.as_object_mut() {
153        obj.remove("$schema");
154        obj.remove("title");
155        obj.remove("definitions");
156    }
157    Ok(json_value)
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163    use schemars::JsonSchema;
164
165    #[derive(JsonSchema, serde::Deserialize)]
166    #[allow(dead_code)]
167    pub struct WeatherParams {
168        /// The city and state, e.g. San Francisco, CA
169        pub location: String,
170        pub unit: Option<UnitEnum>,
171        pub arr: Option<Vec<String>>,
172    }
173
174    #[derive(JsonSchema, serde::Deserialize, PartialEq, Debug)]
175    #[serde(rename_all = "lowercase")]
176    pub enum UnitEnum {
177        Celsius,
178        Fahrenheit,
179    }
180
181    #[test]
182    fn test_parse_input_schema() {
183        let schema = parse_input_schema::<WeatherParams>().unwrap();
184
185        let got = serde_json::to_string(&schema).unwrap();
186        let want = serde_json::json!({
187            "type": "object",
188            "properties": {
189                "location": {
190                    "type": "string",
191                    "description": "The city and state, e.g. San Francisco, CA"
192                },
193                "unit": {
194                    "type": "string",
195                    "enum": ["celsius", "fahrenheit"]
196                },
197                "arr": {
198                    "type": "array",
199                    "items": {
200                        "type": "string"
201                    }
202                }
203            },
204            "required": ["location"]
205        })
206        .to_string();
207        assert_eq!(want, got, "Expected: {} Got: {}", want, got);
208    }
209
210    #[test]
211    fn test_get_tool() {
212        let tool = get_tool::<WeatherParams, _, _>(
213            "get_weather",
214            Some("Get the current weather in a given location"),
215        )
216        .unwrap();
217
218        assert_eq!(tool.name, "get_weather");
219        assert_eq!(
220            tool.description,
221            Some(Cow::Borrowed("Get the current weather in a given location"))
222        );
223
224        // Verify input_schema structure
225        let schema = &tool.input_schema;
226        assert_eq!(schema.get("type").and_then(|v| v.as_str()), Some("object"));
227        assert!(schema.get("properties").is_some());
228        assert!(schema.get("required").is_some());
229    }
230
231    #[test]
232    fn test_deserialize_params() {
233        let s = r#"{"location":"Boston, MA","unit":"celsius"}"#;
234        let got: WeatherParams = serde_json::from_str(s).unwrap();
235        assert_eq!(
236            "Boston, MA", got.location,
237            "Expected: Boston, MA Got: {}",
238            got.location
239        );
240        assert_eq!(
241            Some(UnitEnum::Celsius),
242            got.unit,
243            "Expected: celsius got: {:?}",
244            got.unit
245        );
246    }
247
248    #[test]
249    fn test_define_tool_macro() {
250        define_tool!(
251            WEATHER_TOOL,
252            "get_weather",
253            "Get the current weather",
254            WeatherParams
255        );
256
257        let tool = get_weather_tool().unwrap();
258        assert_eq!(tool.name, "get_weather");
259        assert_eq!(
260            tool.description,
261            Some(Cow::Borrowed("Get the current weather"))
262        );
263        assert!(tool.input_schema.get("properties").is_some());
264        assert!(tool.input_schema.get("required").is_some());
265    }
266
267    #[test]
268    fn test_nullable_handling() {
269        #[derive(JsonSchema)]
270        #[allow(dead_code)]
271        struct TestNullable {
272            required_field: String,
273            optional_field: Option<String>,
274        }
275
276        let schema = parse_input_schema::<TestNullable>().unwrap();
277
278        // Verify that optional fields don't have "null" in their type
279        let props = schema.get("properties").unwrap().as_object().unwrap();
280        let optional = props.get("optional_field").unwrap();
281
282        // Should not have null type due to AddNullable transform
283        if let Some(type_val) = optional.get("type") {
284            assert_ne!(type_val.as_str(), Some("null"));
285        }
286
287        // Only required_field should be in required array
288        let required = schema.get("required").unwrap().as_array().unwrap();
289        assert_eq!(required.len(), 1);
290        assert_eq!(required[0].as_str(), Some("required_field"));
291    }
292}
magi_tool/lib.rs

magi_tool/
lib.rs
