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 { type_name: self.type_name.clone(), description: self.description.clone(), items: if self.items.is_some() { Option::from(Box::new(items)) } else { None } }
40    }
41}
42
43#[derive(Debug, Clone, Deserialize, Serialize)]
44struct JsonItem {
45    #[serde(rename = "type")]
46    type_name: Option<String>,
47    properties: HashMap<String, ItemType>,
48    #[serde(skip_serializing_if = "Option::is_none")]
49    required: Option<Vec<String>>,
50    #[serde(rename = "additionalProperties")]
51    additional_properties: bool,
52}
53
54impl JsonItem {
55    fn add_property<T: AsRef<str>>(&mut self, prop_name: T, item: ItemType) {
56        self.properties.insert(prop_name.as_ref().to_string(), item.clone());
57        if self.required.is_none() {
58            self.required = Option::from(vec![prop_name.as_ref().to_string()]);
59        } else {
60            let mut required = self.required.clone().unwrap();
61            required.push(prop_name.as_ref().to_string());
62            self.required = Option::from(required);
63        }
64    }
65
66    fn add_array<T: AsRef<str>>(&mut self, prop_name: T, items: JsonItem) {
67        let mut prop = ItemType::new("array", "");
68        prop.items = Option::from(Box::new(items));
69        self.properties.insert(prop_name.as_ref().to_string(), prop);
70        self.required = Option::from(vec![prop_name.as_ref().to_string()]);
71    }
72}
73
74impl Default for JsonItem {
75    fn default() -> Self {
76        Self { type_name: Some("object".to_string()), properties: HashMap::new(), required: None, additional_properties: false }
77    }
78}
79
80#[derive(Debug, Clone, Default, Deserialize, Serialize)]
81pub struct Schema {
82    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
83    type_name: Option<String>,
84    #[serde(skip_serializing_if = "Option::is_none")]
85    name: Option<String>,
86    #[serde(skip_serializing_if = "Option::is_none")]
87    schema: Option<JsonItem>,
88}
89
90impl Schema {
91    pub fn responses_text_schema() -> Self {
92        Self { type_name: Some("text".to_string()), name: None, schema: None }
93    }
94
95    pub fn responses_json_schema<T: AsRef<str>>(name: T) -> Self {
96        Self { type_name: Some("json_schema".to_string()), name: Some(name.as_ref().to_string()), schema: Some(JsonItem::default()) }
97    }
98
99    pub fn chat_json_schema<T: AsRef<str>>(name: T) -> Self {
100        Self { type_name: None, name: Some(name.as_ref().to_string()), schema: Some(JsonItem::default()) }
101    }
102
103    pub fn add_property<T: AsRef<str>, U: AsRef<str>, V: AsRef<str>>(&mut self, prop_name: T, type_name: U, description: V) {
104        let new_item = ItemType::new(type_name, description);
105        self.schema.as_mut().unwrap().add_property(prop_name, new_item);
106    }
107
108    pub fn add_array<T: AsRef<str>, U: AsRef<str>>(&mut self, prop_name: T, items: Vec<(U, U)>) {
109        let mut array_item = JsonItem::default();
110        for (name, description) in items.iter() {
111            let item = ItemType::new("string", description.as_ref());
112            array_item.add_property(name.as_ref(), item);
113        }
114        self.schema.as_mut().unwrap().add_array(prop_name, array_item);
115    }
116}
openai_tools/common/structured_output.rs

openai_tools/common/
structured_output.rs
