Skip to main content

tinyagent_macros/
lib.rs

1mod completion_schema;
2mod schema_extraction;
3
4use proc_macro::TokenStream;
5use quote::quote;
6use syn::parse_macro_input;
7
8/// Defines the `tool!` macro for declaring tools.
9/// Generates a `Tool` impl with JSON Schema from `params`
10/// and wires an async closure as the executor.
11#[proc_macro]
12pub fn tool(input: TokenStream) -> TokenStream {
13    let input = parse_macro_input!(input as ToolDefinition);
14
15    let name = input.name;
16    let description = input.description;
17    let params_type = input.params_type;
18    let execute_body = input.execute_body;
19
20    // Convert snake_case to PascalCase for struct name
21    // tbd if needed long term idk
22    let struct_name = name
23        .value()
24        .split('_')
25        .filter(|part| !part.is_empty())
26        .map(|part| {
27            let mut chars = part.chars();
28            match chars.next() {
29                Some(first) => first.to_uppercase().chain(chars).collect::<String>(),
30                None => String::new(),
31            }
32        })
33        .collect::<String>();
34
35    let tool_struct = quote::format_ident!("{}", struct_name);
36
37    let expanded = quote! {
38        #[derive(Debug)]
39        pub struct #tool_struct;
40
41        impl tiny_agent_rs::tools::Tool for #tool_struct {
42            fn name(&self) -> &'static str {
43                #name
44            }
45
46            fn description(&self) -> &'static str {
47                #description
48            }
49
50            fn parameters_schema(&self) -> serde_json::Value {
51                let schema = schemars::schema_for!(#params_type);
52                serde_json::to_value(&schema.schema).unwrap_or_else(|_| {
53                    serde_json::json!({
54                        "type": "object",
55                        "properties": {},
56                        "required": []
57                    })
58                })
59            }
60
61            fn execute(
62                &self,
63                parameters: serde_json::Value,
64            ) -> std::pin::Pin<
65                Box<
66                    dyn std::future::Future<Output = Result<serde_json::Value, tiny_agent_rs::AgentError>>
67                        + Send
68                        + '_,
69                >,
70            > {
71                Box::pin(async move {
72                    let params: #params_type = serde_json::from_value(parameters)
73                        .map_err(|e| tiny_agent_rs::AgentError::ToolExecution(
74                            format!("Invalid parameters for {}: {}", #name, e)
75                        ))?;
76
77                    let handler = #execute_body;
78                    handler(params)
79                        .await
80                        .map_err(|e| tiny_agent_rs::AgentError::ToolExecution(e))
81                })
82            }
83        }
84    };
85
86    TokenStream::from(expanded)
87}
88
89struct ToolDefinition {
90    name: syn::LitStr,
91    description: syn::LitStr,
92    params_type: syn::Type,
93    execute_body: syn::ExprClosure,
94}
95
96fn parse_named_assignment<T: syn::parse::Parse>(
97    input: syn::parse::ParseStream,
98    keyword: &str,
99) -> syn::Result<T> {
100    let ident: syn::Ident = input.parse()?;
101    if ident != keyword {
102        return Err(syn::Error::new_spanned(
103            ident,
104            format!("expected '{keyword}'"),
105        ));
106    }
107    input.parse::<syn::Token![=]>()?;
108    let value = input.parse::<T>()?;
109    input.parse::<syn::Token![,]>()?;
110    Ok(value)
111}
112
113impl syn::parse::Parse for ToolDefinition {
114    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
115        let name = parse_named_assignment::<syn::LitStr>(input, "name")?;
116        let description = parse_named_assignment::<syn::LitStr>(input, "description")?;
117        let params_type = parse_named_assignment::<syn::Type>(input, "params")?;
118        let execute_body: syn::ExprClosure = input.parse()?;
119
120        Ok(ToolDefinition {
121            name,
122            description,
123            params_type,
124            execute_body,
125        })
126    }
127}
128
129#[proc_macro_attribute]
130pub fn completion_schema(attr: TokenStream, item: TokenStream) -> TokenStream {
131    completion_schema::completion_schema(attr, item)
132}