kind_openai_schema_impl/
lib.rs

1mod enum_gen;
2mod struct_gen;
3mod utils;
4
5use proc_macro::TokenStream;
6use quote::quote;
7use struct_gen::GenSegment;
8use syn::{parse_macro_input, Data, DeriveInput};
9
10/// Places an associated function on a struct that returns an `&'static str` containing its OpenAI-compatible JSON schema.
11#[proc_macro_derive(OpenAISchema)]
12pub fn openai_schema_derive(input: TokenStream) -> TokenStream {
13    let input = parse_macro_input!(input as DeriveInput);
14
15    match generate_openai_schema(&input) {
16        Ok(expanded) => expanded.into(),
17        Err(err) => err.to_compile_error().into(),
18    }
19}
20
21fn generate_openai_schema(input: &DeriveInput) -> Result<proc_macro2::TokenStream, syn::Error> {
22    let name = &input.ident;
23    // this is the top-level docstring of the struct for the schema description.
24    // individual field docstrings are also extracted.
25    let description = utils::get_description(&input.attrs);
26    let repr = utils::has_repr_attr(&input.attrs)?;
27
28    if utils::has_top_level_serde_attr(&input.attrs) {
29        return Err(syn::Error::new_spanned(
30            &input.ident,
31            "Top-level serde attrs are not supported",
32        ));
33    }
34
35    match &input.data {
36        Data::Struct(data) => {
37            let tokens = struct_gen::handle_struct(data, name, description)?
38                .into_iter()
39                .map(|seg| match seg {
40                    GenSegment::Quote(subordinate_get_schema_method_call) => quote! {
41                        s.push_str(&#subordinate_get_schema_method_call);
42                    },
43                    GenSegment::StringLit(s) => quote! { s.push_str(&#s); },
44                });
45
46            Ok(quote! {
47                impl ::kind_openai::OpenAISchema for #name {
48                    fn openai_schema() -> ::kind_openai::GeneratedOpenAISchema {
49                        use ::kind_openai::SubordinateOpenAISchema;
50                        let mut s = ::std::string::String::new();
51                        #(#tokens)*
52                        s.into()
53                    }
54                }
55            })
56        }
57        Data::Enum(data) => {
58            let schema = serde_json::to_string(&enum_gen::handle_enum(data, repr, description)?)
59                .map_err(|err| syn::Error::new_spanned(&input.ident, err.to_string()))?;
60
61            Ok(quote! {
62                impl ::kind_openai::SubordinateOpenAISchema for #name {
63                    fn subordinate_openai_schema() -> &'static str {
64                        #schema
65                    }
66                }
67            })
68        }
69        _ => Err(syn::Error::new_spanned(
70            &input.ident,
71            "Only structs and enums with unit variants are supported",
72        )),
73    }
74}