openai_tools/common/
structured_output.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4#[derive(Debug, Clone, Deserialize, Serialize)]
5struct ItemType {
6    #[serde(rename = "type")]
7    type_name: String,
8    #[serde(skip_serializing_if = "Option::is_none")]
9    description: Option<String>,
10    #[serde(skip_serializing_if = "Option::is_none")]
11    items: Option<Box<JsonItem>>,
12}
13
14impl ItemType {
15    pub fn new<T: AsRef<str>, U: AsRef<str>>(type_name: T, description: U) -> Self {
16        Self {
17            type_name: type_name.as_ref().to_string(),
18            description: match description.as_ref() {
19                "" => None,
20                _ => Some(description.as_ref().to_string()),
21            },
22            items: None,
23        }
24    }
25
26    pub fn clone(&self) -> Self {
27        let mut items: JsonItem = JsonItem::default();
28        if let Some(item) = &self.items {
29            let mut _properties: HashMap<String, ItemType> = HashMap::default();
30            for (key, value) in item.properties.iter() {
31                _properties.insert(key.clone(), value.clone());
32            }
33            items.type_name = item.type_name.clone();
34            items.properties = _properties;
35            items.required = item.required.clone();
36            items.additional_properties = item.additional_properties;
37        }
38
39        Self {
40            type_name: self.type_name.clone(),
41            description: self.description.clone(),
42            items: if self.items.is_some() { Option::from(Box::new(items)) } else { None },
43        }
44    }
45}
46
47#[derive(Debug, Clone, Deserialize, Serialize)]
48struct JsonItem {
49    #[serde(rename = "type")]
50    type_name: Option<String>,
51    properties: HashMap<String, ItemType>,
52    #[serde(skip_serializing_if = "Option::is_none")]
53    required: Option<Vec<String>>,
54    #[serde(rename = "additionalProperties")]
55    additional_properties: bool,
56}
57
58impl JsonItem {
59    fn add_property<T: AsRef<str>>(&mut self, prop_name: T, item: ItemType) {
60        self.properties.insert(prop_name.as_ref().to_string(), item.clone());
61        if self.required.is_none() {
62            self.required = Some(vec![]);
63        }
64        self.required.as_mut().unwrap().push(prop_name.as_ref().to_string());
65    }
66
67    fn add_array<T: AsRef<str>>(&mut self, prop_name: T, items: JsonItem) {
68        let mut prop = ItemType::new("array", "");
69        prop.items = Option::from(Box::new(items));
70        self.properties.insert(prop_name.as_ref().to_string(), prop);
71        if self.required.is_none() {
72            self.required = Some(vec![]);
73        }
74        self.required.as_mut().unwrap().push(prop_name.as_ref().to_string());
75    }
76}
77
78impl Default for JsonItem {
79    fn default() -> Self {
80        Self { type_name: Some("object".to_string()), properties: HashMap::new(), required: None, additional_properties: false }
81    }
82}
83
84#[derive(Debug, Clone, Default, Deserialize, Serialize)]
85pub struct Schema {
86    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
87    type_name: Option<String>,
88    #[serde(skip_serializing_if = "Option::is_none")]
89    name: Option<String>,
90    #[serde(skip_serializing_if = "Option::is_none")]
91    schema: Option<JsonItem>,
92}
93
94impl Schema {
95    pub fn responses_text_schema() -> Self {
96        Self { type_name: Some("text".to_string()), name: None, schema: None }
97    }
98
99    pub fn responses_json_schema<T: AsRef<str>>(name: T) -> Self {
100        Self { type_name: Some("json_schema".to_string()), name: Some(name.as_ref().to_string()), schema: Some(JsonItem::default()) }
101    }
102
103    pub fn chat_json_schema<T: AsRef<str>>(name: T) -> Self {
104        Self { type_name: None, name: Some(name.as_ref().to_string()), schema: Some(JsonItem::default()) }
105    }
106
107    pub fn add_property<T: AsRef<str>, U: AsRef<str>, V: AsRef<str>>(&mut self, prop_name: T, type_name: U, description: V) {
108        let new_item = ItemType::new(type_name, description);
109        self.schema.as_mut().unwrap().add_property(prop_name, new_item);
110    }
111
112    pub fn add_array<T: AsRef<str>, U: AsRef<str>>(&mut self, prop_name: T, items: Vec<(U, U)>) {
113        let mut array_item = JsonItem::default();
114        for (name, description) in items.iter() {
115            let item = ItemType::new("string", description.as_ref());
116            array_item.add_property(name.as_ref(), item);
117        }
118        self.schema.as_mut().unwrap().add_array(prop_name, array_item);
119    }
120}
openai_tools/common/structured_output.rs

openai_tools/common/
structured_output.rs
