rig_derive/
lib.rs

1extern crate proc_macro;
2
3use convert_case::{Case, Casing};
4use proc_macro::TokenStream;
5use quote::{format_ident, quote};
6use std::{collections::HashMap, ops::Deref};
7use syn::{
8    Attribute, DeriveInput, Expr, ExprLit, Ident, Lit, Meta, PathArguments, ReturnType, Token,
9    Type,
10    parse::{Parse, ParseStream},
11    parse_macro_input,
12    punctuated::Punctuated,
13};
14
15mod basic;
16mod custom;
17mod embed;
18
19pub(crate) const EMBED: &str = "embed";
20
21pub(crate) fn rig_core_path() -> proc_macro2::TokenStream {
22    match proc_macro_crate::crate_name("rig-core") {
23        Ok(proc_macro_crate::FoundCrate::Itself) => quote!(crate),
24        Ok(proc_macro_crate::FoundCrate::Name(name)) => {
25            let ident = format_ident!("{name}");
26            quote!(::#ident)
27        }
28        Err(_) => match proc_macro_crate::crate_name("rig") {
29            Ok(proc_macro_crate::FoundCrate::Itself) => quote!(crate),
30            Ok(proc_macro_crate::FoundCrate::Name(name)) => {
31                let ident = format_ident!("{name}");
32                quote!(::#ident)
33            }
34            Err(_) => quote!(::rig_core),
35        },
36    }
37}
38
39//References:
40//<https://doc.rust-lang.org/book/ch19-06-macros.html#how-to-write-a-custom-derive-macro>
41//<https://doc.rust-lang.org/reference/procedural-macros.html>
42/// A macro that allows you to implement the `rig::embedding::Embed` trait by deriving it.
43/// Usage can be found below:
44///
45/// ```text
46/// use rig::Embed;
47/// use rig_derive::Embed;
48///
49/// #[derive(Embed)]
50/// struct Foo {
51///     id: String,
52///     #[embed] // this helper shows which field to embed
53///     description: String
54///}
55/// ```
56#[proc_macro_derive(Embed, attributes(embed))]
57pub fn derive_embedding_trait(item: TokenStream) -> TokenStream {
58    let mut input = parse_macro_input!(item as DeriveInput);
59
60    embed::expand_derive_embedding(&mut input)
61        .unwrap_or_else(syn::Error::into_compile_error)
62        .into()
63}
64
65struct MacroArgs {
66    name: Option<String>,
67    description: Option<String>,
68    param_descriptions: HashMap<String, String>,
69    required: Option<Vec<String>>,
70}
71
72fn parse_string_literal(expr: &Expr, field_name: &str) -> syn::Result<String> {
73    match expr {
74        Expr::Lit(ExprLit {
75            lit: Lit::Str(lit_str),
76            ..
77        }) => Ok(lit_str.value()),
78        _ => Err(syn::Error::new_spanned(
79            expr,
80            format!("`{field_name}` must be a string literal"),
81        )),
82    }
83}
84
85fn validate_explicit_tool_name(name: &str, expr: &Expr) -> syn::Result<()> {
86    if name.is_empty() || name.len() > 64 {
87        return Err(syn::Error::new_spanned(
88            expr,
89            "`name` must be between 1 and 64 characters long",
90        ));
91    }
92
93    let mut chars = name.chars();
94    let Some(first_char) = chars.next() else {
95        return Err(syn::Error::new_spanned(
96            expr,
97            "`name` must be between 1 and 64 characters long",
98        ));
99    };
100
101    if !first_char.is_ascii_alphabetic() && first_char != '_' {
102        return Err(syn::Error::new_spanned(
103            expr,
104            "`name` must start with an ASCII letter or underscore",
105        ));
106    }
107
108    if chars.any(|ch| !ch.is_ascii_alphanumeric() && ch != '_' && ch != '-') {
109        return Err(syn::Error::new_spanned(
110            expr,
111            "`name` may only contain ASCII letters, digits, underscores, or hyphens",
112        ));
113    }
114
115    Ok(())
116}
117
118impl Parse for MacroArgs {
119    fn parse(input: ParseStream) -> syn::Result<Self> {
120        let mut name = None;
121        let mut description = None;
122        let mut param_descriptions = HashMap::new();
123        let mut required = None;
124
125        // If the input is empty, return default values
126        if input.is_empty() {
127            return Ok(MacroArgs {
128                name,
129                description,
130                param_descriptions,
131                required,
132            });
133        }
134
135        let meta_list: Punctuated<Meta, Token![,]> = Punctuated::parse_terminated(input)?;
136
137        for meta in meta_list {
138            match meta {
139                Meta::NameValue(nv) => {
140                    let ident = nv.path.get_ident().ok_or_else(|| {
141                        syn::Error::new_spanned(
142                            &nv.path,
143                            "unsupported top-level #[rig_tool] argument",
144                        )
145                    })?;
146
147                    match ident.to_string().as_str() {
148                        "name" => {
149                            let parsed_name = parse_string_literal(&nv.value, "name")?;
150                            validate_explicit_tool_name(&parsed_name, &nv.value)?;
151                            name = Some(parsed_name);
152                        }
153                        "description" => {
154                            description = Some(parse_string_literal(&nv.value, "description")?);
155                        }
156                        _ => {
157                            return Err(syn::Error::new_spanned(
158                                &nv.path,
159                                format!("unsupported top-level #[rig_tool] argument `{}`", ident),
160                            ));
161                        }
162                    }
163                }
164                Meta::List(list) => {
165                    let ident = list.path.get_ident().ok_or_else(|| {
166                        syn::Error::new_spanned(
167                            &list.path,
168                            "unsupported top-level #[rig_tool] argument",
169                        )
170                    })?;
171
172                    match ident.to_string().as_str() {
173                        "params" => {
174                            let nested: Punctuated<Meta, Token![,]> =
175                                list.parse_args_with(Punctuated::parse_terminated)?;
176
177                            for meta in nested {
178                                if let Meta::NameValue(nv) = meta
179                                    && let Expr::Lit(ExprLit {
180                                        lit: Lit::Str(lit_str),
181                                        ..
182                                    }) = nv.value
183                                {
184                                    let Some(param_ident) = nv.path.get_ident() else {
185                                        return Err(syn::Error::new_spanned(
186                                            &nv.path,
187                                            "parameter descriptions must use identifier keys",
188                                        ));
189                                    };
190                                    let param_name = param_ident.to_string();
191                                    param_descriptions.insert(param_name, lit_str.value());
192                                }
193                            }
194                        }
195                        "required" => {
196                            let required_variables: Punctuated<Ident, Token![,]> =
197                                list.parse_args_with(Punctuated::parse_terminated)?;
198
199                            required = Some(
200                                required_variables
201                                    .into_iter()
202                                    .map(|x| x.to_string())
203                                    .collect(),
204                            );
205                        }
206                        _ => {
207                            return Err(syn::Error::new_spanned(
208                                &list.path,
209                                format!("unsupported top-level #[rig_tool] argument `{}`", ident),
210                            ));
211                        }
212                    }
213                }
214                Meta::Path(path) => {
215                    let message = if let Some(ident) = path.get_ident() {
216                        format!("unsupported top-level #[rig_tool] argument `{ident}`")
217                    } else {
218                        "unsupported top-level #[rig_tool] argument".to_string()
219                    };
220
221                    return Err(syn::Error::new_spanned(path, message));
222                }
223            }
224        }
225
226        Ok(MacroArgs {
227            name,
228            description,
229            param_descriptions,
230            required,
231        })
232    }
233}
234
235/// Extract doc comment text from `#[doc = "..."]` attributes.
236fn extract_doc_comment(attrs: &[Attribute]) -> Option<String> {
237    let lines: Vec<String> = attrs
238        .iter()
239        .filter_map(|attr| {
240            if !attr.path().is_ident("doc") {
241                return None;
242            }
243            if let Meta::NameValue(nv) = &attr.meta
244                && let Expr::Lit(ExprLit {
245                    lit: Lit::Str(s), ..
246                }) = &nv.value
247            {
248                return Some(s.value());
249            }
250            None
251        })
252        .collect();
253
254    if lines.is_empty() {
255        return None;
256    }
257
258    Some(
259        lines
260            .iter()
261            .map(|l| l.strip_prefix(' ').unwrap_or(l))
262            .collect::<Vec<_>>()
263            .join("\n")
264            .trim()
265            .to_string(),
266    )
267}
268
269/// Check if a type is `Option<T>`.
270fn is_option_type(ty: &Type) -> bool {
271    if let Type::Path(type_path) = ty
272        && let Some(segment) = type_path.path.segments.last()
273    {
274        return segment.ident == "Option";
275    }
276    false
277}
278
279fn result_type_tokens(
280    return_type: &ReturnType,
281) -> syn::Result<(proc_macro2::TokenStream, proc_macro2::TokenStream)> {
282    let ReturnType::Type(_, ty) = return_type else {
283        return Err(syn::Error::new_spanned(
284            return_type,
285            "function must have a return type of Result<T, E>",
286        ));
287    };
288
289    let Type::Path(type_path) = ty.deref() else {
290        return Err(syn::Error::new_spanned(
291            ty,
292            "return type must be Result<T, E>",
293        ));
294    };
295
296    let Some(last_segment) = type_path.path.segments.last() else {
297        return Err(syn::Error::new_spanned(
298            &type_path.path,
299            "return type must be Result<T, E>",
300        ));
301    };
302
303    if last_segment.ident != "Result" {
304        return Err(syn::Error::new_spanned(
305            &last_segment.ident,
306            "return type must be Result<T, E>",
307        ));
308    }
309
310    let PathArguments::AngleBracketed(args) = &last_segment.arguments else {
311        return Err(syn::Error::new_spanned(
312            &last_segment.arguments,
313            "expected angle-bracketed type parameters for Result<T, E>",
314        ));
315    };
316
317    let mut generic_args = args.args.iter();
318    let Some(output) = generic_args.next() else {
319        return Err(syn::Error::new_spanned(
320            &args.args,
321            "expected Result<T, E> with exactly two type parameters",
322        ));
323    };
324    let Some(error) = generic_args.next() else {
325        return Err(syn::Error::new_spanned(
326            &args.args,
327            "expected Result<T, E> with exactly two type parameters",
328        ));
329    };
330
331    if generic_args.next().is_some() {
332        return Err(syn::Error::new_spanned(
333            &args.args,
334            "expected Result<T, E> with exactly two type parameters",
335        ));
336    }
337
338    Ok((quote!(#output), quote!(#error)))
339}
340
341/// A procedural macro that transforms a function into a `rig::tool::Tool` that can be used with a `rig::agent::Agent`.
342///
343/// # Examples
344///
345/// Basic usage:
346/// ```text
347/// use rig_derive::rig_tool;
348///
349/// #[rig_tool]
350/// fn add(a: i32, b: i32) -> Result<i32, rig::tool::ToolError> {
351///     Ok(a + b)
352/// }
353/// ```
354///
355/// With description:
356/// ```text
357/// use rig_derive::rig_tool;
358///
359/// #[rig_tool(description = "Perform basic arithmetic operations")]
360/// fn calculator(x: i32, y: i32, operation: String) -> Result<i32, rig::tool::ToolError> {
361///     match operation.as_str() {
362///         "add" => Ok(x + y),
363///         "subtract" => Ok(x - y),
364///         "multiply" => Ok(x * y),
365///         "divide" => Ok(x / y),
366///         _ => Err(rig::tool::ToolError::ToolCallError("Unknown operation".into())),
367///     }
368/// }
369/// ```
370///
371/// With a custom tool name:
372/// ```text
373/// use rig_derive::rig_tool;
374///
375/// // Explicit names must be string literals that start with an ASCII letter
376/// // or `_`, may contain ASCII letters, digits, `_`, or `-`, and be at most
377/// // 64 characters long.
378/// #[rig_tool(name = "search-docs", description = "Search the documentation")]
379/// fn search_docs_impl(query: String) -> Result<String, rig::tool::ToolError> {
380///     Ok(format!("Searching docs for {query}"))
381/// }
382/// ```
383///
384/// With parameter descriptions:
385/// ```text
386/// use rig_derive::rig_tool;
387///
388/// #[rig_tool(
389///     description = "A tool that performs string operations",
390///     params(
391///         text = "The input text to process",
392///         operation = "The operation to perform (uppercase, lowercase, reverse)"
393///     )
394/// )]
395/// fn string_processor(text: String, operation: String) -> Result<String, rig::tool::ToolError> {
396///     match operation.as_str() {
397///         "uppercase" => Ok(text.to_uppercase()),
398///         "lowercase" => Ok(text.to_lowercase()),
399///         "reverse" => Ok(text.chars().rev().collect()),
400///         _ => Err(rig::tool::ToolError::ToolCallError("Unknown operation".into())),
401///     }
402/// }
403/// ```
404#[proc_macro_attribute]
405pub fn rig_tool(args: TokenStream, input: TokenStream) -> TokenStream {
406    let args = parse_macro_input!(args as MacroArgs);
407    let input_fn = parse_macro_input!(input as syn::ItemFn);
408
409    // Extract function details
410    let fn_name = &input_fn.sig.ident;
411    let fn_name_str = fn_name.to_string();
412    let tool_name = args.name.clone().unwrap_or_else(|| fn_name_str.clone());
413    let vis = &input_fn.vis;
414    let is_async = input_fn.sig.asyncness.is_some();
415
416    // Build a cleaned copy of the function with doc attrs stripped from parameters,
417    // since `#[doc]` on function parameters is not allowed by the compiler.
418    let cleaned_fn = {
419        let mut f = input_fn.clone();
420        for arg in f.sig.inputs.iter_mut() {
421            if let syn::FnArg::Typed(pat_type) = arg {
422                pat_type.attrs.retain(|a| !a.path().is_ident("doc"));
423            }
424        }
425        f
426    };
427
428    // Extract return type and get Output and Error types from Result<T, E>
429    let return_type = &input_fn.sig.output;
430    let (output_type, error_type) = match result_type_tokens(return_type) {
431        Ok(types) => types,
432        Err(error) => return error.into_compile_error().into(),
433    };
434
435    // Generate PascalCase struct name from the function name
436    let struct_name = format_ident!("{}", { fn_name_str.to_case(Case::Pascal) });
437
438    // Tool description: explicit attribute > doc comment > default
439    let fn_doc = extract_doc_comment(&input_fn.attrs);
440    let tool_description = match args.description {
441        Some(desc) => quote! { #desc.to_string() },
442        None => match fn_doc {
443            Some(doc) => quote! { #doc.to_string() },
444            None => quote! { format!("Function to {}", Self::NAME) },
445        },
446    };
447
448    // Extract parameter names, doc comments, and build struct field tokens
449    let mut param_names = Vec::new();
450    let mut field_tokens = Vec::new();
451
452    for arg in input_fn.sig.inputs.iter() {
453        if let syn::FnArg::Typed(pat_type) = arg
454            && let syn::Pat::Ident(param_ident) = &*pat_type.pat
455        {
456            let param_name = &param_ident.ident;
457            let param_name_str = param_name.to_string();
458            let ty = &pat_type.ty;
459
460            // Determine the description for this field:
461            // explicit params() > parameter doc comment > default
462            let field_doc_attr =
463                if let Some(explicit) = args.param_descriptions.get(&param_name_str) {
464                    // Explicit override via params() — use #[schemars(description = "...")]
465                    quote! { #[schemars(description = #explicit)] }
466                } else if let Some(doc) = extract_doc_comment(&pat_type.attrs) {
467                    // Doc comment on the parameter — propagate as #[doc = "..."]
468                    quote! { #[doc = #doc] }
469                } else {
470                    // Default fallback
471                    let default_desc = format!("Parameter {param_name_str}");
472                    quote! { #[schemars(description = #default_desc)] }
473                };
474
475            // Auto-add #[serde(default)] for Option<T> fields
476            let serde_default = if is_option_type(ty) {
477                quote! { #[serde(default)] }
478            } else {
479                quote! {}
480            };
481
482            field_tokens.push(quote! {
483                #field_doc_attr
484                #serde_default
485                #vis #param_name: #ty
486            });
487
488            param_names.push(param_name);
489        }
490    }
491
492    // Default required to all parameters only when required(...) was omitted.
493    let required_args: Vec<String> = args
494        .required
495        .unwrap_or_else(|| param_names.iter().map(|n| n.to_string()).collect());
496
497    let params_struct_name = format_ident!("{}Parameters", struct_name);
498    let static_name = format_ident!("{}", fn_name_str.to_uppercase());
499
500    // Generate the call implementation based on whether the function is async
501    let call_impl = if is_async {
502        quote! {
503            async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
504                #fn_name(#(args.#param_names,)*).await
505            }
506        }
507    } else {
508        quote! {
509            async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
510                #fn_name(#(args.#param_names,)*)
511            }
512        }
513    };
514
515    let rig_core = rig_core_path();
516    let schemars_crate = format!("{}::schemars", rig_core.to_string().replace(' ', ""));
517    let expanded = quote! {
518        #[derive(serde::Deserialize, #rig_core::schemars::JsonSchema)]
519        #[schemars(crate = #schemars_crate)]
520        #vis struct #params_struct_name {
521            #(#field_tokens,)*
522        }
523
524        #cleaned_fn
525
526        #[derive(Default)]
527        #vis struct #struct_name;
528
529        impl #rig_core::tool::Tool for #struct_name {
530            const NAME: &'static str = #tool_name;
531
532            type Args = #params_struct_name;
533            type Output = #output_type;
534            type Error = #error_type;
535
536            fn name(&self) -> String {
537                #tool_name.to_string()
538            }
539
540            fn description(&self) -> String {
541                #tool_description.to_string()
542            }
543
544            fn parameters(&self) -> serde_json::Value {
545                let mut schema = serde_json::to_value(
546                    #rig_core::schemars::schema_for!(#params_struct_name)
547                ).expect("schema serialization");
548                schema["required"] = serde_json::json!([#(#required_args),*]);
549                schema
550            }
551
552            #call_impl
553        }
554
555        #vis static #static_name: #struct_name = #struct_name;
556    };
557
558    TokenStream::from(expanded)
559}
rig_derive/lib.rs

rig_derive/
lib.rs
