subako-function-sdk 0.2.0

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use subako_function_sdk::{MANIFEST_FLAG, ToolError, function, package};

#[derive(Deserialize, JsonSchema)]
struct GreetInput {
    /// Who to greet.
    name: String,
    #[serde(default)]
    shout: bool,
}

#[derive(Serialize, JsonSchema)]
struct Greeting {
    text: String,
}

/// Greet someone by name.
#[function]
fn greet(input: GreetInput) -> Result<Greeting, ToolError> {
    let text = format!("hello {}", input.name);
    Ok(Greeting {
        text: match input.shout {
            true => text.to_uppercase(),
            false => text,
        },
    })
}

mod tools {
    use super::*;

    #[derive(Deserialize, JsonSchema)]
    pub struct Nothing {}

    /// Always fails.
    #[function]
    pub fn fail(_: Nothing) -> Result<Greeting, String> {
        Err("no".to_string())
    }
}

package! {
    name: "hello",
    description: concat!("Says ", "hello."),
    functions: [greet, tools::fail],
    main: false,
}

fn execute(args: &[&str], stdin: &str) -> (Value, u8) {
    let args: Vec<String> = args.iter().map(ToString::to_string).collect();
    let (out, code) = package().execute(&args, stdin.as_bytes());
    (serde_json::from_slice(&out).unwrap(), code)
}

#[test]
fn the_manifest_comes_from_the_functions() {
    let (manifest, code) = execute(&["hello", MANIFEST_FLAG], "");
    assert_eq!(code, 0);
    assert_eq!(
        manifest,
        json!({
            "name": "hello",
            "description": "Says hello.",
            "functions": [
                {
                    "name": "greet",
                    "description": "Greet someone by name.",
                    "input_schema": {
                        "type": "object",
                        "properties": {
                            "name": {"type": "string", "description": "Who to greet."},
                            "shout": {"type": "boolean", "default": false},
                        },
                        "required": ["name"],
                    },
                    "output_schema": {
                        "type": "object",
                        "properties": {"text": {"type": "string"}},
                        "required": ["text"],
                    },
                },
                {
                    "name": "fail",
                    "description": "Always fails.",
                    "input_schema": {"type": "object"},
                    "output_schema": {
                        "type": "object",
                        "properties": {"text": {"type": "string"}},
                        "required": ["text"],
                    },
                },
            ],
        })
    );
}

#[test]
fn a_call_runs_the_named_function() {
    assert_eq!(
        execute(&["hello", "greet"], r#"{"name":"ada","shout":true}"#),
        (json!({"text": "HELLO ADA"}), 0)
    );
}

#[test]
fn a_function_error_reaches_the_model() {
    assert_eq!(
        execute(&["hello", "fail"], "{}"),
        (json!({"error": "no"}), 1)
    );
}

#[test]
fn bad_input_and_unknown_tools_are_tool_errors() {
    let (out, code) = execute(&["hello", "greet"], r#"{"shout":true}"#);
    assert_eq!(code, 1);
    assert!(out["error"].as_str().unwrap().starts_with("invalid input"));
    assert_eq!(
        execute(&["hello", "nope"], "{}"),
        (json!({"error": "unknown tool: nope"}), 1)
    );
}

#[test]
fn the_function_stays_callable() {
    let greeting = greet(GreetInput {
        name: "ada".to_string(),
        shout: false,
    })
    .unwrap();
    assert_eq!(greeting.text, "hello ada");
    assert_eq!(
        package()
            .functions()
            .iter()
            .map(|f| f.name())
            .collect::<Vec<_>>(),
        ["greet", "fail"]
    );
}