gemini_proc_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{
4    Attribute, Data, DeriveInput, Fields, FnArg, ItemFn, Lit, Meta, Pat, Type, parse_macro_input,
5};
6
7/// Attribute macro to mark a function as callable by Gemini.
8///
9/// This macro generates a schema for the function and an `execute` method to call it
10/// with JSON arguments.
11///
12/// # Requirements
13/// - Function arguments must be owned types that implement `GeminiSchema` (e.g., `String`, `i32`, `bool`).
14/// - References are not supported.
15/// - The function can be `async` and can return a `Result` (the `Ok` value must implement `Serialize`).
16///
17/// # Example
18/// ```ignore
19/// #[gemini_function]
20/// /// Returns the current weather for a location.
21/// fn get_weather(location: String) -> String {
22///     format!("The weather in {} is sunny.", location)
23/// }
24/// ```
25#[proc_macro_attribute]
26pub fn gemini_function(_attr: TokenStream, item: TokenStream) -> TokenStream {
27    let mut input_fn = parse_macro_input!(item as ItemFn);
28    let fn_name = &input_fn.sig.ident;
29    let args_struct_name = syn::Ident::new(&format!("{}_args", fn_name), fn_name.span());
30    let fn_description = extract_doc_comments(&input_fn.attrs);
31
32    let mut properties = Vec::new();
33    let mut required = Vec::new();
34    let mut param_names = Vec::new();
35    let mut param_types = Vec::new();
36
37    for arg in input_fn.sig.inputs.iter_mut() {
38        if let FnArg::Typed(pat_type) = arg {
39            if let Pat::Ident(pat_ident) = &*pat_type.pat {
40                let param_name = pat_ident.ident.clone();
41                let param_name_str = param_name.to_string();
42                let param_type = (*pat_type.ty).clone();
43                let param_desc = extract_doc_comments(&pat_type.attrs);
44
45                if has_reference(&param_type) {
46                    return syn::Error::new_spanned(
47                        &param_type,
48                        "references are not supported in gemini_function. Use owned types like String instead.",
49                    )
50                    .to_compile_error()
51                    .into();
52                }
53
54                // Remove doc attributes from the function signature so it compiles
55                pat_type.attrs.retain(|attr| !attr.path().is_ident("doc"));
56
57                let is_optional = is_option(&param_type);
58
59                properties.push(quote! {
60                    let mut schema = <#param_type as GeminiSchema>::gemini_schema();
61                    if !#param_desc.is_empty() {
62                        if let Some(obj) = schema.as_object_mut() {
63                            obj.insert("description".to_string(), serde_json::json!(#param_desc));
64                        }
65                    }
66                    props.insert(#param_name_str.to_string(), schema);
67                });
68
69                if !is_optional {
70                    required.push(param_name_str);
71                }
72
73                param_names.push(param_name);
74                param_types.push(param_type);
75            }
76        }
77    }
78
79    let fn_name_str = fn_name.to_string();
80    let is_async = input_fn.sig.asyncness.is_some();
81    let call_await = if is_async {
82        quote! { .await }
83    } else {
84        quote! {}
85    };
86
87    let is_result = match &input_fn.sig.output {
88        syn::ReturnType::Default => false,
89        syn::ReturnType::Type(_, ty) => {
90            let s = quote!(#ty).to_string();
91            s.contains("Result")
92        }
93    };
94
95    let result_handling = if is_result {
96        quote! {
97            match result {
98                Ok(v) => Ok(serde_json::json!(v)),
99                Err(e) => Err(e.to_string()),
100            }
101        }
102    } else {
103        quote! {
104            Ok(serde_json::json!(result))
105        }
106    };
107
108    let expanded = quote! {
109        #input_fn
110
111        #[allow(non_camel_case_types)]
112        pub struct #fn_name { }
113
114        #[allow(non_camel_case_types)]
115        #[derive(gemini_client_api::serde::Deserialize)]
116        pub struct #args_struct_name {
117            #(pub #param_names: #param_types,)*
118        }
119
120        impl GeminiSchema for #fn_name {
121            fn gemini_schema() -> serde_json::Value {
122                use serde_json::{json, Map};
123                let mut props = Map::new();
124                #(#properties)*
125
126                json!({
127                    "name": #fn_name_str,
128                    "description": #fn_description,
129                    "parameters": {
130                        "type": "OBJECT",
131                        "properties": props,
132                        "required": [#(#required),*]
133                    }
134                })
135            }
136        }
137
138        impl #fn_name {
139            pub async fn execute(args: serde_json::Value) -> Result<serde_json::Value, String> {
140                use gemini_client_api::serde::Deserialize;
141                let args = #args_struct_name::deserialize(&args).map_err(|e| e.to_string())?;
142                let result = #fn_name(#(args.#param_names),*) #call_await;
143                #result_handling
144            }
145            pub fn parse_arguments(args: &serde_json::Value) -> Result<(#(#param_types,)*), serde_json::Error>
146            {
147                use gemini_client_api::serde::Deserialize;
148                let args = #args_struct_name::deserialize(args)?;
149                Ok((#(args.#param_names,)*))
150            }
151        }
152    };
153
154    TokenStream::from(expanded)
155}
156
157/// Macro to execute function calls requested by Gemini and update the session history, with a custom callback for results.
158///
159/// # Usage
160/// `execute_function_calls_with_callback!(session, callback, function1, function2, ...)`
161///
162/// The `callback` should be a closure or function that takes `(String, Result<serde_json::Value, String>)`
163/// and returns `serde_json::Value`.
164/// (function_name, result) is passed to it
165///
166/// # Returns
167/// A `Vec<Option<Result<serde_json::Value, String>>>` containing the results of each function call.
168/// - `Some(Ok(value))` if the function was called and succeeded.
169/// - `Some(Err(err))` if the function was called but failed.
170/// - `None` if the function was not called.
171///
172/// # Note
173/// The `session` is automatically updated with the `FunctionResponse`.
174/// The `callback` is invoked with the result of the function execution (whether `Ok` or `Err`)
175/// and its return value is used to update the session.
176#[proc_macro]
177pub fn execute_function_calls_with_callback(input: TokenStream) -> TokenStream {
178    use syn::parse::{Parse, ParseStream};
179    use syn::{Expr, Token};
180
181    struct ExecuteWithCallbackInput {
182        session: Expr,
183        _comma1: Token![,],
184        callback: Expr,
185        _comma2: Token![,],
186        functions: syn::punctuated::Punctuated<syn::Path, Token![,]>,
187    }
188
189    impl Parse for ExecuteWithCallbackInput {
190        fn parse(input: ParseStream) -> syn::Result<Self> {
191            Ok(ExecuteWithCallbackInput {
192                session: input.parse()?,
193                _comma1: input.parse()?,
194                callback: input.parse()?,
195                _comma2: input.parse()?,
196                functions: input.parse_terminated(syn::Path::parse, Token![,])?,
197            })
198        }
199    }
200
201    let input = parse_macro_input!(input as ExecuteWithCallbackInput);
202    generate_execute_logic(&input.session, &input.callback, &input.functions)
203}
204
205/// Attribute macro to derive the `GeminiSchema` trait for a struct or enum.
206///
207/// This allows the type to be used in structured output (`set_json_mode`) or as a parameter
208/// in a `gemini_function`.
209///
210/// # Requirements
211/// - For structs: only named fields are supported.
212/// - For enums: only unit variants (no data) are supported.
213/// - Field/variant types must also implement `GeminiSchema`.
214/// - Doc comments on fields and variants are extracted as descriptions in the schema.
215///
216/// # Example
217/// ```ignore
218/// #[gemini_schema]
219/// struct SearchResult {
220///     /// The title of the page.
221///     title: String,
222///     /// The URL of the page.
223///     url: String,
224/// }
225/// ```
226#[proc_macro_attribute]
227pub fn gemini_schema(_attr: TokenStream, item: TokenStream) -> TokenStream {
228    let input = parse_macro_input!(item as DeriveInput);
229    let name = &input.ident;
230    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
231    let description = extract_doc_comments(&input.attrs);
232
233    let expanded = match &input.data {
234        Data::Struct(data) => {
235            let mut properties = Vec::new();
236            let mut required = Vec::new();
237
238            match &data.fields {
239                Fields::Named(fields) => {
240                    for field in &fields.named {
241                        let field_name = field.ident.as_ref().unwrap();
242                        let field_name_str = field_name.to_string();
243                        let field_type = &field.ty;
244                        let field_desc = extract_doc_comments(&field.attrs);
245
246                        if has_reference(field_type) {
247                            return syn::Error::new_spanned(
248                                field_type,
249                                "references are not supported in gemini_schema. Use owned types instead.",
250                            )
251                            .to_compile_error()
252                            .into();
253                        }
254
255                        let is_optional = is_option(field_type);
256
257                        properties.push(quote! {
258                            let mut schema = <#field_type as GeminiSchema>::gemini_schema();
259                            if !#field_desc.is_empty() {
260                                if let Some(obj) = schema.as_object_mut() {
261                                    obj.insert("description".to_string(), serde_json::json!(#field_desc));
262                                }
263                            }
264                            props.insert(#field_name_str.to_string(), schema);
265                        });
266
267                        if !is_optional {
268                            required.push(field_name_str);
269                        }
270                    }
271                }
272                _ => panic!("gemini_schema only supports named fields in structs"),
273            }
274
275            quote! {
276                impl #impl_generics GeminiSchema for #name #ty_generics #where_clause {
277                    fn gemini_schema() -> serde_json::Value {
278                        use serde_json::{json, Map};
279                        let mut props = Map::new();
280                        #(#properties)*
281
282                        let mut schema = json!({
283                            "type": "OBJECT",
284                            "properties": props,
285                            "required": [#(#required),*]
286                        });
287
288                        if !#description.is_empty() {
289                            if let Some(obj) = schema.as_object_mut() {
290                                obj.insert("description".to_string(), json!(#description));
291                            }
292                        }
293                        schema
294                    }
295                }
296            }
297        }
298        Data::Enum(data) => {
299            let mut variants = Vec::new();
300            for variant in &data.variants {
301                if !matches!(variant.fields, Fields::Unit) {
302                    panic!("gemini_schema only supports unit variants in enums");
303                }
304                variants.push(variant.ident.to_string());
305            }
306
307            quote! {
308                impl #impl_generics GeminiSchema for #name #ty_generics #where_clause {
309                    fn gemini_schema() -> serde_json::Value {
310                        use serde_json::json;
311                        let mut schema = json!({
312                            "type": "STRING",
313                            "enum": [#(#variants),*]
314                        });
315
316                        if !#description.is_empty() {
317                            if let Some(obj) = schema.as_object_mut() {
318                                obj.insert("description".to_string(), json!(#description));
319                            }
320                        }
321                        schema
322                    }
323                }
324            }
325        }
326        _ => panic!("gemini_schema only supports structs and enums"),
327    };
328
329    let output = quote! {
330        #input
331        #expanded
332    };
333
334    TokenStream::from(output)
335}
336
337fn extract_doc_comments(attrs: &[Attribute]) -> String {
338    let mut doc_comments = Vec::new();
339    for attr in attrs {
340        if attr.path().is_ident("doc") {
341            if let Meta::NameValue(nv) = &attr.meta {
342                if let syn::Expr::Lit(expr_lit) = &nv.value {
343                    if let Lit::Str(lit_str) = &expr_lit.lit {
344                        doc_comments.push(lit_str.value().trim().to_string());
345                    }
346                }
347            }
348        }
349    }
350    doc_comments.join("\n")
351}
352
353fn is_option(ty: &Type) -> bool {
354    if let Type::Path(tp) = ty {
355        if let Some(seg) = tp.path.segments.last() {
356            return seg.ident == "Option";
357        }
358    }
359    false
360}
361
362fn has_reference(ty: &Type) -> bool {
363    match ty {
364        Type::Reference(_) => true,
365        Type::Path(tp) => {
366            for seg in &tp.path.segments {
367                if let syn::PathArguments::AngleBracketed(ab) = &seg.arguments {
368                    for arg in &ab.args {
369                        if let syn::GenericArgument::Type(inner) = arg {
370                            if has_reference(inner) {
371                                return true;
372                            }
373                        }
374                    }
375                }
376            }
377            false
378        }
379        _ => false,
380    }
381}
382
383/// Macro to execute function calls requested by Gemini and update the session history.
384///
385/// # Usage
386/// `execute_function_calls!(session, function1, function2, ...)`
387///
388/// # Returns
389/// A `Vec<Option<Result<serde_json::Value, String>>>` containing the results of each function call.
390/// The length of the vector matches the number of functions provided.
391/// - `Some(Ok(value))` if the function was called and succeeded.
392/// - `Some(Err(err))` if the function was called but failed.
393/// - `None` if the function was not called.
394///
395/// # Note
396/// The `session` is automatically updated with the `FunctionResponse` for successful calls.
397/// If a function call fails, the error is converted to a JSON object `{"Error": error_message}`
398/// and sent to the session as the function response.
399#[proc_macro]
400pub fn execute_function_calls(input: TokenStream) -> TokenStream {
401    use syn::parse::{Parse, ParseStream};
402    use syn::{Expr, Token};
403
404    struct ExecuteInput {
405        session: Expr,
406        _comma: Token![,],
407        functions: syn::punctuated::Punctuated<syn::Path, Token![,]>,
408    }
409
410    impl Parse for ExecuteInput {
411        fn parse(input: ParseStream) -> syn::Result<Self> {
412            Ok(ExecuteInput {
413                session: input.parse()?,
414                _comma: input.parse()?,
415                functions: input.parse_terminated(syn::Path::parse, Token![,])?,
416            })
417        }
418    }
419
420    let input = parse_macro_input!(input as ExecuteInput);
421    let callback: Expr = syn::parse_quote! {
422        |_name: String, result: Result<gemini_client_api::serde_json::Value, String>| {
423            match result {
424                Ok(value) => value,
425                Err(e) => gemini_client_api::serde_json::json!({"Error": e}),
426            }
427        }
428    };
429
430    generate_execute_logic(&input.session, &callback, &input.functions)
431}
432
433fn generate_execute_logic(
434    session: &syn::Expr,
435    callback: &syn::Expr,
436    functions: &syn::punctuated::Punctuated<syn::Path, syn::Token![,]>,
437) -> TokenStream {
438    let num_funcs = functions.len();
439
440    let match_arms = functions.iter().enumerate().map(|(i, path)| {
441        let name_str = path.segments.last().unwrap().ident.to_string();
442        quote! {
443            #name_str => {
444                let args = call.args().clone().unwrap_or(gemini_client_api::serde_json::json!({}));
445                let fut: gemini_client_api::futures::future::BoxFuture<'static, (usize, String, Result<gemini_client_api::serde_json::Value, String>)> = Box::pin(async move {
446                    (#i, #name_str.to_string(), #path::execute(args).await)
447                });
448                futures.push(fut);
449            }
450        }
451    });
452
453    let expanded = quote! {
454        {
455            let mut results_array = vec![None; #num_funcs];
456            // Define callback here to ensure it's available
457            let mut result_callback = #callback;
458
459            if let Some(chat) = #session.get_last_chat() {
460                let mut futures = Vec::new();
461                for part in chat.parts() {
462                    if let gemini_client_api::gemini::types::request::PartType::FunctionCall(call) = part.data() {
463                        match call.name().as_str() {
464                            #(#match_arms)*
465                            _ => {}
466                        }
467                    }
468                }
469                if !futures.is_empty() {
470                    let results = gemini_client_api::futures::future::join_all(futures).await;
471                    for (idx, name, res) in results {
472                        // Invoke callback regardless of success or failure
473                        let val_to_add = result_callback(name.clone(), res.clone());
474
475                        if let Err(e) = #session.add_function_response(name.clone(), val_to_add) {
476                             results_array[idx] = Some(Err(format!(
477                                "failed to add function response for `{}`: {}",
478                                name, e
479                            )));
480                            continue;
481                        }
482                        results_array[idx] = Some(res);
483                    }
484                }
485            }
486            results_array
487        }
488    };
489
490    TokenStream::from(expanded)
491}
gemini_proc_macros/lib.rs

gemini_proc_macros/
lib.rs
