1mod completion_schema;
2mod schema_extraction;
3
4use proc_macro::TokenStream;
5use quote::quote;
6use syn::parse_macro_input;
7
8#[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 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}