Skip to main content

lc_tools_derive/
lib.rs

1#![warn(missing_docs)]
2// lc-tools-derive/src/lib.rs
3//! Procedural macro for deriving BaseTool implementations from functions.
4//!
5//! # Example
6//!
7//! ```rust,ignore
8//! use lc_tools::{tool, BaseTool, Tool, ToolError};
9//!
10//! #[tool(description = "Useful for arithmetic calculations")]
11//! fn calculator(
12//!     #[param(desc = "The mathematical expression to evaluate")]
13//!     expression: String,
14//! ) -> Result<f64, ToolError> {
15//!     expression
16//!         .parse::<f64>()
17//!         .map_err(|e| ToolError::ExecutionFailed(e.to_string()))
18//! }
19//! ```
20//!
21//! This expands to:
22//! - `CalculatorTool` struct
23//! - `CalculatorInput` struct with `Deserialize` + `JsonSchema`
24//! - `impl BaseTool for CalculatorTool`
25//! - `impl Tool for CalculatorTool`
26//! - The original `calculator` function is preserved
27
28use 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
36/// Attribute for individual parameters.
37const PARAM_ATTR: &str = "param";
38
39/// The `#[tool]` procedural macro.
40///
41/// Transforms a function into a full Tool implementation.
42///
43/// # Attributes
44///
45/// - `#[tool(description = "...")]` — Required. The tool description shown to the LLM.
46/// - `#[param(desc = "...")]` — Optional per-parameter. Adds description to the JSON schema.
47///
48/// # Parameter Rules
49///
50/// - `String`, `i64`, `f64`, `bool` → required in schema
51/// - `Option<T>` → optional in schema
52/// - `Vec<T>` → array in schema
53#[proc_macro_attribute]
54pub fn tool(attr: TokenStream, item: TokenStream) -> TokenStream {
55    let func = parse_macro_input!(item as ItemFn);
56
57    // Parse the attribute as `description = "..."`
58    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
69/// Parse `#[tool(description = "...")]` attribute tokens.
70fn 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    // Parse as `description = "..."`
79    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    // 1. Extract all information from the function BEFORE mutating it
99    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    // 2. Extract parameters from function signature
105    let params = extract_params(&func.sig)?;
106    let field_names: Vec<Ident> = params.iter().map(|p| p.name.clone()).collect();
107
108    // 3. Determine the output type from the function return type
109    let output_type = match &func.sig.output {
110        syn::ReturnType::Default => quote! { () },
111        syn::ReturnType::Type(_, ty) => {
112            // If it's Result<T, E>, extract T
113            if let Some(inner) = extract_result_ok(ty) {
114                quote! { #inner }
115            } else {
116                quote! { #ty }
117            }
118        }
119    };
120
121    // 3b. F5:函数返回 `Result<_, ToolError>` 时,`invoke` 直接透传原错误
122    // (参数错 / 业务错原样保留,不再统一压平成 `ExecutionFailed`);返回
123    // 其他错误类型时才包 `ExecutionFailed`。此为 breaking:错误语义变化。
124    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    // 4. Generate Input struct fields
131    let input_fields = generate_input_fields(&params);
132
133    // 5. Generate field-level schemars attributes for descriptions
134    let input_field_attrs = generate_field_attrs(&params);
135
136    // 6. Remove #[param] attributes from the original function so the compiler
137    //    doesn't complain about unknown attributes
138    strip_param_attrs(&mut func);
139
140    // 7. Generate the full expanded code
141    let expanded = quote! {
142        // Preserve the original function (with #[param] attrs stripped)
143        #func
144
145        /// Auto-generated Tool struct.
146        #[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        /// Auto-generated Input struct.
162        #[derive(serde::Deserialize, schemars::JsonSchema)]
163        pub struct #input_struct_name {
164            #(#input_field_attrs)*
165            #(#input_fields)*
166        }
167
168        // Implement Tool trait (type-safe version)
169        #[::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        // Implement BaseTool trait (string version, for Agent)
181        #[::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                // F5:序列化失败不再静默回退 Debug 文本(会把 Rust 内部结构喂给 LLM),
198                // 而是返回 ExecutionFailed 错误,让上层明确感知输出无法序列化。
199                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
214/// Parameter info extracted from function signature.
215struct ParamInfo {
216    name: Ident,
217    ty: Type,
218    desc: Option<String>,
219}
220
221/// Extract parameter info from the function signature.
222fn extract_params(sig: &Signature) -> Result<Vec<ParamInfo>> {
223    let mut params = Vec::new();
224
225    for arg in &sig.inputs {
226        // Skip self parameter
227        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            // Extract #[param(desc = "...")] attribute
238            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
251/// Extract `T` from `Result<T, E>`. Returns None if not a Result type.
252fn 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
268/// 判断函数返回类型是否为 `Result<_, ToolError>`(错误类型最后一个路径段为
269/// `ToolError`)。F5:宏据此决定 `invoke` 是否直接透传原错误。
270///
271/// 只能是语法级判断:任何以 `ToolError` 结尾的错误类型都被视为库的 `ToolError`。
272/// 常见写法均命中——裸 `ToolError`、`lc_core::tools::ToolError`、`tools::ToolError`
273/// 或用户 `use` 进来的别名;返回 `anyhow::Error` / `String` 等其他错误类型时不命中。
274fn 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    // 取出 `Result<_, E>` 的第二个泛型参数作为错误类型
288    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
309/// Extract `desc` from `#[param(desc = "...")]`.
310fn 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
329/// Generate Input struct fields.
330fn 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
343/// Generate schemars field attributes for descriptions.
344///
345/// Uses `#[doc = "..."]` instead of `#[schemars(description = "...")]` because
346/// schemars 0.8 automatically extracts descriptions from doc comments, and using
347/// `#[schemars(description)]` alongside the derive causes "duplicate attribute" errors
348/// when there are multiple fields.
349fn 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
364/// Convert snake_case to PascalCase.
365fn 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
377/// Remove all `#[param(...)]` attributes from function parameters.
378/// This prevents the compiler from complaining about unknown attributes
379/// when the original function is emitted.
380fn 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}