Skip to main content

headless_lms_utils/
json_schema_types.rs

1use indexmap::IndexMap;
2use serde::{Deserialize, Serialize};
3
4#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
5#[serde(rename_all = "snake_case")]
6pub enum JSONType {
7    JsonSchema,
8    Object,
9    Array,
10    String,
11    Number,
12    Integer,
13    Boolean,
14}
15
16/// Defines the shape of a JSON object for the LLM, used both for structured output
17/// and for the parameters of a tool definition.
18#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
19#[serde(rename_all = "camelCase")]
20pub struct Schema {
21    #[serde(rename = "type")]
22    /// Type of the schema, should be Object
23    pub type_field: JSONType,
24    /// Order-preserving, and deliberately not a `HashMap`: `RandomState` reseeds per map instance,
25    /// so a `HashMap` here serializes its keys in a different order on nearly every request. Tool
26    /// definitions and structured output schemas sit at the front of the prompt, and Azure's
27    /// prompt cache matches an exact prefix, so that alone misses the cache on every request.
28    pub properties: IndexMap<String, SchemaPropertyType>,
29    /// All 'properties' keys must be included in this 'required' list
30    pub required: Vec<String>,
31    /// additionalProperties should always be 'false'
32    pub additional_properties: bool,
33    /// Explains the object to the LLM.
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub description: Option<String>,
36}
37
38impl Schema {
39    /// A strict-mode object schema with every property key required.
40    ///
41    /// Azure's `strict: true` tool/response schemas have no notion of an optional property —
42    /// a property a tool doesn't strictly need still has to be listed here, with "Optional"
43    /// said in its own description instead.
44    pub fn strict_object(
45        properties: IndexMap<String, SchemaPropertyType>,
46        description: Option<&str>,
47    ) -> Self {
48        let required = properties.keys().cloned().collect();
49        Self {
50            type_field: JSONType::Object,
51            properties,
52            required,
53            additional_properties: false,
54            description: description.map(str::to_string),
55        }
56    }
57}
58
59#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
60#[serde(untagged)]
61pub enum SchemaPropertyType {
62    ArrayProperty(ArrayProperty),
63    Object(Schema),
64    Item(JsonItem),
65}
66
67#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
68pub struct ArrayProperty {
69    #[serde(rename = "type")]
70    pub type_field: JSONType,
71    pub items: ArrayItem,
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub description: Option<String>,
74}
75
76#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
77#[serde(untagged)]
78pub enum ArrayItem {
79    Schema(Schema),
80    JsonItem(JsonItem),
81}
82
83/// A scalar value in a [Schema]: the property types that have no inner shape.
84#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
85pub struct JsonItem {
86    #[serde(rename = "type")]
87    pub type_field: JSONType,
88    #[serde(skip_serializing_if = "Option::is_none")]
89    pub description: Option<String>,
90}
91
92/// An array of plain strings, the one composite property that tool parameter and structured output
93/// schemas keep asking for. `description` explains the array to the LLM; the items carry none of
94/// their own.
95pub fn string_array_property(description: Option<&str>) -> SchemaPropertyType {
96    SchemaPropertyType::ArrayProperty(ArrayProperty {
97        type_field: JSONType::Array,
98        items: ArrayItem::JsonItem(JsonItem {
99            type_field: JSONType::String,
100            description: None,
101        }),
102        description: description.map(str::to_string),
103    })
104}