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, ToolSpecExecutor};
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> FunctionTool<Args>
52where
53    Args: Send + Sync + 'static,
54{
55    pub fn to_executor(&self) -> std::sync::Arc<dyn crate::tools::ToolExecutor> {
56        ToolSpecExecutor::new(self.as_tool_spec()).into_arc()
57    }
58}
59
60impl<Args> Tool for FunctionTool<Args>
61where
62    Args: Send + Sync + 'static,
63{
64    fn name(&self) -> &str {
65        &self.name
66    }
67
68    fn description(&self) -> &str {
69        &self.description
70    }
71
72    fn parameters_schema(&self) -> &Value {
73        &self.parameters_schema
74    }
75
76    fn as_tool_spec(&self) -> ToolSpec {
77        let name = self.name.clone();
78        let description = self.description.clone();
79        let parameters_schema = self.parameters_schema.clone();
80        let handler = self.handler.clone();
81        let mut spec = ToolSpec::new(
82            name.clone(),
83            description.clone(),
84            Arc::new(
85                move |context: &mut ToolContext, arguments: &ToolArguments| {
86                    let raw_arguments = Value::Object(arguments.clone().into_iter().collect());
87                    let call_context = ToolCallContext {
88                        run: RunContext {
89                            run_id: context.task_id.clone(),
90                            agent_name: context
91                                .metadata
92                                .get("agent_name")
93                                .and_then(Value::as_str)
94                                .unwrap_or_default()
95                                .to_string(),
96                            model: None,
97                            workspace: Some(context.workspace.clone()),
98                            metadata: context.metadata.clone(),
99                        },
100                        tool_call_id: String::new(),
101                        tool_name: name.clone(),
102                        raw_arguments: raw_arguments.clone(),
103                        metadata: context.metadata.clone(),
104                        app_state: None,
105                    };
106                    match handler(call_context, raw_arguments) {
107                        Ok(output) => output.to_result(""),
108                        Err(error) => ToolOutput::error(error).to_result(""),
109                    }
110                },
111            ),
112        );
113        spec.schema = serde_json::json!({
114            "type": "function",
115            "function": {
116                "name": self.name,
117                "description": self.description,
118                "parameters": parameters_schema,
119            }
120        });
121        spec
122    }
123}
124
125pub struct FunctionToolBuilder<Args = Value> {
126    name: String,
127    description: String,
128    parameters_schema: Value,
129    handler: Option<DynamicFunctionHandler>,
130    _args: PhantomData<fn() -> Args>,
131}
132
133impl<Args> FunctionToolBuilder<Args> {
134    pub fn description(mut self, description: impl Into<String>) -> Self {
135        self.description = description.into();
136        self
137    }
138
139    pub fn json_schema(mut self, schema: Value) -> Self {
140        self.parameters_schema = schema;
141        self
142    }
143
144    pub fn handler<NextArgs, F, Fut>(self, handler: F) -> FunctionToolBuilder<NextArgs>
145    where
146        NextArgs: DeserializeOwned + Send + Sync + 'static,
147        F: Fn(ToolCallContext, NextArgs) -> Fut + Send + Sync + 'static,
148        Fut: Future<Output = Result<ToolOutput, String>> + Send + 'static,
149    {
150        let handler = Arc::new(handler);
151        FunctionToolBuilder {
152            name: self.name,
153            description: self.description,
154            parameters_schema: self.parameters_schema,
155            handler: Some(Arc::new(move |context, raw_arguments| {
156                let args = serde_json::from_value::<NextArgs>(raw_arguments)
157                    .map_err(|error| format!("invalid tool arguments: {error}"))?;
158                block_on_tool_future(handler(context, args))
159            })),
160            _args: PhantomData,
161        }
162    }
163
164    pub fn build(self) -> Result<FunctionTool<Args>, String> {
165        if self.name.trim().is_empty() {
166            return Err("tool name cannot be empty".to_string());
167        }
168        let Some(handler) = self.handler else {
169            return Err("tool handler is required".to_string());
170        };
171        Ok(FunctionTool {
172            name: self.name,
173            description: self.description,
174            parameters_schema: self.parameters_schema,
175            handler,
176            _args: PhantomData,
177        })
178    }
179}
180
181fn block_on_tool_future<F>(future: F) -> Result<ToolOutput, String>
182where
183    F: Future<Output = Result<ToolOutput, String>> + Send + 'static,
184{
185    if let Ok(handle) = tokio::runtime::Handle::try_current() {
186        if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
187            tokio::task::block_in_place(|| handle.block_on(future))
188        } else {
189            std::thread::spawn(move || {
190                tokio::runtime::Builder::new_current_thread()
191                    .enable_all()
192                    .build()
193                    .map_err(|error| error.to_string())?
194                    .block_on(future)
195            })
196            .join()
197            .map_err(|_| "tool handler thread panicked".to_string())?
198        }
199    } else {
200        tokio::runtime::Builder::new_current_thread()
201            .enable_all()
202            .build()
203            .map_err(|error| error.to_string())?
204            .block_on(future)
205    }
206}