vex_macros/
lib.rs

1extern crate proc_macro;
2use proc_macro::TokenStream;
3use quote::{format_ident, quote};
4use syn::{parse_macro_input, DeriveInput, ItemFn};
5
6/// Auto-implements the `Job` trait for a struct.
7///
8/// Usage:
9/// ```ignore
10/// #[derive(VexJob)]
11/// struct MyJob { ... }
12/// ```
13#[proc_macro_derive(VexJob, attributes(job))]
14pub fn derive_vex_job(input: TokenStream) -> TokenStream {
15    let input = parse_macro_input!(input as DeriveInput);
16    let name = input.ident;
17
18    // Use struct name as job name (simplified - no attribute parsing)
19    let job_name = name.to_string().to_lowercase();
20    let max_retries = 3u32;
21
22    let expanded = quote! {
23        #[async_trait::async_trait]
24        impl vex_queue::job::Job for #name {
25            fn name(&self) -> &str {
26                #job_name
27            }
28
29            async fn execute(&mut self) -> vex_queue::job::JobResult {
30                self.run().await
31            }
32
33            fn max_retries(&self) -> u32 {
34                #max_retries
35            }
36
37            fn backoff_strategy(&self) -> vex_queue::job::BackoffStrategy {
38                vex_queue::job::BackoffStrategy::Exponential {
39                    initial_secs: 1,
40                    multiplier: 2.0
41                }
42            }
43        }
44    };
45
46    TokenStream::from(expanded)
47}
48
49/// Generates a ToolDefinition constant for an LLM tool function.
50///
51/// Usage:
52/// ```ignore
53/// #[vex_tool]
54/// fn web_search(query: String) -> String { ... }
55/// ```
56///
57/// Generates a `WEB_SEARCH_TOOL` constant.
58#[proc_macro_attribute]
59pub fn vex_tool(_args: TokenStream, item: TokenStream) -> TokenStream {
60    let input = parse_macro_input!(item as ItemFn);
61    let fn_name = &input.sig.ident;
62    let tool_name = fn_name.to_string();
63    let tool_desc = "Auto-generated tool"; // Can be improved by parsing doc comments
64
65    let mut props_map = std::collections::HashMap::new();
66    let mut req_list = Vec::new();
67
68    for arg in &input.sig.inputs {
69        if let syn::FnArg::Typed(pat_type) = arg {
70            if let syn::Pat::Ident(pat_ident) = &*pat_type.pat {
71                let arg_name = pat_ident.ident.to_string();
72                let arg_type = &*pat_type.ty;
73
74                let (json_type, is_optional) = match get_json_type_static(arg_type) {
75                    Some((t, opt)) => (t, opt),
76                    None => ("string", false),
77                };
78
79                props_map.insert(arg_name.clone(), json_type);
80                if !is_optional {
81                    req_list.push(arg_name);
82                }
83            }
84        }
85    }
86
87    let parameters = if props_map.is_empty() {
88        "{}".to_string()
89    } else {
90        let mut props_vec: Vec<_> = props_map.iter().collect();
91        props_vec.sort_by_key(|a| a.0); // Sort for deterministic output
92
93        let props_str = props_vec
94            .iter()
95            .map(|(k, v)| format!("\"{}\":{{\"type\":\"{}\"}}", k, v))
96            .collect::<Vec<_>>()
97            .join(",");
98
99        let mut req_vec = req_list.clone();
100        req_vec.sort();
101
102        let req_str = req_vec
103            .iter()
104            .map(|k| format!("\"{}\"", k))
105            .collect::<Vec<_>>()
106            .join(",");
107        format!(
108            "{{\"type\":\"object\",\"properties\":{{{}}},\"required\":[{}]}}",
109            props_str, req_str
110        )
111    };
112
113    let const_name = format_ident!("{}_TOOL", tool_name.to_uppercase());
114
115    let expanded = quote! {
116        #input
117
118        pub const #const_name: vex_llm::ToolDefinition = vex_llm::ToolDefinition {
119            name: #tool_name,
120            description: #tool_desc,
121            parameters: #parameters,
122        };
123    };
124
125    TokenStream::from(expanded)
126}
127
128fn get_json_type_static(ty: &syn::Type) -> Option<(&'static str, bool)> {
129    match ty {
130        syn::Type::Path(tp) => {
131            let last = tp.path.segments.last()?;
132            let ident = last.ident.to_string();
133            match ident.as_str() {
134                "String" | "str" => Some(("string", false)),
135                "i32" | "i64" | "u32" | "u64" | "isize" | "usize" => Some(("integer", false)),
136                "f32" | "f64" => Some(("number", false)),
137                "bool" => Some(("boolean", false)),
138                "Option" => {
139                    if let syn::PathArguments::AngleBracketed(args) = &last.arguments {
140                        if let Some(syn::GenericArgument::Type(inner)) = args.args.first() {
141                            let (inner_type, _) = get_json_type_static(inner)?;
142                            return Some((inner_type, true));
143                        }
144                    }
145                    None
146                }
147                _ => None,
148            }
149        }
150        _ => None,
151    }
152}
153
154/// Instruments an agent function with tracing.
155///
156/// Usage:
157/// ```ignore
158/// #[instrument_agent]
159/// async fn think(&self) { ... }
160/// ```
161#[proc_macro_attribute]
162pub fn instrument_agent(_args: TokenStream, item: TokenStream) -> TokenStream {
163    let input = parse_macro_input!(item as ItemFn);
164    let fn_name = &input.sig.ident;
165    let block = &input.block;
166    let sig = &input.sig;
167    let vis = &input.vis;
168    let attrs = &input.attrs;
169
170    let expanded = quote! {
171        #(#attrs)*
172        #vis #sig {
173            let span = tracing::info_span!(
174                stringify!(#fn_name),
175                agent_id = %self.id,
176                generation = %self.generation
177            );
178            let _enter = span.enter();
179            #block
180        }
181    };
182
183    TokenStream::from(expanded)
184}
vex_macros/lib.rs

vex_macros/
lib.rs
