kernelx_macros/

lib.rs

use proc_macro::TokenStream;
use proc_macro_crate::{crate_name, FoundCrate};
use quote::quote;
use syn::{parse::Parse, parse::ParseStream, parse_macro_input, DeriveInput, Ident, Token, Type};

fn get_crate_name() -> proc_macro2::TokenStream {
    let found_crate =
        crate_name("kernelx-core").expect("kernelx-core should be present in `Cargo.toml`");
    match found_crate {
        FoundCrate::Itself => quote!(crate),
        FoundCrate::Name(name) => {
            let ident = proc_macro2::Ident::new(&name, proc_macro2::Span::call_site());
            quote!(#ident)
        }
    }
}

// Schema derive macro
#[proc_macro_attribute]
pub fn lmp_schema(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let expanded = quote! {
        #[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
        #[serde(rename_all = "camelCase")]
        #input
    };
    TokenStream::from(expanded)
}

// Provider configuration macro
#[proc_macro]
pub fn provider(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ProviderDef);
    // let crate_name = get_crate_name();

    let provider_type = &input.provider_type;
    let api_base = &input.api_base;
    let api_key = &input.api_key;
    let models = &input.models;

    // Check if the type is a path with segments
    let provider_path = match provider_type {
        Type::Path(type_path) if type_path.path.segments.len() > 1 => {
            quote!(#provider_type)
        }
        _ => {
            quote!(crate::#provider_type)
        }
    };

    let expanded = quote! {
        thread_local! {
            static PROVIDER: std::sync::OnceLock<#provider_path> = {
                let cell = std::sync::OnceLock::new();
                let provider = #provider_path::builder()
                    .api_base(#api_base)
                    .api_key(#api_key)
                    .models(#models)
                    .build()
                    .expect("Failed to create provider");
                cell.set(provider).unwrap();
                cell
            };
        }

        fn get_provider() -> #provider_path {
            PROVIDER.with(|p| p.get().unwrap().clone())
        }
    };

    TokenStream::from(expanded)
}

// Provider definition parser
struct ProviderDef {
    provider_type: syn::Type,
    api_base: syn::LitStr,
    api_key: syn::LitStr,
    models: syn::Expr, // Using Expr to capture the entire models![] macro invocation
}

impl Parse for ProviderDef {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let provider_type: syn::Type = input.parse()?;

        input.parse::<Token![,]>()?;
        input.parse::<syn::Ident>()?; // parse "api_base:"
        input.parse::<Token![:]>()?;
        let api_base: syn::LitStr = input.parse()?;

        input.parse::<Token![,]>()?;
        input.parse::<syn::Ident>()?; // parse "api_key:"
        input.parse::<Token![:]>()?;
        let api_key: syn::LitStr = input.parse()?;

        input.parse::<Token![,]>()?;
        input.parse::<syn::Ident>()?; // parse "models:"
        input.parse::<Token![:]>()?;
        let models: syn::Expr = input.parse()?;

        Ok(ProviderDef {
            provider_type,
            api_base,
            api_key,
            models,
        })
    }
}

// LMP function macro
#[proc_macro_attribute]
pub fn lmp(args: TokenStream, input: TokenStream) -> TokenStream {
    let args = parse_macro_input!(args as LmpArgs);
    let input_fn = parse_macro_input!(input as syn::ItemFn);

    let system_prompt = input_fn
        .attrs
        .iter()
        .filter(|attr| attr.path().is_ident("doc"))
        .filter_map(|attr| {
            if let syn::Meta::NameValue(meta) = &attr.meta {
                if let syn::Expr::Lit(syn::ExprLit {
                    lit: syn::Lit::Str(s),
                    ..
                }) = &meta.value
                {
                    Some(s.value())
                } else {
                    None
                }
            } else {
                None
            }
        })
        .collect::<Vec<_>>()
        .join("\n");

    let system_prompt_literal = syn::LitStr::new(&system_prompt, proc_macro2::Span::call_site());
    let model = args.model;
    let temperature = args.temperature.unwrap_or(0.0);
    let max_tokens = args.max_tokens.unwrap_or(512);

    let fn_vis = &input_fn.vis;
    let fn_sig = &input_fn.sig;
    let original_body = &input_fn.block;
    let crate_name = get_crate_name();

    let expanded = if let Some(rf_type) = args.response_format {
        quote! {
            #fn_vis #fn_sig {
                use #crate_name::{StructuredLM, Result, Error};

                let prompt = { #original_body };
                let model = get_provider().get_model::<dyn StructuredLM>(#model)?
                    .system_prompt(#system_prompt_literal)
                    .temperature(#temperature)
                    .max_tokens(#max_tokens);

                let schema = schemars::schema_for!(#rf_type);
                let json = model.structured_complete(&prompt, &serde_json::to_value(schema.schema)?)
                    .await?;
                serde_json::from_value(json).map_err(Error::Serialization)
            }
        }
    } else {
        quote! {
            #fn_vis #fn_sig {
                use #crate_name::{LLM, Result};

                let prompt = { #original_body };
                let model = get_provider().get_model::<dyn LLM>(#model)?
                    .system_prompt(#system_prompt_literal)
                    .temperature(#temperature)
                    .max_tokens(#max_tokens);

                model.complete(&prompt).await
            }
        }
    };

    TokenStream::from(expanded)
}

// Args parser for LMP
#[derive(Default)]
struct LmpArgs {
    model: String,
    temperature: Option<f32>,
    max_tokens: Option<u32>,
    response_format: Option<Type>,
}

impl Parse for LmpArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut args = LmpArgs::default();
        while !input.is_empty() {
            let key: Ident = input.parse()?;
            input.parse::<Token![=]>()?;
            match key.to_string().as_str() {
                "model" => args.model = input.parse::<syn::LitStr>()?.value(),
                "temperature" => {
                    args.temperature = Some(input.parse::<syn::LitFloat>()?.base10_parse()?)
                }
                "max_tokens" => {
                    args.max_tokens = Some(input.parse::<syn::LitInt>()?.base10_parse()?)
                }
                "response_format" => args.response_format = Some(input.parse()?),
                _ => return Err(syn::Error::new(key.span(), format!("Unknown key: {}", key))),
            }
            if input.peek(Token![,]) {
                input.parse::<Token![,]>()?;
            }
        }
        Ok(args)
    }
}

#[proc_macro_attribute]
pub fn json_schema(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as DeriveInput);
    let name = &input.ident;

    let expanded = quote! {
        #[derive(serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
        #input

        impl #name {
            pub fn schema() -> serde_json::Value {
                let schema = schemars::schema_for!(Self);
                serde_json::to_value(schema.schema).expect("Failed to serialize schema")
            }

            pub fn from_json(value: serde_json::Value) -> Result<Self> {
                serde_json::from_value(value).map_err(Error::Serialization)
            }
        }
    };

    TokenStream::from(expanded)
}