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