subako-function-sdk 0.2.0

//! A package as a value: its functions, the manifest they describe, and the
//! `main` that serves them.

use std::io::Write as _;
use std::process::ExitCode;

use schemars::JsonSchema;
use schemars::generate::{Contract, SchemaSettings};
use serde::Serialize;
use serde::de::DeserializeOwned;

use crate::{Input, Output, ToolError};

/// The argument that prints the manifest instead of calling a function.
pub const MANIFEST_FLAG: &str = "--subako-manifest";

/// One function a package serves as a tool.
pub struct Function {
    name: String,
    description: String,
    input_schema: fn(Contract) -> serde_json::Value,
    output_schema: fn(Contract) -> serde_json::Value,
    call: Box<dyn Fn(Input) -> Result<Output, ToolError>>,
}

impl Function {
    /// `f` as the tool `name`. The schemas come from `I` and `O`.
    pub fn new<I, O, E>(
        name: impl Into<String>,
        description: impl Into<String>,
        f: impl Fn(I) -> Result<O, E> + 'static,
    ) -> Self
    where
        I: DeserializeOwned + JsonSchema,
        O: Serialize + JsonSchema,
        E: Into<ToolError>,
    {
        Function {
            name: name.into(),
            description: description.into(),
            input_schema: schema::<I>,
            output_schema: schema::<O>,
            call: Box::new(move |input| Output::json(&f(input.parse()?).map_err(Into::into)?)),
        }
    }

    pub fn name(&self) -> &str {
        &self.name
    }
}

/// What `subako-package.json` holds.
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct Manifest {
    pub name: String,
    pub description: String,
    pub functions: Vec<ManifestFunction>,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct ManifestFunction {
    pub name: String,
    pub description: String,
    pub input_schema: serde_json::Value,
    pub output_schema: serde_json::Value,
}

/// A package: its name, its description, and the functions it serves.
pub struct Package {
    name: String,
    description: String,
    functions: Vec<Function>,
}

impl Package {
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        functions: Vec<Function>,
    ) -> Self {
        Package {
            name: name.into(),
            description: description.into(),
            functions,
        }
    }

    pub fn functions(&self) -> &[Function] {
        &self.functions
    }

    pub fn manifest(&self) -> Manifest {
        Manifest {
            name: self.name.clone(),
            description: self.description.clone(),
            functions: self
                .functions
                .iter()
                .map(|f| ManifestFunction {
                    name: f.name.clone(),
                    description: f.description.clone(),
                    input_schema: (f.input_schema)(Contract::Deserialize),
                    output_schema: (f.output_schema)(Contract::Serialize),
                })
                .collect(),
        }
    }

    /// Calls the function named `tool`.
    pub fn call(&self, tool: &str, input: Input) -> Result<Output, ToolError> {
        match self.functions.iter().find(|f| f.name == tool) {
            Some(function) => (function.call)(input),
            None => Err(ToolError::unknown_tool(tool)),
        }
    }

    /// The whole process as a pure function of its arguments and stdin: the
    /// manifest under [`MANIFEST_FLAG`], otherwise a call.
    pub fn execute(&self, args: &[String], stdin: &[u8]) -> (Vec<u8>, u8) {
        match args.get(1).map(String::as_str) {
            Some(MANIFEST_FLAG) => (self.manifest_json(), 0),
            _ => crate::execute(args, stdin, |tool, input| self.call(tool, input)),
        }
    }

    /// The program's `main`. The manifest is printed without reading stdin,
    /// so `module --subako-manifest` needs no input.
    // COVERAGE(untestable): reads the real process's arguments and stdio; `execute` is its testable half.
    pub fn run(&self) -> ExitCode {
        let args: Vec<String> = std::env::args().collect();
        match args.get(1).map(String::as_str) {
            Some(MANIFEST_FLAG) => {
                let mut stdout = std::io::stdout().lock();
                match stdout
                    .write_all(&self.manifest_json())
                    .and_then(|()| stdout.flush())
                {
                    Ok(()) => ExitCode::SUCCESS,
                    Err(_) => ExitCode::from(2),
                }
            }
            _ => crate::run(|tool, input| self.call(tool, input)),
        }
    }

    fn manifest_json(&self) -> Vec<u8> {
        // A manifest is plain data, so it always serializes.
        let mut json = serde_json::to_vec_pretty(&self.manifest()).unwrap_or_default();
        json.push(b'\n');
        json
    }
}

/// `T`'s JSON Schema as `contract` reads or writes it, with every subschema
/// inline so a model reads one self-contained object.
fn schema<T: JsonSchema>(contract: Contract) -> serde_json::Value {
    let mut schema = SchemaSettings::draft2020_12()
        .with(|s| {
            s.inline_subschemas = true;
            s.meta_schema = None;
            s.contract = contract;
        })
        .into_generator()
        .into_root_schema_for::<T>();
    schema.remove("title");
    schema.to_value()
}