Skip to main content

rig_derive/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use syn::{DeriveInput, parse_macro_input};
5
6mod embed;
7mod resolve;
8mod tool;
9
10//References:
11//<https://doc.rust-lang.org/book/ch19-06-macros.html#how-to-write-a-custom-derive-macro>
12//<https://doc.rust-lang.org/reference/procedural-macros.html>
13/// A macro that allows you to implement the `rig::embedding::Embed` trait by deriving it.
14/// Usage can be found below:
15///
16/// ```text
17/// use rig::Embed;
18/// use rig_derive::Embed;
19///
20/// #[derive(Embed)]
21/// struct Foo {
22///     id: String,
23///     #[embed] // this helper shows which field to embed
24///     description: String
25///}
26/// ```
27#[proc_macro_derive(Embed, attributes(embed))]
28pub fn derive_embedding_trait(item: TokenStream) -> TokenStream {
29    let mut input = parse_macro_input!(item as DeriveInput);
30
31    embed::expand_derive_embedding(&mut input)
32        .unwrap_or_else(syn::Error::into_compile_error)
33        .into()
34}
35
36/// A procedural macro that transforms a function into a portable
37/// `rig_core::tool::PortableTool`, or into the classic contextual
38/// `rig::tool::Tool` when the function accepts classic runtime context.
39///
40/// # Examples
41///
42/// Basic usage:
43/// ```text
44/// use rig_derive::rig_tool;
45///
46/// #[rig_tool]
47/// fn add(a: i32, b: i32) -> Result<i32, rig::tool::ToolExecutionError> {
48///     Ok(a + b)
49/// }
50/// ```
51///
52/// With description:
53/// ```text
54/// use rig_derive::rig_tool;
55///
56/// #[rig_tool(description = "Perform basic arithmetic operations")]
57/// fn calculator(x: i32, y: i32, operation: String) -> Result<i32, rig::tool::ToolExecutionError> {
58///     match operation.as_str() {
59///         "add" => Ok(x + y),
60///         "subtract" => Ok(x - y),
61///         "multiply" => Ok(x * y),
62///         "divide" => Ok(x / y),
63///         _ => Err(rig::tool::ToolExecutionError::other("Unknown operation")),
64///     }
65/// }
66/// ```
67///
68/// With a custom tool name:
69/// ```text
70/// use rig_derive::rig_tool;
71///
72/// // Explicit names must be string literals that start with an ASCII letter
73/// // or `_`, may contain ASCII letters, digits, `_`, or `-`, and be at most
74/// // 64 characters long.
75/// #[rig_tool(name = "search-docs", description = "Search the documentation")]
76/// fn search_docs_impl(query: String) -> Result<String, rig::tool::ToolExecutionError> {
77///     Ok(format!("Searching docs for {query}"))
78/// }
79/// ```
80///
81/// With parameter descriptions:
82/// ```text
83/// use rig_derive::rig_tool;
84///
85/// #[rig_tool(
86///     description = "A tool that performs string operations",
87///     params(
88///         text = "The input text to process",
89///         operation = "The operation to perform (uppercase, lowercase, reverse)"
90///     )
91/// )]
92/// fn string_processor(text: String, operation: String) -> Result<String, rig::tool::ToolExecutionError> {
93///     match operation.as_str() {
94///         "uppercase" => Ok(text.to_uppercase()),
95///         "lowercase" => Ok(text.to_lowercase()),
96///         "reverse" => Ok(text.chars().rev().collect()),
97///         _ => Err(rig::tool::ToolExecutionError::other("Unknown operation")),
98///     }
99/// }
100/// ```
101///
102/// # Required parameters
103///
104/// Required-ness is derived from the parameter types: every non-`Option`
105/// parameter is required, and `Option<T>` parameters are optional (absent
106/// fields deserialize to `None`). An explicit `required(...)` list overrides
107/// this. A parameter *omitted* from an explicit list is deserialized with
108/// `#[serde(default)]`, so its type must be `Option<T>` or implement
109/// `Default` — the advertised schema and the deserializer always agree.
110/// Listing an `Option<T>` parameter is a compile error (schemars and serde
111/// would both silently ignore the directive). Names in `params(...)` and
112/// `required(...)` must match actual parameters.
113///
114/// ```text
115/// use rig_derive::rig_tool;
116///
117/// // `b` is advertised as optional and defaults to `0` when omitted.
118/// #[rig_tool(required(a))]
119/// fn add(a: i64, b: i64) -> Result<i64, rig::tool::ToolExecutionError> {
120///     Ok(a + b)
121/// }
122/// ```
123///
124/// # Execution context
125///
126/// ```text
127/// use rig::tool::ToolContext;
128/// use rig_derive::rig_tool;
129///
130/// #[rig_tool]
131/// fn current_user(
132///     // The marker is required for imported names and type aliases. A fully
133///     // qualified `&mut rig::tool::ToolContext` — including under a renamed
134///     // dependency — is recognized directly.
135///     #[rig(context)] context: &mut ToolContext,
136///     greeting: String,
137/// ) -> Result<String, rig::tool::ToolExecutionError> {
138///     let user = context
139///         .get::<String>()
140///         .map(String::as_str)
141///         .unwrap_or("guest");
142///     Ok(format!("{greeting}, {user}!"))
143/// }
144/// ```
145#[proc_macro_attribute]
146pub fn rig_tool(args: TokenStream, input: TokenStream) -> TokenStream {
147    let args = parse_macro_input!(args as tool::args::MacroArgs);
148    let input_fn = parse_macro_input!(input as syn::ItemFn);
149
150    tool::expand::expand_rig_tool(args, input_fn)
151        .unwrap_or_else(syn::Error::into_compile_error)
152        .into()
153}