openai_ng/proto/
tool.rs

1use crate::error::*;
2use serde_with::skip_serializing_none;
3use smart_default::SmartDefault;
4use std::collections::HashMap;
5
6#[skip_serializing_none]
7#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
8pub struct ToolCall {
9    pub id: Option<String>,
10    #[serde(rename = "type")]
11    pub typ: Option<String>,
12    pub function: Function,
13}
14
15impl From<Function> for ToolCall {
16    fn from(f: Function) -> Self {
17        ToolCall {
18            id: None,
19            typ: Some("function".to_string()),
20            function: f,
21        }
22    }
23}
24
25impl ToolCall {
26    pub fn builder() -> ToolCallBuilder {
27        ToolCallBuilder::default()
28    }
29}
30
31#[derive(Debug, Clone, SmartDefault)]
32pub struct ToolCallBuilder {
33    pub id: Option<String>,
34    #[default(Some("function".to_string()))]
35    typ: Option<String>,
36    function: Option<Function>,
37}
38
39impl ToolCallBuilder {
40    pub fn with_function(mut self, function: impl Into<Function>) -> Self {
41        self.function = Some(function.into());
42        self
43    }
44
45    pub fn build(self) -> Result<ToolCall> {
46        let Self { id, typ, function } = self;
47        let typ = typ.ok_or(Error::ToolCallBuild)?;
48        let function = function.ok_or(Error::ToolCallBuild)?;
49        Ok(ToolCall {
50            id,
51            typ: Some(typ),
52            function,
53        })
54    }
55}
56
57#[skip_serializing_none]
58#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
59pub struct Function {
60    pub name: Option<String>,
61    pub description: Option<String>,
62    pub parameters: Option<Parameters>,
63    pub arguments: Option<String>,
64}
65
66pub mod serde_value {
67
68    use serde::de::{self, Deserialize, DeserializeOwned, Deserializer};
69    use serde::ser::{self, Serialize, Serializer};
70    use serde_json;
71
72    pub fn serialize<T, S>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
73    where
74        T: Serialize,
75        S: Serializer,
76    {
77        let j = serde_json::to_string(value).map_err(ser::Error::custom)?;
78        j.serialize(serializer)
79    }
80
81    pub fn deserialize<'de, T, D>(deserializer: D) -> Result<T, D::Error>
82    where
83        T: DeserializeOwned,
84        D: Deserializer<'de>,
85    {
86        let j = String::deserialize(deserializer)?;
87        serde_json::from_str(&j).map_err(de::Error::custom)
88    }
89}
90
91impl Function {
92    pub fn builder() -> FunctionBuilder {
93        FunctionBuilder::default()
94    }
95}
96
97#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, SmartDefault)]
98pub struct FunctionBuilder {
99    pub name: Option<String>,
100    pub description: Option<String>,
101    pub parameters: Option<Parameters>,
102    pub arguments: Option<String>,
103}
104
105impl FunctionBuilder {
106    pub fn with_name(mut self, name: impl Into<String>) -> Self {
107        self.name = Some(name.into());
108        self
109    }
110
111    pub fn with_description(mut self, description: impl Into<String>) -> Self {
112        self.description = Some(description.into());
113        self
114    }
115
116    pub fn with_parameters(mut self, parameters: Parameters) -> Self {
117        self.parameters = Some(parameters);
118        self
119    }
120
121    pub fn build(self) -> Result<Function> {
122        let Self {
123            name,
124            description,
125            parameters,
126            arguments,
127        } = self;
128
129        let name = name.ok_or(Error::ToolCallFunctionBuild)?;
130
131        Ok(Function {
132            name: Some(name),
133            description,
134            parameters,
135            arguments,
136        })
137    }
138}
139
140#[skip_serializing_none]
141#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
142#[allow(non_camel_case_types)]
143#[serde(untagged)]
144pub enum Argument {
145    string(String),
146    number(f64),
147    integer(i64),
148    boolean(bool),
149    array(Vec<Argument>),
150    object(HashMap<String, Argument>),
151}
152
153macro_rules! impl_argument_as_value {
154    ($fun: ident, $i: ident, $typ: ty) => {
155        impl Argument {
156            pub fn $fun(&self) -> Option<&$typ> {
157                match self {
158                    Self::$i(inner) => Some(&inner),
159                    _ => None,
160                }
161            }
162        }
163    };
164}
165
166impl_argument_as_value!(as_string, string, String);
167impl_argument_as_value!(as_number, number, f64);
168impl_argument_as_value!(as_integer, integer, i64);
169impl_argument_as_value!(as_boolean, boolean, bool);
170impl_argument_as_value!(as_array, array, Vec<Argument>);
171impl_argument_as_value!(as_object, object, HashMap<String, Argument>);
172
173#[skip_serializing_none]
174#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
175pub struct Parameters {
176    #[serde(rename = "type")]
177    pub typ: String,
178    pub properties: HashMap<String, ParameterProperty>,
179    #[serde(default)]
180    pub required: Vec<String>,
181}
182
183impl Parameters {
184    pub fn builder() -> ParametersBuilder {
185        ParametersBuilder::default()
186    }
187}
188
189#[derive(Debug, Clone, SmartDefault)]
190pub struct ParametersBuilder {
191    #[default(Some("object".to_string()))]
192    typ: Option<String>,
193    properties: HashMap<String, ParameterProperty>,
194    required: Vec<String>,
195}
196
197impl ParametersBuilder {
198    pub fn add_property(
199        mut self,
200        name: impl Into<String>,
201        property: ParameterProperty,
202        required: bool,
203    ) -> Self {
204        let name = name.into();
205        self.properties.insert(name.clone(), property);
206        if required {
207            self.required.push(name);
208        }
209        self
210    }
211
212    pub fn build(self) -> Result<Parameters> {
213        let Self {
214            typ,
215            properties,
216            required,
217        } = self;
218
219        let typ = typ.ok_or(Error::ToolCallParametersBuild)?;
220
221        Ok(Parameters {
222            typ,
223            properties,
224            required,
225        })
226    }
227}
228
229#[skip_serializing_none]
230#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
231pub struct ParameterProperty {
232    #[serde(rename = "type")]
233    pub typ: Option<ParameterType>,
234    pub description: String,
235    pub items: Option<HashMap<String, String>>,
236}
237
238impl ParameterProperty {
239    pub fn builder() -> ParameterPropertyBuilder {
240        ParameterPropertyBuilder::default()
241    }
242}
243
244#[derive(Debug, Clone, SmartDefault)]
245pub struct ParameterPropertyBuilder {
246    typ: Option<ParameterType>,
247    description: Option<String>,
248    items: Option<HashMap<String, String>>,
249}
250
251impl ParameterPropertyBuilder {
252    pub fn with_type(mut self, typ: ParameterType) -> Self {
253        self.typ = Some(typ);
254        self
255    }
256
257    pub fn with_description(mut self, description: impl Into<String>) -> Self {
258        self.description = Some(description.into());
259        self
260    }
261
262    pub fn with_items(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
263        if self.items.is_none() {
264            self.items = Some(HashMap::new());
265        }
266
267        self.items
268            .as_mut()
269            .unwrap()
270            .insert(key.into(), value.into());
271        self
272    }
273
274    pub fn build(self) -> Result<ParameterProperty> {
275        let Self {
276            typ,
277            description,
278            items,
279        } = self;
280
281        let typ = typ.ok_or(Error::ToolCallParametersBuild)?;
282        let description = description.ok_or(Error::ToolCallParametersBuild)?;
283
284        Ok(ParameterProperty {
285            typ: Some(typ),
286            description,
287            items,
288        })
289    }
290}
291
292#[skip_serializing_none]
293#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
294#[allow(non_camel_case_types)]
295pub enum ParameterType {
296    string,
297    number,
298    integer,
299    boolean,
300    array,
301    object,
302}