llm_chain_macros/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use syn::{
5    parse_macro_input, DeriveInput, Ident, LitStr,
6    __private::{quote::quote, TokenStream2},
7};
8
9fn literal_from_ident(ident: Ident) -> LitStr {
10    let value = ident.to_string();
11    let span = ident.span();
12    LitStr::new(&value, span)
13}
14
15#[proc_macro_derive(Describe, attributes(purpose))]
16pub fn derive_describe(input: TokenStream) -> TokenStream {
17    let input = parse_macro_input!(input as DeriveInput);
18    let syn::Data::Struct(described_struct) = input.data else {
19        panic!("Can only describe structs")
20    };
21
22    // Parse field attrs
23    let pairs: Vec<(LitStr, LitStr)> = described_struct
24        .fields
25        .iter()
26        .map(|field| {
27            let ident = field
28                .ident
29                .clone()
30                .expect("All struct fields must be named");
31            (
32                literal_from_ident(ident),
33                field
34                    .attrs
35                    .iter()
36                    .filter(|attr| {
37                        attr.path().segments.len() == 1
38                            && attr.path().segments[0].ident == "purpose"
39                    })
40                    .nth(0)
41                    .expect("All fields on the string must have a purpose annotation")
42                    .parse_args::<LitStr>()
43                    .expect("Purpose must be a single string literal"),
44            )
45        })
46        .collect();
47    if pairs.len() == 0 {
48        panic!("You need to annotate each field");
49    }
50
51    // Generate FormatParts from fields
52    let mut format_parts = Vec::new();
53    for (key, purpose) in pairs.into_iter() {
54        let gen: TokenStream2 = quote! {
55            FormatPart {
56                key: #key.to_string(),
57                purpose: #purpose.to_string()
58            }
59        }
60        .into();
61
62        format_parts.push(gen);
63    }
64
65    // Implement trait using generated FormatParts
66    let name = &input.ident;
67
68    let gen: TokenStream = quote! (
69        impl Describe for #name {
70            fn describe() -> Format {
71                 Format {
72                    parts: vec![
73                        #(#format_parts),*
74                    ]
75                }
76            }
77        }
78    )
79    .into();
80
81    gen
82}