1use std::{future::Future, marker::PhantomData};
2
3use runifold_core::{CapabilityId, EffectClass, RiskLevel};
4use schemars::{JsonSchema, schema_for};
5use serde::{Serialize, de::DeserializeOwned};
6
7use crate::{Tool, ToolContext, ToolDescriptor, ToolError, ToolErrorKind, ToolFuture, ToolOutput};
8
9pub struct FunctionTool<Input, Output, Handler> {
12 descriptor: ToolDescriptor,
13 handler: Handler,
14 types: PhantomData<fn(Input) -> Output>,
15}
16
17impl<Input, Output, Handler> FunctionTool<Input, Output, Handler>
18where
19 Input: JsonSchema,
20 Output: JsonSchema,
21{
22 pub fn new(name: impl Into<String>, description: impl Into<String>, handler: Handler) -> Self {
28 Self {
29 descriptor: ToolDescriptor {
30 id: CapabilityId::new(),
31 name: name.into(),
32 version: "1".into(),
33 description: description.into(),
34 input_schema: schema_for!(Input).to_value(),
35 output_schema: schema_for!(Output).to_value(),
36 effect: EffectClass::Pure,
37 risk: RiskLevel::Low,
38 metadata: std::collections::BTreeMap::new(),
39 },
40 handler,
41 types: PhantomData,
42 }
43 }
44
45 #[must_use]
47 pub const fn capability_id(mut self, id: CapabilityId) -> Self {
48 self.descriptor.id = id;
49 self
50 }
51
52 #[must_use]
54 pub fn version(mut self, version: impl Into<String>) -> Self {
55 self.descriptor.version = version.into();
56 self
57 }
58
59 #[must_use]
61 pub const fn effect(mut self, effect: EffectClass) -> Self {
62 self.descriptor.effect = effect;
63 self
64 }
65
66 #[must_use]
68 pub const fn risk(mut self, risk: RiskLevel) -> Self {
69 self.descriptor.risk = risk;
70 self
71 }
72
73 #[must_use]
75 pub fn metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
76 self.descriptor.metadata.insert(key.into(), value);
77 self
78 }
79}
80
81impl<Input, Output, Handler> std::fmt::Debug for FunctionTool<Input, Output, Handler> {
82 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83 formatter
84 .debug_struct("FunctionTool")
85 .field("descriptor", &self.descriptor)
86 .finish_non_exhaustive()
87 }
88}
89
90impl<Input, Output, Handler, HandlerFuture> Tool for FunctionTool<Input, Output, Handler>
91where
92 Input: DeserializeOwned + JsonSchema + Send + 'static,
93 Output: JsonSchema + Serialize + Send + 'static,
94 Handler: Fn(Input, ToolContext) -> HandlerFuture + Send + Sync,
95 HandlerFuture: Future<Output = Result<Output, ToolError>> + Send + 'static,
96{
97 fn descriptor(&self) -> &ToolDescriptor {
98 &self.descriptor
99 }
100
101 fn invoke(
102 &self,
103 input: serde_json::Value,
104 context: ToolContext,
105 ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
106 let input = match serde_json::from_value(input) {
107 Ok(input) => input,
108 Err(error) => {
109 return Box::pin(async move {
110 Err(ToolError::local(
111 ToolErrorKind::InvalidInput,
112 format!("typed Tool input is invalid: {error}"),
113 ))
114 });
115 }
116 };
117 let future = (self.handler)(input, context);
118 Box::pin(async move {
119 let output = future.await?;
120 let value = serde_json::to_value(output).map_err(|error| {
121 ToolError::local(
122 ToolErrorKind::InvalidOutput,
123 format!("typed Tool output cannot be serialized: {error}"),
124 )
125 })?;
126 Ok(ToolOutput::model_visible(value))
127 })
128 }
129}
130
131#[cfg(test)]
132mod tests {
133 use std::sync::{
134 Arc,
135 atomic::{AtomicUsize, Ordering},
136 };
137
138 use runifold_core::{Budget, BudgetTracker, CapabilitySet, RunContext};
139 use schemars::JsonSchema;
140 use serde::{Deserialize, Serialize};
141 use serde_json::json;
142
143 use super::FunctionTool;
144 use crate::{Tool, ToolErrorKind, ToolRegistry};
145
146 #[derive(Deserialize, JsonSchema)]
147 struct AddInput {
148 left: i64,
149 right: i64,
150 }
151
152 #[derive(JsonSchema, Serialize)]
153 struct AddOutput {
154 sum: i64,
155 }
156
157 #[test]
158 fn typed_function_generates_schemas_and_runs_through_registry() {
159 let calls = Arc::new(AtomicUsize::new(0));
160 let observed = calls.clone();
161 let tool = Arc::new(FunctionTool::new(
162 "add",
163 "adds two integers",
164 move |input: AddInput, _context| {
165 let observed = observed.clone();
166 async move {
167 observed.fetch_add(1, Ordering::SeqCst);
168 Ok(AddOutput {
169 sum: input.left + input.right,
170 })
171 }
172 },
173 ));
174 let descriptor = tool.descriptor();
175 assert_eq!(
176 descriptor.input_schema["required"],
177 json!(["left", "right"])
178 );
179 assert_eq!(descriptor.output_schema["required"], json!(["sum"]));
180
181 let mut capabilities = CapabilitySet::new();
182 capabilities.grant(descriptor.capability());
183 let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
184 let mut registry = ToolRegistry::new();
185 registry.register(tool).unwrap();
186
187 let output = futures_executor::block_on(registry.invoke(
188 "add",
189 json!({"left": 2, "right": 3}),
190 &run,
191 ))
192 .unwrap();
193
194 assert_eq!(output.value, json!({"sum": 5}));
195 assert_eq!(calls.load(Ordering::SeqCst), 1);
196 }
197
198 #[test]
199 fn invalid_typed_input_never_calls_handler() {
200 let calls = Arc::new(AtomicUsize::new(0));
201 let observed = calls.clone();
202 let tool = Arc::new(FunctionTool::new(
203 "add",
204 "adds two integers",
205 move |_input: AddInput, _context| {
206 let observed = observed.clone();
207 async move {
208 observed.fetch_add(1, Ordering::SeqCst);
209 Ok(AddOutput { sum: 0 })
210 }
211 },
212 ));
213 let mut capabilities = CapabilitySet::new();
214 capabilities.grant(tool.descriptor().capability());
215 let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
216 let mut registry = ToolRegistry::new();
217 registry.register(tool).unwrap();
218
219 let error = futures_executor::block_on(registry.invoke("add", json!({"left": 2}), &run))
220 .unwrap_err();
221
222 assert_eq!(error.kind, ToolErrorKind::InvalidInput);
223 assert_eq!(calls.load(Ordering::SeqCst), 0);
224 }
225}