llm-schema 0.1.0

A Rust library for type-safe schema conversion between Rust types and LLM.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, DeriveInput, Data, Fields, Type, Meta};

#[proc_macro_derive(LlmSchema, attributes(llmschem))]
pub fn derive_llm_schema(input: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(input as DeriveInput);
    let name = &ast.ident;

    let schema = generate_schema(&ast);

    let expanded = quote! {
        impl #name {
            pub fn llm_schema() -> ::serde_json::Value {
                #schema
            }
        }
    };

    TokenStream::from(expanded)
}

fn generate_schema(ast: &DeriveInput) -> proc_macro2::TokenStream {
    let name = &ast.ident;
    let fields = match &ast.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(fields) => fields.named.iter(),
            _ => unimplemented!("Only named fields are supported"),
        },
        _ => unimplemented!("Only structs are supported"),
    };

    let mut properties = quote! {};
    let mut required_fields = Vec::new();

    for field in fields {
        let field_name = &field.ident;
        let field_name_str = field_name.as_ref().unwrap().to_string();

        // Check for #[llmschem(require)] attribute
        let mut is_required = false; // Changed to false by default to match previous behavior
        for attr in &field.attrs {
            if attr.path().is_ident("llmschem") {
                attr.parse_nested_meta(|meta| {
                    if meta.path.is_ident("require") {
                        is_required = true;
                    }
                    Ok(())
                }).unwrap_or_default();
            }
        }

        // Handle Option<T> types
        if let Type::Path(type_path) = &field.ty {
            if let Some(segment) = type_path.path.segments.last() {
                if segment.ident == "Option" {
                    is_required = false;
                }
            }
        }

        let type_def = get_type_definition(&field.ty);

        properties.extend(quote! {
            #field_name_str: {
                "type": #type_def
            },
        });

        if is_required {
            required_fields.push(field_name_str);
        }
    }

    let required_array = if !required_fields.is_empty() {
        let required = required_fields.iter().map(|s| quote! { #s });
        quote! {
            schema["required"] = ::serde_json::json!([#(#required),*]);
        }
    } else {
        quote! {}
    };

    quote! {
        {
            let mut schema = ::serde_json::json!({
                "type": "object",
                "properties": {
                    #properties
                }
            });

            #required_array

            schema
        }
    }
}

fn get_type_definition(ty: &Type) -> &'static str {
    match ty {
        Type::Path(type_path) => {
            if let Some(segment) = type_path.path.segments.last() {
                match segment.ident.to_string().as_str() {
                    "String" => "string",
                    "str" => "string",
                    "bool" => "boolean",
                    "f32" | "f64" | "i8" | "i16" | "i32" | "i64" | "u8" | "u16" | "u32" | "u64" => "number",
                    "Option" => "null", // Option<T> will be handled in parent function
                    _ => panic!("Unsupported type for LLM schema"),
                }
            } else {
                panic!("Invalid type path");
            }
        }
        _ => panic!("Only simple types are supported"),
    }
}