rskit_tool/
from_fn.rs

1//! Auto-wiring: create tools from async functions.
2
3use async_trait::async_trait;
4use rskit_errors::{AppError, AppResult, ErrorCode};
5use rskit_schema::{CompiledSchema, ValidationResult};
6use schemars::JsonSchema;
7use serde::{Serialize, de::DeserializeOwned};
8
9use crate::callable::Callable;
10use crate::context::Context;
11use crate::definition::Definition;
12use crate::io::{ToolInput, ToolMetadata, ToolOutput, ToolSchema};
13use crate::result::ToolResult;
14
15/// Create a tool from an async handler function.
16///
17/// Generates JSON Schemas from the input and output type parameters using schemars,
18/// and wraps the handler in a type-erased `Callable`.
19///
20/// The handler receives a `Context` reference and deserialized input, returning a `ToolResult`.
21///
22/// # Example
23///
24/// ```rust,ignore
25/// use rskit_tool::from_fn;
26///
27/// let tool = from_fn("search", "Search the web", |_ctx, input: SearchInput| async move {
28///     Ok(rskit_tool::result::text_result(&format!("found: {}", input.query)))
29/// })?;
30/// ```
31pub fn from_fn<I, F, Fut>(name: &str, description: &str, handler: F) -> AppResult<Box<dyn Callable>>
32where
33    I: DeserializeOwned + JsonSchema + Send + 'static,
34    F: Fn(Context, I) -> Fut + Send + Sync + 'static,
35    Fut: Future<Output = AppResult<ToolResult>> + Send + 'static,
36{
37    let input_schema = rskit_schema::generate::<I>()?;
38    let input_validator = rskit_schema::compile(&input_schema)?;
39
40    let def = Definition {
41        name: name.to_string(),
42        description: description.to_string(),
43        input_schema: ToolSchema::new(input_schema)?,
44        output_schema: None,
45        annotations: crate::Annotations::default(),
46        envelope: crate::Envelope::default(),
47    };
48
49    Ok(Box::new(FnTool {
50        definition: def,
51        input_validator,
52        handler,
53        _phantom: std::marker::PhantomData,
54    }))
55}
56
57struct FnTool<F, I> {
58    definition: Definition,
59    input_validator: CompiledSchema,
60    handler: F,
61    _phantom: std::marker::PhantomData<fn(I)>,
62}
63
64#[async_trait]
65impl<I, F, Fut> Callable for FnTool<F, I>
66where
67    I: DeserializeOwned + Send + 'static,
68    F: Fn(Context, I) -> Fut + Send + Sync + 'static,
69    Fut: Future<Output = AppResult<ToolResult>> + Send + 'static,
70{
71    fn definition(&self) -> &Definition {
72        &self.definition
73    }
74
75    fn validate(&self, input: &ToolInput) -> ValidationResult {
76        self.input_validator.validate(input.as_json())
77    }
78
79    async fn call(&self, ctx: &Context, input: ToolInput) -> AppResult<ToolResult> {
80        let typed_input: I = serde_json::from_value(input.into_json()).map_err(|e| {
81            AppError::new(ErrorCode::InvalidInput, format!("invalid tool input: {e}"))
82        })?;
83
84        (self.handler)(ctx.clone(), typed_input).await
85    }
86}
87
88/// Convenience: create a tool from a simpler handler that returns a Serialize type.
89///
90/// The handler takes only the input (no Context) and returns `AppResult<O>`.
91/// The output is automatically serialized to JSON.
92pub fn from_fn_simple<I, O, F, Fut>(
93    name: &str,
94    description: &str,
95    handler: F,
96) -> AppResult<Box<dyn Callable>>
97where
98    I: DeserializeOwned + JsonSchema + Send + 'static,
99    O: Serialize + Send + 'static,
100    F: Fn(I) -> Fut + Send + Sync + 'static,
101    Fut: Future<Output = AppResult<O>> + Send + 'static,
102{
103    let input_schema = rskit_schema::generate::<I>()?;
104    let input_validator = rskit_schema::compile(&input_schema)?;
105
106    let def = Definition {
107        name: name.to_string(),
108        description: description.to_string(),
109        input_schema: ToolSchema::new(input_schema)?,
110        output_schema: None,
111        annotations: crate::Annotations::default(),
112        envelope: crate::Envelope::default(),
113    };
114
115    Ok(Box::new(SimpleFnTool {
116        definition: def,
117        input_validator,
118        handler,
119        _phantom: std::marker::PhantomData,
120    }))
121}
122
123struct SimpleFnTool<F, I, O> {
124    definition: Definition,
125    input_validator: CompiledSchema,
126    handler: F,
127    _phantom: std::marker::PhantomData<fn(I) -> O>,
128}
129
130#[async_trait]
131impl<I, O, F, Fut> Callable for SimpleFnTool<F, I, O>
132where
133    I: DeserializeOwned + Send + 'static,
134    O: Serialize + Send + 'static,
135    F: Fn(I) -> Fut + Send + Sync + 'static,
136    Fut: Future<Output = AppResult<O>> + Send + 'static,
137{
138    fn definition(&self) -> &Definition {
139        &self.definition
140    }
141
142    fn validate(&self, input: &ToolInput) -> ValidationResult {
143        self.input_validator.validate(input.as_json())
144    }
145
146    async fn call(&self, _ctx: &Context, input: ToolInput) -> AppResult<ToolResult> {
147        let typed_input: I = serde_json::from_value(input.into_json()).map_err(|e| {
148            AppError::new(ErrorCode::InvalidInput, format!("invalid tool input: {e}"))
149        })?;
150
151        let output = (self.handler)(typed_input).await?;
152
153        let json = ToolOutput::from_serializable(&output)?;
154        let content = serde_json::to_string(&output).map_err(|e| {
155            AppError::new(
156                ErrorCode::Internal,
157                format!("failed to serialize tool output content: {e}"),
158            )
159            .with_cause(e)
160        })?;
161
162        Ok(ToolResult {
163            output: Some(json),
164            content,
165            is_error: false,
166            metadata: ToolMetadata::new(),
167        })
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    use schemars::JsonSchema;
175    use serde::{Deserialize, Serialize};
176
177    #[derive(Deserialize, JsonSchema)]
178    struct Input {
179        value: i32,
180    }
181
182    #[derive(Serialize)]
183    struct Output {
184        value: i32,
185    }
186
187    #[tokio::test]
188    async fn from_fn_exposes_definition_validate_and_call() {
189        let tool = from_fn(
190            "double",
191            "Double",
192            |_ctx: Context, input: Input| async move {
193                Ok(crate::text_result(&(input.value * 2).to_string()))
194            },
195        )
196        .expect("tool should build");
197
198        assert_eq!(tool.definition().name, "double");
199        assert!(
200            tool.validate(&ToolInput::new(serde_json::json!({"value": 2})).unwrap())
201                .valid
202        );
203        let result = tool
204            .call(
205                &Context::new(),
206                ToolInput::new(serde_json::json!({"value": 3})).unwrap(),
207            )
208            .await
209            .expect("call should succeed");
210        assert_eq!(result.text(), "6");
211    }
212
213    #[tokio::test]
214    async fn from_fn_rejects_deserialisation_errors() {
215        let tool = from_fn(
216            "double",
217            "Double",
218            |_ctx: Context, _input: Input| async move { Ok(crate::text_result("unused")) },
219        )
220        .expect("tool should build");
221
222        let err = tool
223            .call(
224                &Context::new(),
225                ToolInput::new(serde_json::json!({"value": "bad"})).unwrap(),
226            )
227            .await
228            .expect_err("invalid input should fail");
229        assert_eq!(err.code(), ErrorCode::InvalidInput);
230    }
231
232    #[tokio::test]
233    async fn from_fn_simple_exposes_definition_validate_and_serialises_output() {
234        let tool = from_fn_simple("double", "Double", |input: Input| async move {
235            Ok(Output {
236                value: input.value * 2,
237            })
238        })
239        .expect("tool should build");
240
241        assert_eq!(tool.definition().description, "Double");
242        assert!(
243            tool.validate(&ToolInput::new(serde_json::json!({"value": 1})).unwrap())
244                .valid
245        );
246        let result = tool
247            .call(
248                &Context::new(),
249                ToolInput::new(serde_json::json!({"value": 4})).unwrap(),
250            )
251            .await
252            .expect("call should succeed");
253        assert_eq!(result.output.as_ref().expect("output")["value"], 8);
254        assert_eq!(result.content, r#"{"value":8}"#);
255    }
256
257    #[tokio::test]
258    async fn from_fn_simple_rejects_deserialisation_errors() {
259        let tool = from_fn_simple("double", "Double", |input: Input| async move {
260            Ok(Output { value: input.value })
261        })
262        .expect("tool should build");
263
264        let err = tool
265            .call(
266                &Context::new(),
267                ToolInput::new(serde_json::json!({"value": []})).unwrap(),
268            )
269            .await
270            .expect_err("invalid input should fail");
271        assert_eq!(err.code(), ErrorCode::InvalidInput);
272    }
273}
rskit_tool/from_fn.rs

rskit_tool/
from_fn.rs
