Skip to main content

lash_lashlang_runtime/
typed_output.rs

1//! The RLM typed-output vocabulary: parsing the `output` argument that
2//! RLM tools (`llm.query`, `agents.spawn`) accept into a JSON Schema.
3
4use serde_json::{Value, json};
5
6use crate::LASH_TYPE_KEY;
7
8/// Parse a tool's `output` argument into the JSON Schema it requests.
9///
10/// Two shapes are accepted:
11/// - a record of field-name → type-descriptor strings (`"str"`, `"int"`,
12///   `"float"`, `"bool"`, `"record"`, or `"list[...]"` of those), compiled
13///   into a strict object schema; or
14/// - a Lashlang `Type { ... }` literal — a single-field
15///   `{"$lash_type": <schema>}` wrapper as produced by the Lashlang
16///   compiler — whose inner schema is passed through after validation.
17///
18/// Returns `Ok(None)` when `output` is absent or `null` (the tool falls back
19/// to its untyped default).
20pub fn parse_output_schema(value: Option<&Value>) -> Result<Option<Value>, String> {
21    let Some(value) = value else {
22        return Ok(None);
23    };
24    if value.is_null() {
25        return Ok(None);
26    }
27    let output = value.as_object().ok_or_else(|| {
28        "invalid `output`: expected a record describing the typed shape".to_string()
29    })?;
30    if output.is_empty() {
31        return Err("at least one output field is required".to_string());
32    }
33
34    if output.len() == 1
35        && let Some(schema) = output.get(LASH_TYPE_KEY)
36    {
37        validate_lash_type_schema(schema)?;
38        return Ok(Some(schema.clone()));
39    }
40
41    let mut properties = serde_json::Map::new();
42    let mut required = Vec::new();
43    for (name, descriptor) in output {
44        let type_str = descriptor
45            .as_str()
46            .ok_or_else(|| format!("field `{name}`: type descriptor must be a string"))?;
47        properties.insert(name.clone(), type_descriptor_to_json_schema(type_str)?);
48        required.push(Value::String(name.clone()));
49    }
50    Ok(Some(json!({
51        "type": "object",
52        "properties": properties,
53        "required": required,
54        "additionalProperties": false,
55    })))
56}
57
58fn validate_lash_type_schema(schema: &Value) -> Result<(), String> {
59    let object = schema
60        .as_object()
61        .ok_or_else(|| "Type schema must be a JSON object".to_string())?;
62    let kind = object
63        .get("type")
64        .and_then(Value::as_str)
65        .ok_or_else(|| "Type schema missing `type` field".to_string())?;
66    match kind {
67        "object" | "array" | "string" | "integer" | "number" | "boolean" => Ok(()),
68        other => Err(format!("unsupported Type schema kind `{other}`")),
69    }
70}
71
72fn type_descriptor_to_json_schema(descriptor: &str) -> Result<Value, String> {
73    let scalar = |ty: &str| -> Result<Value, String> {
74        match ty {
75            "str" | "string" => Ok(json!({"type": "string"})),
76            "int" | "integer" => Ok(json!({"type": "integer"})),
77            "float" | "number" => Ok(json!({"type": "number"})),
78            "bool" | "boolean" => Ok(json!({"type": "boolean"})),
79            "record" | "dict" | "object" => {
80                Ok(json!({"type": "object", "additionalProperties": true}))
81            }
82            other => Err(format!("unknown scalar type `{other}`")),
83        }
84    };
85    let trimmed = descriptor.trim();
86    if let Some(inner) = trimmed
87        .strip_prefix("list[")
88        .and_then(|rest| rest.strip_suffix(']'))
89    {
90        return Ok(json!({
91            "type": "array",
92            "items": scalar(inner.trim())?,
93        }));
94    }
95    scalar(trimmed)
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn output_schema_supports_scalars_and_lists() {
104        let schema = parse_output_schema(Some(&json!({
105            "answer": "str",
106            "count": "int",
107            "items": "list[str]"
108        })))
109        .expect("schema")
110        .expect("present");
111        assert_eq!(schema["properties"]["answer"]["type"], json!("string"));
112        assert_eq!(schema["properties"]["count"]["type"], json!("integer"));
113        assert_eq!(schema["properties"]["items"]["type"], json!("array"));
114    }
115
116    #[test]
117    fn output_schema_passes_through_lash_type_wrapper() {
118        let inner_schema = json!({
119            "type": "object",
120            "properties": {
121                "name": { "type": "string" },
122                "tags": { "type": "array", "items": { "type": "string" } },
123                "status": { "type": "string", "enum": ["ok", "err"] }
124            },
125            "required": ["name", "tags", "status"],
126            "additionalProperties": false
127        });
128        let wrapped = json!({ LASH_TYPE_KEY: inner_schema.clone() });
129        let schema = parse_output_schema(Some(&wrapped))
130            .expect("schema")
131            .expect("present");
132        assert_eq!(schema, inner_schema);
133    }
134
135    #[test]
136    fn output_schema_rejects_lash_type_without_type_field() {
137        let wrapped = json!({ LASH_TYPE_KEY: {"properties": {}} });
138        let err = parse_output_schema(Some(&wrapped)).expect_err("missing type");
139        assert!(err.contains("type"), "error: {err}");
140    }
141
142    #[test]
143    fn output_schema_accepts_array_top_level_type() {
144        let wrapped = json!({
145            LASH_TYPE_KEY: {
146                "type": "array",
147                "items": {"type": "string"}
148            }
149        });
150        let schema = parse_output_schema(Some(&wrapped))
151            .expect("schema")
152            .expect("present");
153        assert_eq!(schema["type"], json!("array"));
154    }
155}