1#![warn(missing_docs)]
2use proc_macro::TokenStream;
29use proc_macro2::TokenStream as TokenStream2;
30use quote::{format_ident, quote};
31use syn::{
32 parse_macro_input, Attribute, Expr, ExprLit, FnArg, Ident, ItemFn, Lit, Meta, MetaNameValue,
33 Pat, PatType, Result, Signature, Type,
34};
35
36const PARAM_ATTR: &str = "param";
38
39#[proc_macro_attribute]
54pub fn tool(attr: TokenStream, item: TokenStream) -> TokenStream {
55 let func = parse_macro_input!(item as ItemFn);
56
57 let description = match parse_tool_attr(attr) {
59 Ok(d) => d,
60 Err(err) => return err.to_compile_error().into(),
61 };
62
63 match tool_impl(description, func) {
64 Ok(tokens) => tokens.into(),
65 Err(err) => err.to_compile_error().into(),
66 }
67}
68
69fn parse_tool_attr(attr: TokenStream) -> Result<String> {
71 if attr.is_empty() {
72 return Err(syn::Error::new(
73 proc_macro2::Span::call_site(),
74 "#[tool(description = \"...\")] is required",
75 ));
76 }
77
78 let meta: Meta = syn::parse(attr)?;
80 if let Meta::NameValue(MetaNameValue { path, value, .. }) = &meta {
81 if path.is_ident("description") {
82 if let Expr::Lit(ExprLit {
83 lit: Lit::Str(lit), ..
84 }) = value
85 {
86 return Ok(lit.value());
87 }
88 }
89 }
90
91 Err(syn::Error::new(
92 proc_macro2::Span::call_site(),
93 "expected #[tool(description = \"...\")]",
94 ))
95}
96
97fn tool_impl(description: String, mut func: ItemFn) -> Result<TokenStream2> {
98 let func_name_str = func.sig.ident.to_string();
100 let tool_struct_name = format_ident!("{}Tool", to_pascal_case(&func_name_str));
101 let input_struct_name = format_ident!("{}Input", to_pascal_case(&func_name_str));
102 let func_name = func.sig.ident.clone();
103
104 let params = extract_params(&func.sig)?;
106 let field_names: Vec<Ident> = params.iter().map(|p| p.name.clone()).collect();
107
108 let output_type = match &func.sig.output {
110 syn::ReturnType::Default => quote! { () },
111 syn::ReturnType::Type(_, ty) => {
112 if let Some(inner) = extract_result_ok(ty) {
114 quote! { #inner }
115 } else {
116 quote! { #ty }
117 }
118 }
119 };
120
121 let invoke_body = if return_type_is_tool_error(&func.sig.output) {
125 quote! { #func_name(#(#field_names),*) }
126 } else {
127 quote! { #func_name(#(#field_names),*).map_err(|e| ::lc_core::tools::ToolError::ExecutionFailed(e.to_string())) }
128 };
129
130 let input_fields = generate_input_fields(¶ms);
132
133 let input_field_attrs = generate_field_attrs(¶ms);
135
136 strip_param_attrs(&mut func);
139
140 let expanded = quote! {
142 #func
144
145 #[derive(Debug, Clone)]
147 pub struct #tool_struct_name;
148
149 impl ::std::default::Default for #tool_struct_name {
150 fn default() -> Self {
151 Self
152 }
153 }
154
155 impl #tool_struct_name {
156 pub fn new() -> Self {
157 Self
158 }
159 }
160
161 #[derive(serde::Deserialize, schemars::JsonSchema)]
163 pub struct #input_struct_name {
164 #(#input_field_attrs)*
165 #(#input_fields)*
166 }
167
168 #[::async_trait::async_trait]
170 impl ::lc_core::tools::Tool for #tool_struct_name {
171 type Input = #input_struct_name;
172 type Output = #output_type;
173
174 async fn invoke(&self, input: Self::Input) -> ::std::result::Result<Self::Output, ::lc_core::tools::ToolError> {
175 let #input_struct_name { #(#field_names),* } = input;
176 #invoke_body
177 }
178 }
179
180 #[::async_trait::async_trait]
182 impl ::lc_core::tools::BaseTool for #tool_struct_name {
183 fn name(&self) -> &str {
184 #func_name_str
185 }
186
187 fn description(&self) -> &str {
188 #description
189 }
190
191 async fn run(&self, input: ::std::string::String) -> ::std::result::Result<::std::string::String, ::lc_core::tools::ToolError> {
192 let parsed: #input_struct_name = ::serde_json::from_str(&input)
193 .map_err(|e| ::lc_core::tools::ToolError::InvalidInput(format!("JSON parse error: {}", e)))?;
194 let #input_struct_name { #(#field_names),* } = parsed;
195 let result = #func_name(#(#field_names),*)
196 .map_err(|e| ::lc_core::tools::ToolError::ExecutionFailed(e.to_string()))?;
197 let serialized = ::serde_json::to_string(&result)
200 .map_err(|e| ::lc_core::tools::ToolError::ExecutionFailed(format!("Failed to serialize tool output: {}", e)))?;
201 Ok(serialized)
202 }
203
204 fn args_schema(&self) -> ::std::option::Option<::serde_json::Value> {
205 use ::schemars::schema_for;
206 ::serde_json::to_value(schema_for!(#input_struct_name)).ok()
207 }
208 }
209 };
210
211 Ok(expanded)
212}
213
214struct ParamInfo {
216 name: Ident,
217 ty: Type,
218 desc: Option<String>,
219}
220
221fn extract_params(sig: &Signature) -> Result<Vec<ParamInfo>> {
223 let mut params = Vec::new();
224
225 for arg in &sig.inputs {
226 if let FnArg::Receiver(_) = arg {
228 continue;
229 }
230
231 if let FnArg::Typed(PatType { pat, ty, attrs, .. }) = arg {
232 let name = match pat.as_ref() {
233 Pat::Ident(ident) => ident.ident.clone(),
234 _ => continue,
235 };
236
237 let desc = extract_param_desc(attrs);
239
240 params.push(ParamInfo {
241 name,
242 ty: (*(*ty)).clone(),
243 desc,
244 });
245 }
246 }
247
248 Ok(params)
249}
250
251fn extract_result_ok(ty: &Type) -> Option<Type> {
253 if let Type::Path(type_path) = ty {
254 if type_path.path.segments.len() == 1 {
255 let segment = &type_path.path.segments[0];
256 if segment.ident == "Result" {
257 if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
258 if let Some(syn::GenericArgument::Type(inner)) = args.args.first() {
259 return Some(inner.clone());
260 }
261 }
262 }
263 }
264 }
265 None
266}
267
268fn return_type_is_tool_error(ret: &syn::ReturnType) -> bool {
275 let syn::ReturnType::Type(_, ty) = ret else {
276 return false;
277 };
278 let Type::Path(type_path) = &**ty else {
279 return false;
280 };
281 let Some(seg) = type_path.path.segments.last() else {
282 return false;
283 };
284 if seg.ident != "Result" {
285 return false;
286 }
287 let syn::PathArguments::AngleBracketed(args) = &seg.arguments else {
289 return false;
290 };
291 let mut generic = args.args.iter().filter_map(|a| match a {
292 syn::GenericArgument::Type(t) => Some(t),
293 _ => None,
294 });
295 let _ok = generic.next();
296 let Some(err) = generic.next() else {
297 return false;
298 };
299 let Type::Path(err_path) = err else {
300 return false;
301 };
302 err_path
303 .path
304 .segments
305 .last()
306 .is_some_and(|s| s.ident == "ToolError")
307}
308
309fn extract_param_desc(attrs: &[Attribute]) -> Option<String> {
311 for attr in attrs {
312 if attr.path().is_ident(PARAM_ATTR) {
313 let meta: Meta = attr.parse_args().ok()?;
314 if let Meta::NameValue(MetaNameValue { path, value, .. }) = &meta {
315 if path.is_ident("desc") {
316 if let Expr::Lit(ExprLit {
317 lit: Lit::Str(lit), ..
318 }) = value
319 {
320 return Some(lit.value());
321 }
322 }
323 }
324 }
325 }
326 None
327}
328
329fn generate_input_fields(params: &[ParamInfo]) -> Vec<TokenStream2> {
331 params
332 .iter()
333 .map(|p| {
334 let name = &p.name;
335 let ty = &p.ty;
336 quote! {
337 pub #name: #ty,
338 }
339 })
340 .collect()
341}
342
343fn generate_field_attrs(params: &[ParamInfo]) -> Vec<TokenStream2> {
350 params
351 .iter()
352 .map(|p| {
353 if let Some(desc) = &p.desc {
354 quote! {
355 #[doc = #desc]
356 }
357 } else {
358 quote! {}
359 }
360 })
361 .collect()
362}
363
364fn to_pascal_case(s: &str) -> String {
366 s.split('_')
367 .map(|word| {
368 let mut chars = word.chars();
369 match chars.next() {
370 None => String::new(),
371 Some(c) => c.to_uppercase().collect::<String>() + chars.as_str(),
372 }
373 })
374 .collect()
375}
376
377fn strip_param_attrs(func: &mut ItemFn) {
381 for arg in &mut func.sig.inputs {
382 if let FnArg::Typed(pat_type) = arg {
383 pat_type
384 .attrs
385 .retain(|attr| !attr.path().is_ident(PARAM_ATTR));
386 }
387 }
388}