subako_function_sdk/
package.rs

1//! A package as a value: its functions, the manifest they describe, and the
2//! `main` that serves them.
3
4use std::io::Write as _;
5use std::process::ExitCode;
6
7use schemars::JsonSchema;
8use schemars::generate::{Contract, SchemaSettings};
9use serde::Serialize;
10use serde::de::DeserializeOwned;
11
12use crate::{Input, Output, ToolError};
13
14/// The argument that prints the manifest instead of calling a function.
15pub const MANIFEST_FLAG: &str = "--subako-manifest";
16
17/// One function a package serves as a tool.
18pub struct Function {
19    name: String,
20    description: String,
21    input_schema: fn(Contract) -> serde_json::Value,
22    output_schema: fn(Contract) -> serde_json::Value,
23    call: Box<dyn Fn(Input) -> Result<Output, ToolError>>,
24}
25
26impl Function {
27    /// `f` as the tool `name`. The schemas come from `I` and `O`.
28    pub fn new<I, O, E>(
29        name: impl Into<String>,
30        description: impl Into<String>,
31        f: impl Fn(I) -> Result<O, E> + 'static,
32    ) -> Self
33    where
34        I: DeserializeOwned + JsonSchema,
35        O: Serialize + JsonSchema,
36        E: Into<ToolError>,
37    {
38        Function {
39            name: name.into(),
40            description: description.into(),
41            input_schema: schema::<I>,
42            output_schema: schema::<O>,
43            call: Box::new(move |input| Output::json(&f(input.parse()?).map_err(Into::into)?)),
44        }
45    }
46
47    pub fn name(&self) -> &str {
48        &self.name
49    }
50}
51
52/// What `subako-package.json` holds.
53#[derive(Debug, Clone, PartialEq, Serialize)]
54pub struct Manifest {
55    pub name: String,
56    pub description: String,
57    pub functions: Vec<ManifestFunction>,
58}
59
60#[derive(Debug, Clone, PartialEq, Serialize)]
61pub struct ManifestFunction {
62    pub name: String,
63    pub description: String,
64    pub input_schema: serde_json::Value,
65    pub output_schema: serde_json::Value,
66}
67
68/// A package: its name, its description, and the functions it serves.
69pub struct Package {
70    name: String,
71    description: String,
72    functions: Vec<Function>,
73}
74
75impl Package {
76    pub fn new(
77        name: impl Into<String>,
78        description: impl Into<String>,
79        functions: Vec<Function>,
80    ) -> Self {
81        Package {
82            name: name.into(),
83            description: description.into(),
84            functions,
85        }
86    }
87
88    pub fn functions(&self) -> &[Function] {
89        &self.functions
90    }
91
92    pub fn manifest(&self) -> Manifest {
93        Manifest {
94            name: self.name.clone(),
95            description: self.description.clone(),
96            functions: self
97                .functions
98                .iter()
99                .map(|f| ManifestFunction {
100                    name: f.name.clone(),
101                    description: f.description.clone(),
102                    input_schema: (f.input_schema)(Contract::Deserialize),
103                    output_schema: (f.output_schema)(Contract::Serialize),
104                })
105                .collect(),
106        }
107    }
108
109    /// Calls the function named `tool`.
110    pub fn call(&self, tool: &str, input: Input) -> Result<Output, ToolError> {
111        match self.functions.iter().find(|f| f.name == tool) {
112            Some(function) => (function.call)(input),
113            None => Err(ToolError::unknown_tool(tool)),
114        }
115    }
116
117    /// The whole process as a pure function of its arguments and stdin: the
118    /// manifest under [`MANIFEST_FLAG`], otherwise a call.
119    pub fn execute(&self, args: &[String], stdin: &[u8]) -> (Vec<u8>, u8) {
120        match args.get(1).map(String::as_str) {
121            Some(MANIFEST_FLAG) => (self.manifest_json(), 0),
122            _ => crate::execute(args, stdin, |tool, input| self.call(tool, input)),
123        }
124    }
125
126    /// The program's `main`. The manifest is printed without reading stdin,
127    /// so `module --subako-manifest` needs no input.
128    // COVERAGE(untestable): reads the real process's arguments and stdio; `execute` is its testable half.
129    pub fn run(&self) -> ExitCode {
130        let args: Vec<String> = std::env::args().collect();
131        match args.get(1).map(String::as_str) {
132            Some(MANIFEST_FLAG) => {
133                let mut stdout = std::io::stdout().lock();
134                match stdout
135                    .write_all(&self.manifest_json())
136                    .and_then(|()| stdout.flush())
137                {
138                    Ok(()) => ExitCode::SUCCESS,
139                    Err(_) => ExitCode::from(2),
140                }
141            }
142            _ => crate::run(|tool, input| self.call(tool, input)),
143        }
144    }
145
146    fn manifest_json(&self) -> Vec<u8> {
147        // A manifest is plain data, so it always serializes.
148        let mut json = serde_json::to_vec_pretty(&self.manifest()).unwrap_or_default();
149        json.push(b'\n');
150        json
151    }
152}
153
154/// `T`'s JSON Schema as `contract` reads or writes it, with every subschema
155/// inline so a model reads one self-contained object.
156fn schema<T: JsonSchema>(contract: Contract) -> serde_json::Value {
157    let mut schema = SchemaSettings::draft2020_12()
158        .with(|s| {
159            s.inline_subschemas = true;
160            s.meta_schema = None;
161            s.contract = contract;
162        })
163        .into_generator()
164        .into_root_schema_for::<T>();
165    schema.remove("title");
166    schema.to_value()
167}
subako_function_sdk/package.rs

subako_function_sdk/
package.rs
