vv_agent/tools/
function.rs1use 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}