tool_useful/
schema.rs

1//! Schema generation and provider-specific formats.
2
3use serde::{Deserialize, Serialize};
4use serde_json::{json, Value};
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct ParameterSchema {
8    pub name: String,
9    pub description: Option<String>,
10    #[serde(rename = "type")]
11    pub type_name: String,
12    #[serde(default)]
13    pub required: bool,
14    #[serde(skip_serializing_if = "Option::is_none")]
15    pub default: Option<Value>,
16    #[serde(skip_serializing_if = "Option::is_none")]
17    pub enum_values: Option<Vec<Value>>,
18}
19
20impl ParameterSchema {
21    pub fn new(name: impl Into<String>, type_name: impl Into<String>) -> Self {
22        Self {
23            name: name.into(),
24            description: None,
25            type_name: type_name.into(),
26            required: false,
27            default: None,
28            enum_values: None,
29        }
30    }
31
32    pub fn with_description(mut self, description: impl Into<String>) -> Self {
33        self.description = Some(description.into());
34        self
35    }
36
37    pub fn required(mut self) -> Self {
38        self.required = true;
39        self
40    }
41
42    pub fn with_default(mut self, default: Value) -> Self {
43        self.default = Some(default);
44        self
45    }
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct ToolSchema {
50    pub name: String,
51    pub description: String,
52    pub parameters: Vec<ParameterSchema>,
53}
54
55impl ToolSchema {
56    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
57        Self {
58            name: name.into(),
59            description: description.into(),
60            parameters: Vec::new(),
61        }
62    }
63
64    pub fn with_parameter(mut self, param: ParameterSchema) -> Self {
65        self.parameters.push(param);
66        self
67    }
68
69    pub fn with_parameters<I>(mut self, params: I) -> Self
70    where
71        I: IntoIterator<Item = ParameterSchema>,
72    {
73        self.parameters.extend(params);
74        self
75    }
76}
77
78/// Provider-specific schema formats
79pub trait ProviderSchema {
80    fn to_openai_schema(&self) -> Value;
81    fn to_anthropic_schema(&self) -> Value;
82    fn to_gemini_schema(&self) -> Value;
83    fn to_json_schema(&self) -> Value;
84}
85
86impl ProviderSchema for ToolSchema {
87    fn to_openai_schema(&self) -> Value {
88        let mut properties = serde_json::Map::new();
89        let mut required = Vec::new();
90
91        for param in &self.parameters {
92            let mut param_schema = serde_json::Map::new();
93            param_schema.insert("type".to_string(), json!(param.type_name));
94            if let Some(desc) = &param.description {
95                param_schema.insert("description".to_string(), json!(desc));
96            }
97            if let Some(enum_vals) = &param.enum_values {
98                param_schema.insert("enum".to_string(), json!(enum_vals));
99            }
100            properties.insert(param.name.clone(), Value::Object(param_schema));
101            if param.required {
102                required.push(param.name.clone());
103            }
104        }
105
106        json!({
107            "type": "function",
108            "function": {
109                "name": self.name,
110                "description": self.description,
111                "parameters": {
112                    "type": "object",
113                    "properties": properties,
114                    "required": required,
115                }
116            }
117        })
118    }
119
120    fn to_anthropic_schema(&self) -> Value {
121        let mut properties = serde_json::Map::new();
122        let mut required = Vec::new();
123
124        for param in &self.parameters {
125            let mut param_schema = serde_json::Map::new();
126            param_schema.insert("type".to_string(), json!(param.type_name));
127            if let Some(desc) = &param.description {
128                param_schema.insert("description".to_string(), json!(desc));
129            }
130            properties.insert(param.name.clone(), Value::Object(param_schema));
131            if param.required {
132                required.push(param.name.clone());
133            }
134        }
135
136        json!({
137            "name": self.name,
138            "description": self.description,
139            "input_schema": {
140                "type": "object",
141                "properties": properties,
142                "required": required,
143            }
144        })
145    }
146
147    fn to_gemini_schema(&self) -> Value {
148        let mut properties = serde_json::Map::new();
149        let mut required = Vec::new();
150
151        for param in &self.parameters {
152            let mut param_schema = serde_json::Map::new();
153            param_schema.insert("type".to_string(), json!(param.type_name));
154            if let Some(desc) = &param.description {
155                param_schema.insert("description".to_string(), json!(desc));
156            }
157            properties.insert(param.name.clone(), Value::Object(param_schema));
158            if param.required {
159                required.push(param.name.clone());
160            }
161        }
162
163        json!({
164            "name": self.name,
165            "description": self.description,
166            "parameters": {
167                "type": "object",
168                "properties": properties,
169                "required": required,
170            }
171        })
172    }
173
174    fn to_json_schema(&self) -> Value {
175        let mut properties = serde_json::Map::new();
176        let mut required = Vec::new();
177
178        for param in &self.parameters {
179            let mut param_schema = serde_json::Map::new();
180            param_schema.insert("type".to_string(), json!(param.type_name));
181            if let Some(desc) = &param.description {
182                param_schema.insert("description".to_string(), json!(desc));
183            }
184            if let Some(default) = &param.default {
185                param_schema.insert("default".to_string(), default.clone());
186            }
187            properties.insert(param.name.clone(), Value::Object(param_schema));
188            if param.required {
189                required.push(param.name.clone());
190            }
191        }
192
193        json!({
194            "type": "object",
195            "properties": properties,
196            "required": required,
197        })
198    }
199}