Skip to main content

imp_core/tools/
lua.rs

1use imp_llm::ContentBlock;
2use serde_json::{json, Map, Value};
3
4use crate::error::{Error, Result};
5use crate::tools::ToolOutput;
6
7/// Normalize Lua tool parameter definitions into JSON Schema.
8///
9/// Lua extensions may register either a full JSON Schema object or a shorthand
10/// object whose keys are parameter names. The returned value is always a valid
11/// object schema for the tool registry.
12#[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    // Strip the non-standard "required" field from each property definition
35    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
53/// Convert a Lua tool's JSON result into imp's native `ToolOutput`.
54///
55/// Supported result forms:
56/// - a plain string → single text block
57/// - an object with `{ content, details, is_error }`
58/// - `content` as a string, a single content block, or an array of blocks
59pub 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}