1use imp_llm::ContentBlock;
2use serde_json::{json, Map, Value};
3
4use crate::error::{Error, Result};
5use crate::tools::ToolOutput;
6
7#[must_use]
13pub fn parameter_schema_from_lua(params: &Value) -> Value {
14 if looks_like_json_schema(params) {
15 return params.clone();
16 }
17
18 let mut properties = match params {
19 Value::Object(map) => map.clone(),
20 _ => Map::new(),
21 };
22
23 let required: Vec<Value> = properties
24 .iter()
25 .filter_map(|(name, definition)| {
26 definition
27 .get("required")
28 .and_then(Value::as_bool)
29 .filter(|required| *required)
30 .map(|_| Value::String(name.clone()))
31 })
32 .collect();
33
34 for (_name, definition) in properties.iter_mut() {
36 if let Value::Object(ref mut map) = definition {
37 map.remove("required");
38 }
39 }
40
41 let mut schema = json!({
42 "type": "object",
43 "properties": properties,
44 });
45
46 if !required.is_empty() {
47 schema["required"] = Value::Array(required);
48 }
49
50 schema
51}
52
53pub fn tool_output_from_lua_result(result: Value) -> Result<ToolOutput> {
60 match result {
61 Value::Null => Ok(ToolOutput {
62 content: Vec::new(),
63 details: Value::Null,
64 is_error: false,
65 }),
66 Value::String(text) => Ok(ToolOutput::text(text)),
67 Value::Object(mut object) => {
68 let is_error = object
69 .remove("is_error")
70 .or_else(|| object.remove("isError"))
71 .and_then(|value| value.as_bool())
72 .unwrap_or(false);
73 let details = object.remove("details").unwrap_or(Value::Null);
74 let content = parse_lua_content(object.remove("content").unwrap_or(Value::Null))?;
75
76 Ok(ToolOutput {
77 content,
78 details,
79 is_error,
80 })
81 }
82 other => Ok(ToolOutput::text(other.to_string())),
83 }
84}
85
86fn looks_like_json_schema(params: &Value) -> bool {
87 let Some(object) = params.as_object() else {
88 return false;
89 };
90
91 object.contains_key("type")
92 || object.contains_key("properties")
93 || object.contains_key("required")
94 || object.contains_key("anyOf")
95 || object.contains_key("oneOf")
96 || object.contains_key("allOf")
97 || object.contains_key("$ref")
98}
99
100fn parse_lua_content(content: Value) -> Result<Vec<ContentBlock>> {
101 match content {
102 Value::Null => Ok(Vec::new()),
103 Value::String(text) => Ok(vec![ContentBlock::Text { text }]),
104 Value::Array(_) => serde_json::from_value(content).map_err(|error| {
105 Error::Tool(format!("Lua tool returned invalid content blocks: {error}"))
106 }),
107 Value::Object(object) if object.contains_key("type") => {
108 let block: ContentBlock =
109 serde_json::from_value(Value::Object(object)).map_err(|error| {
110 Error::Tool(format!(
111 "Lua tool returned an invalid content block: {error}"
112 ))
113 })?;
114 Ok(vec![block])
115 }
116 Value::Object(object) => Ok(vec![ContentBlock::Text {
117 text: Value::Object(object).to_string(),
118 }]),
119 other => Ok(vec![ContentBlock::Text {
120 text: other.to_string(),
121 }]),
122 }
123}
124
125#[cfg(test)]
126mod tests {
127 use super::*;
128
129 fn extract_text(output: &ToolOutput) -> String {
130 output
131 .content
132 .iter()
133 .filter_map(|block| match block {
134 ContentBlock::Text { text } => Some(text.as_str()),
135 _ => None,
136 })
137 .collect::<Vec<_>>()
138 .join("\n")
139 }
140
141 #[test]
142 fn parameter_schema_wraps_shorthand_definitions() {
143 let schema = parameter_schema_from_lua(&json!({
144 "name": { "type": "string", "required": true },
145 "times": { "type": "number" }
146 }));
147
148 assert_eq!(schema["type"], "object");
149 assert_eq!(schema["properties"]["name"]["type"], "string");
150 assert_eq!(schema["properties"]["times"]["type"], "number");
151 assert_eq!(schema["required"], json!(["name"]));
152 }
153
154 #[test]
155 fn parameter_schema_preserves_full_schema() {
156 let original = json!({
157 "type": "object",
158 "properties": {
159 "path": { "type": "string" }
160 },
161 "required": ["path"]
162 });
163
164 assert_eq!(parameter_schema_from_lua(&original), original);
165 }
166
167 #[test]
168 fn tool_output_parser_accepts_string_content() {
169 let output = tool_output_from_lua_result(json!({
170 "content": "hello from lua",
171 "details": { "source": "test" },
172 "is_error": true
173 }))
174 .unwrap();
175
176 assert!(output.is_error);
177 assert_eq!(output.details["source"], "test");
178 assert_eq!(extract_text(&output), "hello from lua");
179 }
180}