gemini_rust/
tools.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4/// Tool that can be used by the model
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct Tool {
7    /// The function declaration for the tool
8    pub function_declarations: Vec<FunctionDeclaration>,
9}
10
11impl Tool {
12    /// Create a new tool with a single function declaration
13    pub fn new(function_declaration: FunctionDeclaration) -> Self {
14        Self {
15            function_declarations: vec![function_declaration],
16        }
17    }
18
19    /// Create a new tool with multiple function declarations
20    pub fn with_functions(function_declarations: Vec<FunctionDeclaration>) -> Self {
21        Self {
22            function_declarations,
23        }
24    }
25}
26
27/// Declaration of a function that can be called by the model
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct FunctionDeclaration {
30    /// The name of the function
31    pub name: String,
32    /// The description of the function
33    pub description: String,
34    /// The parameters for the function
35    pub parameters: FunctionParameters,
36}
37
38impl FunctionDeclaration {
39    /// Create a new function declaration
40    pub fn new(
41        name: impl Into<String>,
42        description: impl Into<String>,
43        parameters: FunctionParameters,
44    ) -> Self {
45        Self {
46            name: name.into(),
47            description: description.into(),
48            parameters,
49        }
50    }
51}
52
53/// Parameters for a function
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct FunctionParameters {
56    /// The type of the parameters
57    #[serde(rename = "type")]
58    pub param_type: String,
59    /// The properties of the parameters
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub properties: Option<HashMap<String, PropertyDetails>>,
62    /// The required properties
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub required: Option<Vec<String>>,
65}
66
67impl FunctionParameters {
68    /// Create a new object parameter set
69    pub fn object() -> Self {
70        Self {
71            param_type: "OBJECT".to_string(),
72            properties: Some(HashMap::new()),
73            required: Some(Vec::new()),
74        }
75    }
76
77    /// Add a property to the parameters
78    pub fn with_property(
79        mut self,
80        name: impl Into<String>,
81        details: PropertyDetails,
82        required: bool,
83    ) -> Self {
84        let name = name.into();
85        if let Some(props) = &mut self.properties {
86            props.insert(name.clone(), details);
87        }
88        if required {
89            if let Some(req) = &mut self.required {
90                req.push(name);
91            }
92        }
93        self
94    }
95}
96
97/// Details about a property
98#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct PropertyDetails {
100    /// The type of the property
101    #[serde(rename = "type")]
102    pub property_type: String,
103    /// The description of the property
104    pub description: String,
105    /// The enum values if the property is an enum
106    #[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
107    pub enum_values: Option<Vec<String>>,
108    /// The items if the property is an array
109    #[serde(skip_serializing_if = "Option::is_none")]
110    pub items: Option<Box<PropertyDetails>>,
111}
112
113impl PropertyDetails {
114    /// Create a new string property
115    pub fn string(description: impl Into<String>) -> Self {
116        Self {
117            property_type: "STRING".to_string(),
118            description: description.into(),
119            enum_values: None,
120            items: None,
121        }
122    }
123
124    /// Create a new number property
125    pub fn number(description: impl Into<String>) -> Self {
126        Self {
127            property_type: "NUMBER".to_string(),
128            description: description.into(),
129            enum_values: None,
130            items: None,
131        }
132    }
133
134    /// Create a new integer property
135    pub fn integer(description: impl Into<String>) -> Self {
136        Self {
137            property_type: "INTEGER".to_string(),
138            description: description.into(),
139            enum_values: None,
140            items: None,
141        }
142    }
143
144    /// Create a new boolean property
145    pub fn boolean(description: impl Into<String>) -> Self {
146        Self {
147            property_type: "BOOLEAN".to_string(),
148            description: description.into(),
149            enum_values: None,
150            items: None,
151        }
152    }
153
154    /// Create a new array property
155    pub fn array(description: impl Into<String>, items: PropertyDetails) -> Self {
156        Self {
157            property_type: "ARRAY".to_string(),
158            description: description.into(),
159            enum_values: None,
160            items: Some(Box::new(items)),
161        }
162    }
163
164    /// Create a new enum property
165    pub fn enum_type(
166        description: impl Into<String>,
167        enum_values: impl IntoIterator<Item = impl Into<String>>,
168    ) -> Self {
169        Self {
170            property_type: "STRING".to_string(),
171            description: description.into(),
172            enum_values: Some(enum_values.into_iter().map(|s| s.into()).collect()),
173            items: None,
174        }
175    }
176}
177
178/// A function call made by the model
179#[derive(Debug, Clone, Serialize, Deserialize)]
180pub struct FunctionCall {
181    /// The name of the function
182    pub name: String,
183    /// The arguments for the function
184    pub args: serde_json::Value,
185}
186
187impl FunctionCall {
188    /// Create a new function call
189    pub fn new(name: impl Into<String>, args: serde_json::Value) -> Self {
190        Self {
191            name: name.into(),
192            args,
193        }
194    }
195
196    /// Get a parameter from the arguments
197    pub fn get<T: serde::de::DeserializeOwned>(&self, key: &str) -> crate::Result<T> {
198        match &self.args {
199            serde_json::Value::Object(obj) => {
200                if let Some(value) = obj.get(key) {
201                    serde_json::from_value(value.clone())
202                        .map_err(|e| crate::Error::FunctionCallError(format!("Error deserializing parameter {}: {}", key, e)))
203                } else {
204                    Err(crate::Error::FunctionCallError(format!("Missing parameter: {}", key)))
205                }
206            }
207            _ => Err(crate::Error::FunctionCallError("Arguments are not an object".to_string())),
208        }
209    }
210}
211
212/// A response from a function
213#[derive(Debug, Clone, Serialize, Deserialize)]
214pub struct FunctionResponse {
215    /// The name of the function
216    pub name: String,
217    /// The response from the function
218    /// This must be a valid JSON object
219    #[serde(skip_serializing_if = "Option::is_none")]
220    pub response: Option<serde_json::Value>,
221}
222
223impl FunctionResponse {
224    /// Create a new function response with a JSON value
225    pub fn new(name: impl Into<String>, response: serde_json::Value) -> Self {
226        Self {
227            name: name.into(),
228            response: Some(response),
229        }
230    }
231    
232    /// Create a new function response with a string that will be parsed as JSON
233    pub fn from_str(name: impl Into<String>, response: impl Into<String>) -> Result<Self, serde_json::Error> {
234        let json = serde_json::from_str(&response.into())?;
235        Ok(Self {
236            name: name.into(),
237            response: Some(json),
238        })
239    }
240}