Skip to main content

vv_agent/tools/
function.rs

1use std::future::Future;
2use std::marker::PhantomData;
3use std::panic::{catch_unwind, AssertUnwindSafe};
4use std::sync::Arc;
5use std::sync::Mutex;
6use std::time::Duration;
7
8use serde::de::DeserializeOwned;
9use serde_json::Value;
10
11use crate::context::{RunContext, ToolCallContext};
12use crate::tools::{
13    Tool, ToolApprovalRule, ToolContext, ToolEnablementContext, ToolEnablementRule, ToolExposure,
14    ToolMetadata, ToolOutput, ToolSpec, ToolSpecExecutor,
15};
16use crate::types::{Metadata, ToolArguments};
17
18type DynamicFunctionHandler =
19    Arc<dyn Fn(ToolCallContext, Value) -> Result<ToolOutput, String> + Send + Sync>;
20pub type ToolErrorMapper = Arc<dyn Fn(&str) -> String + Send + Sync + 'static>;
21
22pub struct FunctionTool<Args = Value> {
23    name: String,
24    description: String,
25    parameters_schema: Value,
26    handler: DynamicFunctionHandler,
27    strict_schema: bool,
28    exposure: ToolExposure,
29    timeout: Option<Duration>,
30    approval: ToolApprovalRule,
31    tool_metadata: Option<ToolMetadata>,
32    enablement: ToolEnablementRule,
33    error_mapper: Option<ToolErrorMapper>,
34    metadata: Metadata,
35    _args: PhantomData<fn() -> Args>,
36}
37
38impl FunctionTool<Value> {
39    pub fn builder(name: impl Into<String>) -> FunctionToolBuilder<Value> {
40        FunctionToolBuilder {
41            name: name.into(),
42            description: String::new(),
43            parameters_schema: serde_json::json!({
44                "type": "object",
45                "properties": {},
46                "required": []
47            }),
48            handler: None,
49            strict_schema: true,
50            exposure: ToolExposure::Direct,
51            timeout: None,
52            approval: ToolApprovalRule::default(),
53            tool_metadata: None,
54            enablement: ToolEnablementRule::default(),
55            error_mapper: None,
56            metadata: Metadata::new(),
57            _args: PhantomData,
58        }
59    }
60}
61
62impl<Args> Clone for FunctionTool<Args> {
63    fn clone(&self) -> Self {
64        Self {
65            name: self.name.clone(),
66            description: self.description.clone(),
67            parameters_schema: self.parameters_schema.clone(),
68            handler: self.handler.clone(),
69            strict_schema: self.strict_schema,
70            exposure: self.exposure,
71            timeout: self.timeout,
72            approval: self.approval.clone(),
73            tool_metadata: self.tool_metadata.clone(),
74            enablement: self.enablement.clone(),
75            error_mapper: self.error_mapper.clone(),
76            metadata: self.metadata.clone(),
77            _args: PhantomData,
78        }
79    }
80}
81
82impl<Args> FunctionTool<Args> {
83    pub fn metadata(&self) -> &Metadata {
84        &self.metadata
85    }
86}
87
88impl<Args> FunctionTool<Args>
89where
90    Args: Send + Sync + 'static,
91{
92    pub fn to_executor(&self) -> std::sync::Arc<dyn crate::tools::ToolExecutor> {
93        ToolSpecExecutor::new(self.as_tool_spec()).into_arc()
94    }
95}
96
97impl<Args> Tool for FunctionTool<Args>
98where
99    Args: Send + Sync + 'static,
100{
101    fn name(&self) -> &str {
102        &self.name
103    }
104
105    fn description(&self) -> &str {
106        &self.description
107    }
108
109    fn parameters_schema(&self) -> &Value {
110        &self.parameters_schema
111    }
112
113    fn strict_schema(&self) -> bool {
114        self.strict_schema
115    }
116
117    fn exposure(&self) -> ToolExposure {
118        self.exposure
119    }
120
121    fn timeout(&self) -> Option<Duration> {
122        self.timeout
123    }
124
125    fn approval_rule(&self) -> ToolApprovalRule {
126        self.approval.clone()
127    }
128
129    fn tool_metadata(&self) -> Option<&ToolMetadata> {
130        self.tool_metadata.as_ref()
131    }
132
133    fn is_enabled(&self, context: &ToolEnablementContext) -> bool {
134        self.enablement.is_enabled(context)
135    }
136
137    fn as_tool_spec(&self) -> ToolSpec {
138        let name = self.name.clone();
139        let description = self.description.clone();
140        let parameters_schema = self.parameters_schema.clone();
141        let handler = self.handler.clone();
142        let error_mapper = self.error_mapper.clone();
143        let mut spec = ToolSpec::new(
144            name.clone(),
145            description.clone(),
146            Arc::new(
147                move |context: &mut ToolContext, arguments: &ToolArguments| {
148                    let raw_arguments = Value::Object(arguments.clone().into_iter().collect());
149                    let shared_state = Arc::new(Mutex::new(context.shared_state.clone()));
150                    let run_context = context.run_context.clone().unwrap_or_else(|| RunContext {
151                        run_id: context.task_id.clone(),
152                        agent_name: context
153                            .metadata
154                            .get("agent_name")
155                            .and_then(Value::as_str)
156                            .unwrap_or_default()
157                            .to_string(),
158                        model: None,
159                        workspace: Some(context.workspace.clone()),
160                        metadata: context.metadata.clone(),
161                        app_state: context.app_state.clone(),
162                    });
163                    let call_context = ToolCallContext {
164                        run: run_context,
165                        tool_call_id: context.tool_call_id.clone(),
166                        tool_name: context.tool_name.clone(),
167                        raw_arguments: raw_arguments.clone(),
168                        idempotency_key: context.idempotency_key.clone(),
169                        metadata: context.metadata.clone(),
170                        app_state: context.app_state.clone(),
171                        shared_state: shared_state.clone(),
172                    };
173                    let outcome =
174                        catch_unwind(AssertUnwindSafe(|| handler(call_context, raw_arguments)));
175                    context.shared_state = shared_state
176                        .lock()
177                        .unwrap_or_else(std::sync::PoisonError::into_inner)
178                        .clone();
179                    match outcome {
180                        Ok(Ok(output)) => output.to_result(&context.tool_call_id),
181                        Ok(Err(error)) => function_error_result(
182                            &name,
183                            &context.tool_call_id,
184                            &error,
185                            error_mapper.as_ref(),
186                        ),
187                        Err(payload) => function_error_result(
188                            &name,
189                            &context.tool_call_id,
190                            &panic_message(payload),
191                            error_mapper.as_ref(),
192                        ),
193                    }
194                },
195            ),
196        );
197        spec.schema = serde_json::json!({
198            "type": "function",
199            "function": {
200                "name": self.name,
201                "description": self.description,
202                "parameters": parameters_schema,
203                "strict": self.strict_schema,
204            }
205        });
206        spec.strict_schema = self.strict_schema;
207        spec.exposure = self.exposure;
208        spec.timeout = self.timeout;
209        spec.approval = self.approval.clone();
210        spec.tool_metadata = self.tool_metadata.clone();
211        spec.metadata = self.metadata.clone();
212        spec
213    }
214}
215
216pub struct FunctionToolBuilder<Args = Value> {
217    name: String,
218    description: String,
219    parameters_schema: Value,
220    handler: Option<DynamicFunctionHandler>,
221    strict_schema: bool,
222    exposure: ToolExposure,
223    timeout: Option<Duration>,
224    approval: ToolApprovalRule,
225    tool_metadata: Option<ToolMetadata>,
226    enablement: ToolEnablementRule,
227    error_mapper: Option<ToolErrorMapper>,
228    metadata: Metadata,
229    _args: PhantomData<fn() -> Args>,
230}
231
232impl<Args> FunctionToolBuilder<Args> {
233    pub fn description(mut self, description: impl Into<String>) -> Self {
234        self.description = description.into();
235        self
236    }
237
238    pub fn json_schema(mut self, schema: Value) -> Self {
239        self.parameters_schema = schema;
240        self
241    }
242
243    pub fn strict_schema(mut self, strict_schema: bool) -> Self {
244        self.strict_schema = strict_schema;
245        self
246    }
247
248    pub fn exposure(mut self, exposure: ToolExposure) -> Self {
249        self.exposure = exposure;
250        self
251    }
252
253    pub fn timeout(mut self, timeout: Duration) -> Self {
254        self.timeout = Some(timeout);
255        self
256    }
257
258    pub fn needs_approval(mut self, approval: impl Into<ToolApprovalRule>) -> Self {
259        self.approval = approval.into();
260        self
261    }
262
263    pub fn needs_approval_if<F>(mut self, predicate: F) -> Self
264    where
265        F: Fn(&ToolContext, &ToolArguments) -> bool + Send + Sync + 'static,
266    {
267        self.approval = ToolApprovalRule::predicate(predicate);
268        self
269    }
270
271    pub fn tool_metadata(mut self, tool_metadata: ToolMetadata) -> Self {
272        self.tool_metadata = Some(tool_metadata);
273        self
274    }
275
276    pub fn enabled(mut self, enabled: bool) -> Self {
277        self.enablement = ToolEnablementRule::Static(enabled);
278        self
279    }
280
281    pub fn enabled_if<F>(mut self, predicate: F) -> Self
282    where
283        F: Fn(&ToolEnablementContext) -> bool + Send + Sync + 'static,
284    {
285        self.enablement = ToolEnablementRule::predicate(predicate);
286        self
287    }
288
289    pub fn failure_error_function<F>(mut self, mapper: F) -> Self
290    where
291        F: Fn(&str) -> String + Send + Sync + 'static,
292    {
293        self.error_mapper = Some(Arc::new(mapper));
294        self
295    }
296
297    pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
298        self.metadata.insert(key.into(), value);
299        self
300    }
301
302    pub fn handler<NextArgs, F, Fut>(self, handler: F) -> FunctionToolBuilder<NextArgs>
303    where
304        NextArgs: DeserializeOwned + Send + Sync + 'static,
305        F: Fn(ToolCallContext, NextArgs) -> Fut + Send + Sync + 'static,
306        Fut: Future<Output = Result<ToolOutput, String>> + Send + 'static,
307    {
308        let handler = Arc::new(handler);
309        FunctionToolBuilder {
310            name: self.name,
311            description: self.description,
312            parameters_schema: self.parameters_schema,
313            handler: Some(Arc::new(move |context, raw_arguments| {
314                let args = serde_json::from_value::<NextArgs>(raw_arguments)
315                    .map_err(|error| format!("invalid tool arguments: {error}"))?;
316                block_on_tool_future(handler(context, args))
317            })),
318            strict_schema: self.strict_schema,
319            exposure: self.exposure,
320            timeout: self.timeout,
321            approval: self.approval,
322            tool_metadata: self.tool_metadata,
323            enablement: self.enablement,
324            error_mapper: self.error_mapper,
325            metadata: self.metadata,
326            _args: PhantomData,
327        }
328    }
329
330    pub fn build(self) -> Result<FunctionTool<Args>, String> {
331        if self.name.trim().is_empty() {
332            return Err("tool name cannot be empty".to_string());
333        }
334        if self.timeout.is_some_and(|timeout| timeout.is_zero()) {
335            return Err("tool timeout must be greater than zero".to_string());
336        }
337        let Some(handler) = self.handler else {
338            return Err("tool handler is required".to_string());
339        };
340        let tool_metadata = self
341            .tool_metadata
342            .as_ref()
343            .map(ToolMetadata::normalized)
344            .transpose()
345            .map_err(|error| error.to_string())?;
346        let parameters_schema =
347            crate::tools::argument_validation::close_object_schemas(&self.parameters_schema);
348        crate::tools::argument_validation::validator_for_parameters(&parameters_schema)?;
349        Ok(FunctionTool {
350            name: self.name,
351            description: self.description,
352            parameters_schema,
353            handler,
354            strict_schema: self.strict_schema,
355            exposure: self.exposure,
356            timeout: self.timeout,
357            approval: self.approval,
358            tool_metadata,
359            enablement: self.enablement,
360            error_mapper: self.error_mapper,
361            metadata: self.metadata,
362            _args: PhantomData,
363        })
364    }
365}
366
367fn function_error_result(
368    tool_name: &str,
369    tool_call_id: &str,
370    error: &str,
371    mapper: Option<&ToolErrorMapper>,
372) -> crate::types::ToolExecutionResult {
373    let message = mapper
374        .map(|mapper| mapper(error))
375        .unwrap_or_else(|| format!("Tool execution failed ({tool_name}): {error}"));
376    ToolOutput::error(message)
377        .with_code("tool_execution_failed")
378        .to_result(tool_call_id)
379}
380
381fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
382    if let Some(message) = payload.downcast_ref::<&str>() {
383        return (*message).to_string();
384    }
385    if let Some(message) = payload.downcast_ref::<String>() {
386        return message.clone();
387    }
388    "tool handler panicked".to_string()
389}
390
391fn block_on_tool_future<F>(future: F) -> Result<ToolOutput, String>
392where
393    F: Future<Output = Result<ToolOutput, String>> + Send + 'static,
394{
395    if let Ok(handle) = tokio::runtime::Handle::try_current() {
396        if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
397            tokio::task::block_in_place(|| handle.block_on(future))
398        } else {
399            std::thread::spawn(move || {
400                tokio::runtime::Builder::new_current_thread()
401                    .enable_all()
402                    .build()
403                    .map_err(|error| error.to_string())?
404                    .block_on(future)
405            })
406            .join()
407            .map_err(|_| "tool handler thread panicked".to_string())?
408        }
409    } else {
410        tokio::runtime::Builder::new_current_thread()
411            .enable_all()
412            .build()
413            .map_err(|error| error.to_string())?
414            .block_on(future)
415    }
416}