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};
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}