Skip to main content

lc_tools_derive/
lib.rs

1// lc-tools-derive/src/lib.rs
2//! Procedural macro for deriving BaseTool implementations from functions.
3//!
4//! # Example
5//!
6//! ```rust,ignore
7//! use lc_tools::{tool, BaseTool, Tool, ToolError};
8//!
9//! #[tool(description = "Useful for arithmetic calculations")]
10//! fn calculator(
11//!     #[param(desc = "The mathematical expression to evaluate")]
12//!     expression: String,
13//! ) -> Result<f64, ToolError> {
14//!     meval::eval_str(&expression)
15//!         .map_err(|e| ToolError::ExecutionFailed(e.to_string()))
16//! }
17//! ```
18//!
19//! This expands to:
20//! - `CalculatorTool` struct
21//! - `CalculatorInput` struct with `Deserialize` + `JsonSchema`
22//! - `impl BaseTool for CalculatorTool`
23//! - `impl Tool for CalculatorTool`
24//! - The original `calculator` function is preserved
25
26use proc_macro::TokenStream;
27use proc_macro2::TokenStream as TokenStream2;
28use quote::{format_ident, quote};
29use syn::{
30    parse_macro_input, Attribute, Expr, ExprLit, FnArg, Ident, ItemFn, Lit, Meta, MetaNameValue,
31    Pat, PatType, Result, Signature, Type,
32};
33
34/// Attribute for individual parameters.
35const PARAM_ATTR: &str = "param";
36
37/// The `#[tool]` procedural macro.
38///
39/// Transforms a function into a full Tool implementation.
40///
41/// # Attributes
42///
43/// - `#[tool(description = "...")]` — Required. The tool description shown to the LLM.
44/// - `#[param(desc = "...")]` — Optional per-parameter. Adds description to the JSON schema.
45///
46/// # Parameter Rules
47///
48/// - `String`, `i64`, `f64`, `bool` → required in schema
49/// - `Option<T>` → optional in schema
50/// - `Vec<T>` → array in schema
51#[proc_macro_attribute]
52pub fn tool(attr: TokenStream, item: TokenStream) -> TokenStream {
53    let func = parse_macro_input!(item as ItemFn);
54
55    // Parse the attribute as `description = "..."`
56    let description = match parse_tool_attr(attr) {
57        Ok(d) => d,
58        Err(err) => return err.to_compile_error().into(),
59    };
60
61    match tool_impl(description, func) {
62        Ok(tokens) => tokens.into(),
63        Err(err) => err.to_compile_error().into(),
64    }
65}
66
67/// Parse `#[tool(description = "...")]` attribute tokens.
68fn parse_tool_attr(attr: TokenStream) -> Result<String> {
69    if attr.is_empty() {
70        return Err(syn::Error::new(
71            proc_macro2::Span::call_site(),
72            "#[tool(description = \"...\")] is required",
73        ));
74    }
75
76    // Parse as `description = "..."`
77    let meta: Meta = syn::parse(attr)?;
78    if let Meta::NameValue(MetaNameValue { path, value, .. }) = &meta {
79        if path.is_ident("description") {
80            if let Expr::Lit(ExprLit {
81                lit: Lit::Str(lit), ..
82            }) = value
83            {
84                return Ok(lit.value());
85            }
86        }
87    }
88
89    Err(syn::Error::new(
90        proc_macro2::Span::call_site(),
91        "expected #[tool(description = \"...\")]",
92    ))
93}
94
95fn tool_impl(description: String, mut func: ItemFn) -> Result<TokenStream2> {
96    // 1. Extract all information from the function BEFORE mutating it
97    let func_name_str = func.sig.ident.to_string();
98    let tool_struct_name = format_ident!("{}Tool", to_pascal_case(&func_name_str));
99    let input_struct_name = format_ident!("{}Input", to_pascal_case(&func_name_str));
100    let func_name = func.sig.ident.clone();
101
102    // 2. Extract parameters from function signature
103    let params = extract_params(&func.sig)?;
104    let field_names: Vec<Ident> = params.iter().map(|p| p.name.clone()).collect();
105
106    // 3. Determine the output type from the function return type
107    let output_type = match &func.sig.output {
108        syn::ReturnType::Default => quote! { () },
109        syn::ReturnType::Type(_, ty) => {
110            // If it's Result<T, E>, extract T
111            if let Some(inner) = extract_result_ok(ty) {
112                quote! { #inner }
113            } else {
114                quote! { #ty }
115            }
116        }
117    };
118
119    // 4. Generate Input struct fields
120    let input_fields = generate_input_fields(&params);
121
122    // 5. Generate field-level schemars attributes for descriptions
123    let input_field_attrs = generate_field_attrs(&params);
124
125    // 6. Remove #[param] attributes from the original function so the compiler
126    //    doesn't complain about unknown attributes
127    strip_param_attrs(&mut func);
128
129    // 7. Generate the full expanded code
130    let expanded = quote! {
131        // Preserve the original function (with #[param] attrs stripped)
132        #func
133
134        /// Auto-generated Tool struct.
135        #[derive(Debug, Clone)]
136        pub struct #tool_struct_name;
137
138        impl ::std::default::Default for #tool_struct_name {
139            fn default() -> Self {
140                Self
141            }
142        }
143
144        impl #tool_struct_name {
145            pub fn new() -> Self {
146                Self
147            }
148        }
149
150        /// Auto-generated Input struct.
151        #[derive(serde::Deserialize, schemars::JsonSchema)]
152        pub struct #input_struct_name {
153            #(#input_field_attrs)*
154            #(#input_fields)*
155        }
156
157        // Implement Tool trait (type-safe version)
158        #[::async_trait::async_trait]
159        impl ::lc_core::tools::Tool for #tool_struct_name {
160            type Input = #input_struct_name;
161            type Output = #output_type;
162
163            async fn invoke(&self, input: Self::Input) -> ::std::result::Result<Self::Output, ::lc_core::tools::ToolError> {
164                let #input_struct_name { #(#field_names),* } = input;
165                #func_name(#(#field_names),*).map_err(|e| ::lc_core::tools::ToolError::ExecutionFailed(e.to_string()))
166            }
167        }
168
169        // Implement BaseTool trait (string version, for Agent)
170        #[::async_trait::async_trait]
171        impl ::lc_core::tools::BaseTool for #tool_struct_name {
172            fn name(&self) -> &str {
173                #func_name_str
174            }
175
176            fn description(&self) -> &str {
177                #description
178            }
179
180            async fn run(&self, input: ::std::string::String) -> ::std::result::Result<::std::string::String, ::lc_core::tools::ToolError> {
181                let parsed: #input_struct_name = ::serde_json::from_str(&input)
182                    .map_err(|e| ::lc_core::tools::ToolError::InvalidInput(format!("JSON parse error: {}", e)))?;
183                let #input_struct_name { #(#field_names),* } = parsed;
184                let result = #func_name(#(#field_names),*)
185                    .map_err(|e| ::lc_core::tools::ToolError::ExecutionFailed(e.to_string()))?;
186                Ok(::serde_json::to_string(&result)
187                    .unwrap_or_else(|_| format!("{:?}", result)))
188            }
189
190            fn args_schema(&self) -> ::std::option::Option<::serde_json::Value> {
191                use ::schemars::schema_for;
192                ::serde_json::to_value(schema_for!(#input_struct_name)).ok()
193            }
194        }
195    };
196
197    Ok(expanded)
198}
199
200/// Parameter info extracted from function signature.
201struct ParamInfo {
202    name: Ident,
203    ty: Type,
204    desc: Option<String>,
205    #[allow(dead_code)]
206    is_option: bool,
207}
208
209/// Extract parameter info from the function signature.
210fn extract_params(sig: &Signature) -> Result<Vec<ParamInfo>> {
211    let mut params = Vec::new();
212
213    for arg in &sig.inputs {
214        // Skip self parameter
215        if let FnArg::Receiver(_) = arg {
216            continue;
217        }
218
219        if let FnArg::Typed(PatType { pat, ty, attrs, .. }) = arg {
220            let name = match pat.as_ref() {
221                Pat::Ident(ident) => ident.ident.clone(),
222                _ => continue,
223            };
224
225            let is_option = is_option_type(ty);
226
227            // Extract #[param(desc = "...")] attribute
228            let desc = extract_param_desc(attrs);
229
230            params.push(ParamInfo {
231                name,
232                ty: (*(*ty)).clone(),
233                desc,
234                is_option,
235            });
236        }
237    }
238
239    Ok(params)
240}
241
242/// Check if a type is `Option<T>`.
243fn is_option_type(ty: &Type) -> bool {
244    if let Type::Path(type_path) = ty {
245        if type_path.path.segments.len() == 1 {
246            return type_path.path.segments[0].ident == "Option";
247        }
248    }
249    false
250}
251
252/// Extract `T` from `Result<T, E>`. Returns None if not a Result type.
253fn extract_result_ok(ty: &Type) -> Option<Type> {
254    if let Type::Path(type_path) = ty {
255        if type_path.path.segments.len() == 1 {
256            let segment = &type_path.path.segments[0];
257            if segment.ident == "Result" {
258                if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
259                    if let Some(syn::GenericArgument::Type(inner)) = args.args.first() {
260                        return Some(inner.clone());
261                    }
262                }
263            }
264        }
265    }
266    None
267}
268
269/// Extract `desc` from `#[param(desc = "...")]`.
270fn extract_param_desc(attrs: &[Attribute]) -> Option<String> {
271    for attr in attrs {
272        if attr.path().is_ident(PARAM_ATTR) {
273            let meta: Meta = attr.parse_args().ok()?;
274            if let Meta::NameValue(MetaNameValue { path, value, .. }) = &meta {
275                if path.is_ident("desc") {
276                    if let Expr::Lit(ExprLit {
277                        lit: Lit::Str(lit), ..
278                    }) = value
279                    {
280                        return Some(lit.value());
281                    }
282                }
283            }
284        }
285    }
286    None
287}
288
289/// Generate Input struct fields.
290fn generate_input_fields(params: &[ParamInfo]) -> Vec<TokenStream2> {
291    params
292        .iter()
293        .map(|p| {
294            let name = &p.name;
295            let ty = &p.ty;
296            quote! {
297                pub #name: #ty,
298            }
299        })
300        .collect()
301}
302
303/// Generate schemars field attributes for descriptions.
304///
305/// Uses `#[doc = "..."]` instead of `#[schemars(description = "...")]` because
306/// schemars 0.8 automatically extracts descriptions from doc comments, and using
307/// `#[schemars(description)]` alongside the derive causes "duplicate attribute" errors
308/// when there are multiple fields.
309fn generate_field_attrs(params: &[ParamInfo]) -> Vec<TokenStream2> {
310    params
311        .iter()
312        .map(|p| {
313            if let Some(desc) = &p.desc {
314                quote! {
315                    #[doc = #desc]
316                }
317            } else {
318                quote! {}
319            }
320        })
321        .collect()
322}
323
324/// Convert snake_case to PascalCase.
325fn to_pascal_case(s: &str) -> String {
326    s.split('_')
327        .map(|word| {
328            let mut chars = word.chars();
329            match chars.next() {
330                None => String::new(),
331                Some(c) => c.to_uppercase().collect::<String>() + chars.as_str(),
332            }
333        })
334        .collect()
335}
336
337/// Remove all `#[param(...)]` attributes from function parameters.
338/// This prevents the compiler from complaining about unknown attributes
339/// when the original function is emitted.
340fn strip_param_attrs(func: &mut ItemFn) {
341    for arg in &mut func.sig.inputs {
342        if let FnArg::Typed(pat_type) = arg {
343            pat_type.attrs.retain(|attr| !attr.path().is_ident(PARAM_ATTR));
344        }
345    }
346}