1use 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
35const PARAM_ATTR: &str = "param";
37
38#[proc_macro_attribute]
53pub fn tool(attr: TokenStream, item: TokenStream) -> TokenStream {
54 let func = parse_macro_input!(item as ItemFn);
55
56 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
68fn 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 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 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 let params = extract_params(&func.sig)?;
105 let field_names: Vec<Ident> = params.iter().map(|p| p.name.clone()).collect();
106
107 let output_type = match &func.sig.output {
109 syn::ReturnType::Default => quote! { () },
110 syn::ReturnType::Type(_, ty) => {
111 if let Some(inner) = extract_result_ok(ty) {
113 quote! { #inner }
114 } else {
115 quote! { #ty }
116 }
117 }
118 };
119
120 let input_fields = generate_input_fields(¶ms);
122
123 let input_field_attrs = generate_field_attrs(¶ms);
125
126 strip_param_attrs(&mut func);
129
130 let expanded = quote! {
132 #func
134
135 #[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 #[derive(serde::Deserialize, schemars::JsonSchema)]
153 pub struct #input_struct_name {
154 #(#input_field_attrs)*
155 #(#input_fields)*
156 }
157
158 #[::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 #[::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
201struct ParamInfo {
203 name: Ident,
204 ty: Type,
205 desc: Option<String>,
206}
207
208fn extract_params(sig: &Signature) -> Result<Vec<ParamInfo>> {
210 let mut params = Vec::new();
211
212 for arg in &sig.inputs {
213 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 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
238fn 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
255fn 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
275fn 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
289fn 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
310fn 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
323fn 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}