1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
//! JSON Schema optimizer for tool definitions
use serde_json::{Map, Value, json};
use std::collections::HashSet;
/// Schema optimizer for creating LLM-compatible JSON schemas
pub struct SchemaOptimizer;
impl SchemaOptimizer {
/// Create an optimized JSON schema from a raw schema
///
/// This function:
/// - Flattens $ref and $defs
/// - Removes additionalProperties
/// - Ensures all properties are required
pub fn optimize(schema: &Value) -> Value {
let mut result = schema.clone();
Self::optimize_recursive(&mut result);
result
}
fn optimize_recursive(value: &mut Value) {
match value {
Value::Object(obj) => {
// Remove additionalProperties
obj.remove("additionalProperties");
// Handle $ref by inlining the definition
if let Some(ref_val) = obj.get("$ref").and_then(|r| r.as_str())
&& let Some(defs) = obj.get("$defs")
&& let Some(def) = ref_val.strip_prefix("#/$defs/")
&& let Some(resolved) = defs.get(def)
{
*value = resolved.clone();
Self::optimize_recursive(value);
return;
}
// Remove $defs if present
obj.remove("$defs");
// Make all properties required
if let Some(properties) = obj.get("properties").and_then(|p| p.as_object()) {
let all_keys: HashSet<&str> = properties.keys().map(|k| k.as_str()).collect();
obj.insert("required".to_string(), json!(all_keys));
}
// Recursively process nested objects
for (_, v) in obj.iter_mut() {
Self::optimize_recursive(v);
}
}
Value::Array(arr) => {
for item in arr.iter_mut() {
Self::optimize_recursive(item);
}
}
_ => {}
}
}
/// Create a tool definition from a JSON schema
pub fn create_tool_definition(
name: impl Into<String>,
description: impl Into<String>,
schema: Value,
) -> super::ToolDefinition {
let optimized = Self::optimize(&schema);
let parameters = optimized.as_object().cloned().unwrap_or_else(|| {
let mut map = Map::new();
map.insert("type".to_string(), json!("object"));
map.insert("properties".to_string(), json!({}));
map
});
super::ToolDefinition {
name: name.into(),
description: description.into(),
parameters,
strict: true,
}
}
/// Create a minimal schema for a simple string parameter
pub fn string_schema() -> Value {
json!({
"type": "object",
"properties": {
"value": {
"type": "string",
"description": "The string value"
}
},
"required": ["value"]
})
}
/// Create a schema for multiple string parameters
pub fn string_params_schema(params: &[(&str, &str)]) -> Value {
let properties: Map<String, Value> = params
.iter()
.map(|(name, desc)| {
(
name.to_string(),
json!({
"type": "string",
"description": desc
}),
)
})
.collect();
let required: Vec<&str> = params.iter().map(|(name, _)| *name).collect();
json!({
"type": "object",
"properties": properties,
"required": required
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_optimize_schema() {
let schema = json!({
"$ref": "#/$defs/MyType",
"$defs": {
"MyType": {
"type": "object",
"properties": {
"name": { "type": "string" }
},
"additionalProperties": false
}
}
});
let optimized = SchemaOptimizer::optimize(&schema);
assert!(optimized.get("$ref").is_none());
assert!(optimized.get("$defs").is_none());
assert!(optimized.get("additionalProperties").is_none());
assert!(optimized.get("required").is_some());
}
}