cortexai_macros/

lib.rs

//! # Procedural Macros for cortex
//!
//! This crate provides derive macros for easily defining tools and other
//! agent components with minimal boilerplate.
//!
//! ## Tool Derive Macro
//!
//! ```rust,ignore
//! use cortexai_macros::Tool;
//!
//! #[derive(Tool)]
//! #[tool(name = "calculator", description = "Perform mathematical calculations")]
//! struct CalculatorTool {
//!     #[tool(param, required, description = "The mathematical expression to evaluate")]
//!     expression: String,
//!
//!     #[tool(param, description = "Number of decimal places for result")]
//!     precision: Option<u32>,
//! }
//!
//! impl CalculatorTool {
//!     async fn run(&self, expression: String, precision: Option<u32>) -> Result<serde_json::Value, String> {
//!         // Implementation here
//!         Ok(serde_json::json!(42.0))
//!     }
//! }
//! ```

use darling::{ast, FromDeriveInput, FromField};
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse::Parser, parse_macro_input, DeriveInput, FnArg, Ident, ItemFn, Meta, Type};

/// Attributes for the Tool derive macro at struct level
#[derive(Debug, FromDeriveInput)]
#[darling(attributes(tool), supports(struct_named))]
struct ToolArgs {
    ident: Ident,
    data: ast::Data<(), ToolField>,

    /// Tool name (defaults to snake_case of struct name)
    #[darling(default)]
    name: Option<String>,

    /// Tool description
    #[darling(default)]
    description: Option<String>,

    /// Whether the tool is dangerous and requires confirmation
    #[darling(default)]
    dangerous: bool,
}

/// Attributes for tool parameters (struct fields)
#[derive(Debug, FromField)]
#[darling(attributes(tool))]
struct ToolField {
    ident: Option<Ident>,
    ty: Type,

    /// Mark this field as a tool parameter
    #[darling(default)]
    param: bool,

    /// Whether this parameter is required
    #[darling(default)]
    required: bool,

    /// Parameter description
    #[darling(default)]
    description: Option<String>,

    /// Skip this field (not a parameter)
    #[darling(default)]
    skip: bool,
}

/// Derive macro for creating Tool implementations
///
/// # Example
///
/// ```rust,ignore
/// #[derive(Tool)]
/// #[tool(name = "search", description = "Search the web")]
/// struct SearchTool {
///     #[tool(param, required, description = "Search query")]
///     query: String,
///
///     #[tool(param, description = "Maximum results to return")]
///     max_results: Option<u32>,
/// }
///
/// impl SearchTool {
///     // You must implement this method
///     async fn run(&self, query: String, max_results: Option<u32>) -> Result<serde_json::Value, String> {
///         Ok(serde_json::json!({"results": []}))
///     }
/// }
/// ```
#[proc_macro_derive(Tool, attributes(tool))]
pub fn derive_tool(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let args = match ToolArgs::from_derive_input(&input) {
        Ok(args) => args,
        Err(e) => return e.write_errors().into(),
    };

    let expanded = generate_tool_impl(&args);

    TokenStream::from(expanded)
}

fn generate_tool_impl(args: &ToolArgs) -> proc_macro2::TokenStream {
    let struct_name = &args.ident;

    // Generate tool name (default to snake_case of struct name without "Tool" suffix)
    let tool_name = args.name.clone().unwrap_or_else(|| {
        let name = struct_name.to_string();
        let name = name.strip_suffix("Tool").unwrap_or(&name);
        to_snake_case(name)
    });

    let description = args
        .description
        .clone()
        .unwrap_or_else(|| format!("{} tool", tool_name));

    let dangerous = args.dangerous;

    // Extract parameter fields
    let fields = match &args.data {
        ast::Data::Struct(fields) => fields,
        _ => panic!("Tool derive only supports structs"),
    };

    let params: Vec<_> = fields
        .fields
        .iter()
        .filter(|f| f.param && !f.skip)
        .collect();

    // Generate JSON schema for parameters
    let param_properties = generate_param_properties(&params);
    let required_params = generate_required_params(&params);

    // Generate argument extraction code
    let arg_extractions = generate_arg_extractions(&params);

    // Generate parameter names for calling run()
    let param_names: Vec<_> = params.iter().map(|f| f.ident.as_ref().unwrap()).collect();

    quote! {
        #[async_trait::async_trait]
        impl cortexai_core::tool::Tool for #struct_name {
            fn schema(&self) -> cortexai_core::tool::ToolSchema {
                let mut __properties = serde_json::Map::new();
                #(#param_properties)*

                cortexai_core::tool::ToolSchema {
                    name: #tool_name.to_string(),
                    description: #description.to_string(),
                    parameters: serde_json::json!({
                        "type": "object",
                        "properties": serde_json::Value::Object(__properties),
                        "required": [#(#required_params),*]
                    }),
                    dangerous: #dangerous,
                    metadata: std::collections::HashMap::new(),
                    required_scopes: vec![],
                }
            }

            async fn execute(
                &self,
                _context: &cortexai_core::tool::ExecutionContext,
                arguments: serde_json::Value,
            ) -> Result<serde_json::Value, cortexai_core::errors::ToolError> {
                #(#arg_extractions)*

                self.run(#(#param_names),*).await
                    .map_err(|e| cortexai_core::errors::ToolError::ExecutionFailed(e.to_string()))
            }
        }
    }
}

fn generate_param_properties(params: &[&ToolField]) -> Vec<proc_macro2::TokenStream> {
    params
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap().to_string();
            let description = field
                .description
                .clone()
                .unwrap_or_else(|| format!("Parameter: {}", name));
            let effective_ty = unwrap_option_type(&field.ty).unwrap_or(&field.ty);
            let schema_tokens = type_to_json_schema(effective_ty);

            quote! {
                {
                    let mut __prop_schema = #schema_tokens;
                    if let serde_json::Value::Object(ref mut m) = __prop_schema {
                        m.insert("description".to_string(), serde_json::Value::String(#description.to_string()));
                    }
                    __properties.insert(#name.to_string(), __prop_schema);
                }
            }
        })
        .collect()
}

fn generate_required_params(params: &[&ToolField]) -> Vec<proc_macro2::TokenStream> {
    params
        .iter()
        .filter(|f| f.required)
        .map(|field| {
            let name = field.ident.as_ref().unwrap().to_string();
            quote! { #name }
        })
        .collect()
}

fn generate_arg_extractions(params: &[&ToolField]) -> Vec<proc_macro2::TokenStream> {
    params
        .iter()
        .map(|field| {
            let ident = field.ident.as_ref().unwrap();
            let name = ident.to_string();
            let ty = &field.ty;

            if is_option_type(ty) {
                quote! {
                    let #ident: #ty = arguments.get(#name)
                        .and_then(|v| serde_json::from_value(v.clone()).ok());
                }
            } else if field.required {
                quote! {
                    let #ident: #ty = {
                        let val = arguments.get(#name)
                            .ok_or_else(|| cortexai_core::errors::ToolError::InvalidArguments(
                                format!("Missing required parameter: {}", #name)
                            ))?
                            .clone();
                        serde_json::from_value(val)
                            .map_err(|e| cortexai_core::errors::ToolError::InvalidArguments(
                                format!("Invalid type for {}: {}", #name, e)
                            ))?
                    };
                }
            } else {
                // Non-required, non-option field - use default
                quote! {
                    let #ident: #ty = arguments.get(#name)
                        .and_then(|v| serde_json::from_value(v.clone()).ok())
                        .unwrap_or_default();
                }
            }
        })
        .collect()
}

/// Generate a `proc_macro2::TokenStream` that evaluates to a `serde_json::Value`
/// representing the JSON Schema for the given Rust type.
fn type_to_json_schema(ty: &Type) -> proc_macro2::TokenStream {
    let type_name = extract_type_name(ty);

    match type_name.as_str() {
        "String" | "&str" | "str" => quote! { serde_json::json!({"type": "string"}) },
        "bool" => quote! { serde_json::json!({"type": "boolean"}) },
        "f32" | "f64" => quote! { serde_json::json!({"type": "number"}) },
        "i8" | "i16" | "i32" | "i64" | "i128" | "isize"
        | "u8" | "u16" | "u32" | "u64" | "u128" | "usize" => {
            quote! { serde_json::json!({"type": "integer"}) }
        }
        "Vec" | "Array" => {
            let items_schema = extract_first_generic_arg(ty)
                .map(|inner| type_to_json_schema(inner))
                .unwrap_or_else(|| quote! { serde_json::json!({}) });
            quote! {
                {
                    let __items = #items_schema;
                    serde_json::json!({"type": "array", "items": __items})
                }
            }
        }
        "HashMap" | "BTreeMap" => {
            let value_schema = extract_second_generic_arg(ty)
                .map(|inner| type_to_json_schema(inner))
                .unwrap_or_else(|| quote! { serde_json::json!({}) });
            quote! {
                {
                    let __additional = #value_schema;
                    serde_json::json!({"type": "object", "additionalProperties": __additional})
                }
            }
        }
        "Option" => {
            // If Option wasn't already unwrapped, handle it here
            extract_first_generic_arg(ty)
                .map(|inner| type_to_json_schema(inner))
                .unwrap_or_else(|| quote! { serde_json::json!({"type": "object"}) })
        }
        _ => quote! { serde_json::json!({"type": "object"}) },
    }
}

/// Extract the outermost type name (e.g. "Vec" from `Vec<String>`, "String" from `String`)
fn extract_type_name(ty: &Type) -> String {
    if let Type::Path(type_path) = ty {
        if let Some(segment) = type_path.path.segments.last() {
            return segment.ident.to_string();
        }
    }
    String::new()
}

/// Extract the first generic type argument (e.g. `String` from `Vec<String>`)
fn extract_first_generic_arg(ty: &Type) -> Option<&Type> {
    if let Type::Path(type_path) = ty {
        if let Some(segment) = type_path.path.segments.last() {
            if let syn::PathArguments::AngleBracketed(ref args) = segment.arguments {
                if let Some(syn::GenericArgument::Type(inner)) = args.args.first() {
                    return Some(inner);
                }
            }
        }
    }
    None
}

/// Extract the second generic type argument (e.g. `V` from `HashMap<K, V>`)
fn extract_second_generic_arg(ty: &Type) -> Option<&Type> {
    if let Type::Path(type_path) = ty {
        if let Some(segment) = type_path.path.segments.last() {
            if let syn::PathArguments::AngleBracketed(ref args) = segment.arguments {
                if let Some(syn::GenericArgument::Type(inner)) = args.args.iter().nth(1) {
                    return Some(inner);
                }
            }
        }
    }
    None
}

/// If the type is `Option<T>`, return `Some(&T)`. Otherwise `None`.
fn unwrap_option_type(ty: &Type) -> Option<&Type> {
    if extract_type_name(ty) == "Option" {
        extract_first_generic_arg(ty)
    } else {
        None
    }
}

fn is_option_type(ty: &Type) -> bool {
    extract_type_name(ty) == "Option"
}

fn to_snake_case(s: &str) -> String {
    let mut result = String::new();
    for (i, c) in s.chars().enumerate() {
        if c.is_uppercase() {
            if i > 0 {
                result.push('_');
            }
            result.push(c.to_ascii_lowercase());
        } else {
            result.push(c);
        }
    }
    result
}

fn to_pascal_case(s: &str) -> String {
    s.split('_')
        .map(|part| {
            let mut chars = part.chars();
            match chars.next() {
                None => String::new(),
                Some(first) => {
                    let upper: String = first.to_uppercase().collect();
                    upper + &chars.collect::<String>()
                }
            }
        })
        .collect()
}

/// Parsed parameter info from a function argument with `#[param(...)]` attributes.
struct FnParam {
    ident: Ident,
    ty: Type,
    description: Option<String>,
    required: bool,
    name_override: Option<String>,
    is_option: bool,
}

fn parse_param_attrs(attrs: &[syn::Attribute]) -> (Option<String>, bool, Option<String>) {
    let mut description = None;
    let mut required = false;
    let mut name_override = None;

    for attr in attrs {
        if !attr.path().is_ident("param") {
            continue;
        }
        if let Meta::List(meta_list) = &attr.meta {
            let tokens = meta_list.tokens.clone();
            let parser = syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated;
            if let Ok(nested) = parser.parse2(tokens) {
                for meta in &nested {
                    match meta {
                        Meta::Path(path) if path.is_ident("required") => {
                            required = true;
                        }
                        Meta::NameValue(nv) if nv.path.is_ident("description") => {
                            if let syn::Expr::Lit(syn::ExprLit {
                                lit: syn::Lit::Str(s),
                                ..
                            }) = &nv.value
                            {
                                description = Some(s.value());
                            }
                        }
                        Meta::NameValue(nv) if nv.path.is_ident("name") => {
                            if let syn::Expr::Lit(syn::ExprLit {
                                lit: syn::Lit::Str(s),
                                ..
                            }) = &nv.value
                            {
                                name_override = Some(s.value());
                            }
                        }
                        _ => {}
                    }
                }
            }
        }
    }
    (description, required, name_override)
}

fn extract_fn_params(sig: &syn::Signature) -> Vec<FnParam> {
    sig.inputs
        .iter()
        .filter_map(|arg| {
            if let FnArg::Typed(pat_type) = arg {
                if let syn::Pat::Ident(pat_ident) = pat_type.pat.as_ref() {
                    let ident = pat_ident.ident.clone();
                    let ty = *pat_type.ty.clone();
                    let is_option = is_option_type(&ty);
                    let (description, explicit_required, name_override) =
                        parse_param_attrs(&pat_type.attrs);
                    // Default: required unless Option<T>
                    let required = explicit_required || !is_option;
                    return Some(FnParam {
                        ident,
                        ty,
                        description,
                        required,
                        name_override,
                        is_option,
                    });
                }
            }
            None
        })
        .collect()
}

/// Attribute macro for defining tools from async functions.
///
/// Generates a `<PascalCaseFnName>Tool` struct implementing the `Tool` trait.
/// Use `#[param(...)]` on function parameters for schema metadata.
#[proc_macro_attribute]
pub fn tool(attr: TokenStream, item: TokenStream) -> TokenStream {
    let attr_args = proc_macro2::TokenStream::from(attr);
    let input_fn = parse_macro_input!(item as ItemFn);

    match generate_fn_tool(attr_args, &input_fn) {
        Ok(tokens) => tokens.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

fn parse_tool_attr_args(
    tokens: proc_macro2::TokenStream,
) -> syn::Result<(Option<String>, Option<String>)> {
    let mut description = None;
    let mut name_override = None;

    if tokens.is_empty() {
        return Ok((description, name_override));
    }

    let parser = syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated;
    let metas = parser.parse2(tokens)?;

    for meta in &metas {
        if let Meta::NameValue(nv) = meta {
            if nv.path.is_ident("description") {
                if let syn::Expr::Lit(syn::ExprLit {
                    lit: syn::Lit::Str(s),
                    ..
                }) = &nv.value
                {
                    description = Some(s.value());
                }
            } else if nv.path.is_ident("name") {
                if let syn::Expr::Lit(syn::ExprLit {
                    lit: syn::Lit::Str(s),
                    ..
                }) = &nv.value
                {
                    name_override = Some(s.value());
                }
            }
        }
    }

    Ok((description, name_override))
}

fn generate_fn_tool(
    attr_args: proc_macro2::TokenStream,
    input_fn: &ItemFn,
) -> syn::Result<proc_macro2::TokenStream> {
    let (attr_description, attr_name) = parse_tool_attr_args(attr_args)?;

    let fn_name = &input_fn.sig.ident;
    let fn_name_str = fn_name.to_string();

    let struct_name = format_ident!("{}Tool", to_pascal_case(&fn_name_str));
    let tool_name = attr_name.unwrap_or_else(|| fn_name_str.clone());
    let description = attr_description.unwrap_or_else(|| format!("{} tool", tool_name));

    let params = extract_fn_params(&input_fn.sig);

    // Generate properties for JSON schema (using T3's rich schema generation)
    let param_properties: Vec<proc_macro2::TokenStream> = params
        .iter()
        .map(|p| {
            let name = p
                .name_override
                .clone()
                .unwrap_or_else(|| p.ident.to_string());
            let desc = p
                .description
                .clone()
                .unwrap_or_else(|| format!("Parameter: {}", name));
            let effective_ty = unwrap_option_type(&p.ty).unwrap_or(&p.ty);
            let schema_tokens = type_to_json_schema(effective_ty);

            quote! {
                {
                    let mut __prop_schema = #schema_tokens;
                    if let serde_json::Value::Object(ref mut m) = __prop_schema {
                        m.insert("description".to_string(), serde_json::Value::String(#desc.to_string()));
                    }
                    __properties.insert(#name.to_string(), __prop_schema);
                }
            }
        })
        .collect();

    let required_params: Vec<proc_macro2::TokenStream> = params
        .iter()
        .filter(|p| p.required)
        .map(|p| {
            let name = p
                .name_override
                .clone()
                .unwrap_or_else(|| p.ident.to_string());
            quote! { #name }
        })
        .collect();

    // Generate arg extraction code
    let arg_extractions: Vec<proc_macro2::TokenStream> = params
        .iter()
        .map(|p| {
            let ident = &p.ident;
            let name = p
                .name_override
                .clone()
                .unwrap_or_else(|| p.ident.to_string());
            let ty = &p.ty;

            if p.is_option {
                quote! {
                    let #ident: #ty = arguments.get(#name)
                        .and_then(|v| serde_json::from_value(v.clone()).ok());
                }
            } else if p.required {
                quote! {
                    let #ident: #ty = {
                        let val = arguments.get(#name)
                            .ok_or_else(|| cortexai_core::errors::ToolError::InvalidArguments(
                                format!("Missing required parameter: {}", #name)
                            ))?
                            .clone();
                        serde_json::from_value(val)
                            .map_err(|e| cortexai_core::errors::ToolError::InvalidArguments(
                                format!("Invalid type for {}: {}", #name, e)
                            ))?
                    };
                }
            } else {
                quote! {
                    let #ident: #ty = arguments.get(#name)
                        .and_then(|v| serde_json::from_value(v.clone()).ok())
                        .unwrap_or_default();
                }
            }
        })
        .collect();

    let param_idents: Vec<&Ident> = params.iter().map(|p| &p.ident).collect();

    // Strip #[param(...)] attributes from the original function signature
    let mut clean_fn = input_fn.clone();
    for arg in &mut clean_fn.sig.inputs {
        if let FnArg::Typed(pat_type) = arg {
            pat_type.attrs.retain(|a| !a.path().is_ident("param"));
        }
    }

    Ok(quote! {
        #clean_fn

        #[derive(Default)]
        pub struct #struct_name;

        #[async_trait::async_trait]
        impl cortexai_core::tool::Tool for #struct_name {
            fn schema(&self) -> cortexai_core::tool::ToolSchema {
                let mut __properties = serde_json::Map::new();
                #(#param_properties)*

                cortexai_core::tool::ToolSchema {
                    name: #tool_name.to_string(),
                    description: #description.to_string(),
                    parameters: serde_json::json!({
                        "type": "object",
                        "properties": serde_json::Value::Object(__properties),
                        "required": [#(#required_params),*]
                    }),
                    dangerous: false,
                    metadata: std::collections::HashMap::new(),
                    required_scopes: Vec::new(),
                }
            }

            async fn execute(
                &self,
                _context: &cortexai_core::tool::ExecutionContext,
                arguments: serde_json::Value,
            ) -> Result<serde_json::Value, cortexai_core::errors::ToolError> {
                #(#arg_extractions)*

                #fn_name(#(#param_idents),*).await
                    .map_err(|e| cortexai_core::errors::ToolError::ExecutionFailed(e.to_string()))
            }
        }
    })
}