openai_tools/
json_schema.rs

1use fxhash::FxHashMap;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Deserialize, Serialize)]
5pub struct ItemType {
6    #[serde(rename = "type")]
7    pub type_name: String,
8    #[serde(skip_serializing_if = "Option::is_none")]
9    pub description: Option<String>,
10    #[serde(skip_serializing_if = "Option::is_none")]
11    pub items: Option<Box<JsonItem>>,
12}
13
14impl ItemType {
15    pub fn new(type_name: String, description: Option<String>) -> Self {
16        Self {
17            type_name: type_name.to_string(),
18            description: description,
19            items: None,
20        }
21    }
22
23    pub fn clone(&self) -> Self {
24        let mut items: JsonItem = JsonItem::default();
25        if let Some(item) = &self.items {
26            let mut _properties: FxHashMap<String, ItemType> = FxHashMap::default();
27            for (key, value) in item.properties.iter() {
28                _properties.insert(key.clone(), value.clone());
29            }
30            items.type_name = item.type_name.clone();
31            items.properties = _properties;
32            items.required = item.required.clone();
33            items.additional_properties = item.additional_properties;
34        }
35
36        Self {
37            type_name: self.type_name.clone(),
38            description: self.description.clone(),
39            items: if self.items.is_some() {
40                Option::from(Box::new(items))
41            } else {
42                None
43            },
44        }
45    }
46}
47
48#[derive(Debug, Clone, Deserialize, Serialize)]
49pub struct JsonItem {
50    #[serde(rename = "type")]
51    pub type_name: String,
52    pub properties: FxHashMap<String, ItemType>,
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub required: Option<Vec<String>>,
55    #[serde(rename = "additionalProperties")]
56    pub additional_properties: bool,
57}
58
59impl JsonItem {
60    pub fn new(type_name: &str, properties: FxHashMap<String, ItemType>) -> Self {
61        let mut required = Vec::new();
62        for key in properties.keys() {
63            required.push(key.clone());
64        }
65        Self {
66            type_name: type_name.to_string(),
67            properties,
68            required: if required.is_empty() {
69                None
70            } else {
71                Option::from(required)
72            },
73            additional_properties: false,
74        }
75    }
76
77    pub fn default() -> Self {
78        Self {
79            type_name: "object".to_string(),
80            properties: FxHashMap::default(),
81            required: None,
82            additional_properties: false,
83        }
84    }
85
86    pub fn add_property(&mut self, prop_name: String, item: ItemType) {
87        self.properties.insert(prop_name.to_string(), item.clone());
88        if self.required.is_none() {
89            self.required = Option::from(vec![prop_name.to_string()]);
90        } else {
91            let mut required = self.required.clone().unwrap();
92            required.push(prop_name.to_string());
93            self.required = Option::from(required);
94        }
95    }
96
97    pub fn add_array(&mut self, prop_name: String, items: JsonItem) {
98        let mut prop = ItemType::new(String::from("array"), None);
99        prop.items = Option::from(Box::new(items));
100        self.properties.insert(prop_name.to_string(), prop);
101        self.required = Option::from(vec![prop_name.to_string()]);
102    }
103
104    pub fn clone(&self) -> Self {
105        let mut properties: FxHashMap<String, ItemType> = FxHashMap::default();
106        for (key, value) in self.properties.iter() {
107            properties.insert(key.clone(), value.clone());
108        }
109        Self {
110            type_name: self.type_name.clone(),
111            properties: properties,
112            required: self.required.clone(),
113            additional_properties: self.additional_properties,
114        }
115    }
116}
117
118#[derive(Debug, Clone, Deserialize, Serialize)]
119pub struct JsonSchema {
120    pub name: String,
121    pub schema: JsonItem,
122}
123
124impl JsonSchema {
125    pub fn new(name: String) -> Self {
126        let schema = JsonItem::default();
127        Self {
128            name: name.to_string(),
129            schema,
130        }
131    }
132
133    pub fn new_schema(name: String) -> Self {
134        Self {
135            name: name.to_string(),
136            schema: JsonItem::default(),
137        }
138    }
139
140    pub fn add_property(
141        &mut self,
142        prop_name: String,
143        type_name: String,
144        description: Option<String>,
145    ) {
146        let new_item = ItemType::new(type_name, description);
147        self.schema.add_property(prop_name, new_item);
148    }
149
150    pub fn add_array(&mut self, prop_name: String, items: Vec<(String, String)>) {
151        let mut array_item = JsonItem::default();
152        for (name, description) in items.iter() {
153            let item = ItemType::new(String::from("string"), Option::from(description.clone()));
154            array_item.add_property(name.clone(), item);
155        }
156        self.schema.add_array(prop_name, array_item);
157    }
158
159    pub fn clone(&self) -> Self {
160        Self {
161            name: self.name.clone(),
162            schema: self.schema.clone(),
163        }
164    }
165}
openai_tools/json_schema.rs

openai_tools/
json_schema.rs
