gigachat_macro/
lib.rs

1
2use quote::quote;
3use syn::{DeriveInput, Expr, LitStr, Meta, MetaNameValue};
4use proc_macro::TokenStream;
5
6
7#[proc_macro_attribute]
8pub fn gigafunction(attrs: TokenStream, input: TokenStream) -> TokenStream {
9    let mut input = syn::parse_macro_input!(input as DeriveInput);
10    let attrs_meta = syn::parse_macro_input!(attrs as Meta);
11
12    let struct_name = &input.ident;
13    let func_description = match attrs_meta {
14        Meta::NameValue(MetaNameValue { value: Expr::Lit(expr), .. }) => {
15            match &expr.lit {
16                syn::Lit::Str(lit_str) => lit_str.value(),
17                _ => "Функция для LLM".to_string(),
18            }
19        },
20        _ => "Функция для LLM".to_string()
21    };
22
23    let (properties, required) = if let syn::Data::Struct(syn::DataStruct { fields, ..}) = &mut input.data {
24        let mut properties = Vec::new();
25        let mut required = Vec::new();
26
27        for field in fields.iter_mut() {
28            let field_name = match field.ident.as_ref() {
29                Some(name) => name.to_string(),
30                _ => panic!("Поле структуры заданно некорректно.")
31            };
32            let field_type = &field.ty;
33            let field_description = field.attrs.iter()
34                .find_map(|attr| {
35                    if attr.path().is_ident("description") {
36                        return match attr.parse_args::<LitStr>() {
37                            Ok(res) => Some(res.value()),
38                            Err(_) => None,
39                        };
40                    }
41                    
42                    None
43                }).unwrap_or_else(|| format!("Аргумент функции {}", field_name));
44
45            let is_required = field.attrs.iter().any(|attr| attr.path().is_ident("required"));
46            
47            let field_schema = quote! {
48                {
49                    let mut object = schemars::schema_for!(#field_type).schema;
50                    object.metadata().description = Some(#field_description.to_string());
51                    object
52                }
53            };
54            
55            properties.push(quote! {
56                (#field_name, #field_schema)
57            });
58            if is_required {
59                required.push(
60                    quote! { #field_name }
61                )
62            }
63
64            field.attrs.retain(|attr| {
65                !attr.path().is_ident("description") && !attr.path().is_ident("required")
66            });
67        }
68
69        (properties, required)
70    } else {
71        (Vec::new(), Vec::new())
72    };
73
74    let expanded = quote! {
75        #[derive(serde::Serialize, serde::Deserialize, schemars::JsonSchema, Debug)]
76        #input
77
78        impl #struct_name {
79            pub fn schema(&self) -> serde_json::Value {
80                let mut props = schemars::Map::new();
81                
82                #(
83                    let (name, schema) = #properties;
84                    props.insert(name, serde_json::to_value(schema).unwrap());
85                )*
86
87                serde_json::json!({
88                    "name": stringify!(#struct_name).to_lowercase(),
89                    "description": #func_description,
90                    "parameters": {
91                        "type": "object",
92                        "properties": props,
93                        "required": vec![#(#required),*]
94                    }
95                })
96            }
97        }
98    };
99
100    expanded.into()
101}