1use 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
34const PARAM_ATTR: &str = "param";
36
37#[proc_macro_attribute]
52pub fn tool(attr: TokenStream, item: TokenStream) -> TokenStream {
53 let func = parse_macro_input!(item as ItemFn);
54
55 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
67fn 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 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 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 let params = extract_params(&func.sig)?;
104 let field_names: Vec<Ident> = params.iter().map(|p| p.name.clone()).collect();
105
106 let output_type = match &func.sig.output {
108 syn::ReturnType::Default => quote! { () },
109 syn::ReturnType::Type(_, ty) => {
110 if let Some(inner) = extract_result_ok(ty) {
112 quote! { #inner }
113 } else {
114 quote! { #ty }
115 }
116 }
117 };
118
119 let input_fields = generate_input_fields(¶ms);
121
122 let input_field_attrs = generate_field_attrs(¶ms);
124
125 strip_param_attrs(&mut func);
128
129 let expanded = quote! {
131 #func
133
134 #[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 #[derive(serde::Deserialize, schemars::JsonSchema)]
152 pub struct #input_struct_name {
153 #(#input_field_attrs)*
154 #(#input_fields)*
155 }
156
157 #[::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 #[::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
200struct ParamInfo {
202 name: Ident,
203 ty: Type,
204 desc: Option<String>,
205 #[allow(dead_code)]
206 is_option: bool,
207}
208
209fn extract_params(sig: &Signature) -> Result<Vec<ParamInfo>> {
211 let mut params = Vec::new();
212
213 for arg in &sig.inputs {
214 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 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
242fn 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
252fn 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
269fn 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
289fn 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
303fn 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
324fn 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
337fn 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
344 .attrs
345 .retain(|attr| !attr.path().is_ident(PARAM_ATTR));
346 }
347 }
348}