Skip to main content

vv_agent/tools/
function.rs

1use std::future::Future;
2use std::marker::PhantomData;
3use std::sync::Arc;
4
5use serde::de::DeserializeOwned;
6use serde_json::Value;
7
8use crate::context::{RunContext, ToolCallContext};
9use crate::tools::{Tool, ToolContext, ToolOutput, ToolSpec};
10use crate::types::ToolArguments;
11
12type DynamicFunctionHandler =
13    Arc<dyn Fn(ToolCallContext, Value) -> Result<ToolOutput, String> + Send + Sync>;
14
15pub struct FunctionTool<Args = Value> {
16    name: String,
17    description: String,
18    parameters_schema: Value,
19    handler: DynamicFunctionHandler,
20    _args: PhantomData<fn() -> Args>,
21}
22
23impl FunctionTool<Value> {
24    pub fn builder(name: impl Into<String>) -> FunctionToolBuilder<Value> {
25        FunctionToolBuilder {
26            name: name.into(),
27            description: String::new(),
28            parameters_schema: serde_json::json!({
29                "type": "object",
30                "properties": {},
31                "required": []
32            }),
33            handler: None,
34            _args: PhantomData,
35        }
36    }
37}
38
39impl<Args> Clone for FunctionTool<Args> {
40    fn clone(&self) -> Self {
41        Self {
42            name: self.name.clone(),
43            description: self.description.clone(),
44            parameters_schema: self.parameters_schema.clone(),
45            handler: self.handler.clone(),
46            _args: PhantomData,
47        }
48    }
49}
50
51impl<Args> Tool for FunctionTool<Args>
52where
53    Args: Send + Sync + 'static,
54{
55    fn name(&self) -> &str {
56        &self.name
57    }
58
59    fn description(&self) -> &str {
60        &self.description
61    }
62
63    fn parameters_schema(&self) -> &Value {
64        &self.parameters_schema
65    }
66
67    fn as_tool_spec(&self) -> ToolSpec {
68        let name = self.name.clone();
69        let description = self.description.clone();
70        let parameters_schema = self.parameters_schema.clone();
71        let handler = self.handler.clone();
72        let mut spec = ToolSpec::new(
73            name.clone(),
74            description.clone(),
75            Arc::new(
76                move |context: &mut ToolContext, arguments: &ToolArguments| {
77                    let raw_arguments = Value::Object(arguments.clone().into_iter().collect());
78                    let call_context = ToolCallContext {
79                        run: RunContext {
80                            run_id: context.task_id.clone(),
81                            agent_name: context
82                                .metadata
83                                .get("agent_name")
84                                .and_then(Value::as_str)
85                                .unwrap_or_default()
86                                .to_string(),
87                            model: None,
88                            workspace: Some(context.workspace.clone()),
89                            metadata: context.metadata.clone(),
90                        },
91                        tool_call_id: String::new(),
92                        tool_name: name.clone(),
93                        raw_arguments: raw_arguments.clone(),
94                        metadata: context.metadata.clone(),
95                        app_state: None,
96                    };
97                    match handler(call_context, raw_arguments) {
98                        Ok(output) => output.to_result(""),
99                        Err(error) => ToolOutput::error(error).to_result(""),
100                    }
101                },
102            ),
103        );
104        spec.schema = serde_json::json!({
105            "type": "function",
106            "function": {
107                "name": self.name,
108                "description": self.description,
109                "parameters": parameters_schema,
110            }
111        });
112        spec
113    }
114}
115
116pub struct FunctionToolBuilder<Args = Value> {
117    name: String,
118    description: String,
119    parameters_schema: Value,
120    handler: Option<DynamicFunctionHandler>,
121    _args: PhantomData<fn() -> Args>,
122}
123
124impl<Args> FunctionToolBuilder<Args> {
125    pub fn description(mut self, description: impl Into<String>) -> Self {
126        self.description = description.into();
127        self
128    }
129
130    pub fn json_schema(mut self, schema: Value) -> Self {
131        self.parameters_schema = schema;
132        self
133    }
134
135    pub fn handler<NextArgs, F, Fut>(self, handler: F) -> FunctionToolBuilder<NextArgs>
136    where
137        NextArgs: DeserializeOwned + Send + Sync + 'static,
138        F: Fn(ToolCallContext, NextArgs) -> Fut + Send + Sync + 'static,
139        Fut: Future<Output = Result<ToolOutput, String>> + Send + 'static,
140    {
141        let handler = Arc::new(handler);
142        FunctionToolBuilder {
143            name: self.name,
144            description: self.description,
145            parameters_schema: self.parameters_schema,
146            handler: Some(Arc::new(move |context, raw_arguments| {
147                let args = serde_json::from_value::<NextArgs>(raw_arguments)
148                    .map_err(|error| format!("invalid tool arguments: {error}"))?;
149                block_on_tool_future(handler(context, args))
150            })),
151            _args: PhantomData,
152        }
153    }
154
155    pub fn build(self) -> Result<FunctionTool<Args>, String> {
156        if self.name.trim().is_empty() {
157            return Err("tool name cannot be empty".to_string());
158        }
159        let Some(handler) = self.handler else {
160            return Err("tool handler is required".to_string());
161        };
162        Ok(FunctionTool {
163            name: self.name,
164            description: self.description,
165            parameters_schema: self.parameters_schema,
166            handler,
167            _args: PhantomData,
168        })
169    }
170}
171
172fn block_on_tool_future<F>(future: F) -> Result<ToolOutput, String>
173where
174    F: Future<Output = Result<ToolOutput, String>> + Send + 'static,
175{
176    if let Ok(handle) = tokio::runtime::Handle::try_current() {
177        if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
178            tokio::task::block_in_place(|| handle.block_on(future))
179        } else {
180            std::thread::spawn(move || {
181                tokio::runtime::Builder::new_current_thread()
182                    .enable_all()
183                    .build()
184                    .map_err(|error| error.to_string())?
185                    .block_on(future)
186            })
187            .join()
188            .map_err(|_| "tool handler thread panicked".to_string())?
189        }
190    } else {
191        tokio::runtime::Builder::new_current_thread()
192            .enable_all()
193            .build()
194            .map_err(|error| error.to_string())?
195            .block_on(future)
196    }
197}