ollama_rest/models/
json_schema.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4
5/// Function definition
6///
7/// Since 0.3.0
8#[derive(Debug, Serialize, Deserialize)]
9pub struct FunctionDef {
10    pub name: String,
11    pub description: Option<String>,
12    pub parameters: Option<Box<JsonSchema>>,
13}
14
15/// A **partly** implemented JSON Schema enum.
16///
17/// Since 0.3.0
18#[derive(Debug, Serialize, Deserialize)]
19#[serde(tag = "type", rename_all = "lowercase")]
20pub enum JsonSchema {
21    Function {
22        function: FunctionDef,
23    },
24    Integer {
25        description: Option<String>,
26    },
27    Number {
28        description: Option<String>,
29    },
30    Object {
31        properties: BTreeMap<String, JsonSchema>,
32        required: Option<Vec<String>>,
33    },
34    String {
35        description: Option<String>,
36        #[serde(rename = "enum")]
37        enumeration: Option<Vec<String>>,
38    },
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44
45    #[test]
46    fn def_function_schema() {
47        const FUNC_NAME: &'static str = "query_weather";
48        const FUNC_DESC: &'static str = "Get current weather in a specified location.";
49
50        const LOC_DESC: &'static str = "Keywords of the location.";
51
52        let obj = serde_json::from_value::<JsonSchema>(serde_json::json!({
53            "type": "function",
54            "function": {
55                "name": FUNC_NAME,
56                "description": FUNC_DESC,
57                "parameters": {
58                    "type": "object",
59                    "properties": {
60                        "location": {
61                            "type": "string",
62                            "description": LOC_DESC,
63                        },
64                    },
65                    "required": ["location"],
66                },
67            },
68        })).unwrap();
69
70        assert!(matches!(obj, JsonSchema::Function { .. }));
71
72        if let JsonSchema::Function { function } = obj {
73            assert_eq!(function.name, FUNC_NAME);
74
75            assert!(matches!(function.description, Some(_)));
76            assert_eq!(function.description.unwrap(), FUNC_DESC);
77
78            assert!(matches!(function.parameters, Some(_)));
79
80            if let Some(boxed_schema) = function.parameters {
81                let param_schema = *boxed_schema;
82
83                assert!(matches!(param_schema, JsonSchema::Object { .. }));
84                if let JsonSchema::Object { properties, required } = param_schema {
85                    let location_schema = properties.get("location");
86                    assert!(matches!(location_schema, Some(_)));
87                    if let Some(location_schema) = location_schema {
88                        assert!(matches!(location_schema, JsonSchema::String { .. }));
89                        if let JsonSchema::String { description, enumeration } = location_schema {
90                            assert!(matches!(description, Some(_)));
91                            if let Some(description) = description {
92                                assert_eq!(description, LOC_DESC);
93                            }
94
95                            assert!(matches!(enumeration, None));
96                        }
97                    }
98
99                    assert!(matches!(required, Some(_)));
100                    if let Some(required_fields) = required {
101                        assert_eq!(required_fields.len(), 1);
102                        assert_eq!(required_fields[0], "location");
103                    }
104                }
105            }
106        }
107    }
108}
ollama_rest/models/json_schema.rs

ollama_rest/models/
json_schema.rs
