use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{
parse_macro_input, Data, DeriveInput, Expr, Field, Fields, Lit, Meta, Type,
};
use syn::punctuated::Punctuated;
use syn::token::Comma;
#[proc_macro_derive(Tool, attributes(tool))]
pub fn derive_tool(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let fields = match &input.data {
Data::Struct(data) => match &data.fields {
Fields::Named(named) => &named.named,
_ => return err_compile("Tool 只支持具名字段的 struct"),
},
_ => return err_compile("Tool 只能 derive 在 struct 上"),
};
let struct_name = &input.ident;
let (tool_name, tool_desc) = parse_tool_attrs(&input);
let mut field_idents: Vec<TokenStream2> = Vec::new();
let mut schema_entries = Vec::new();
let mut deser_code = Vec::new();
let mut required_fields: Vec<String> = Vec::new();
for field in fields.iter() {
let fname = field.ident.as_ref().unwrap();
let fname_str = fname.to_string();
let desc = field_doc(field).unwrap_or_else(|| format!("字段 {}", fname_str));
let (schema_json, deser, required) = type_info(&field.ty, &fname_str, &desc);
if required {
required_fields.push(fname_str.clone());
}
field_idents.push(quote! { #fname });
schema_entries.push(quote! {
#fname_str: #schema_json
});
deser_code.push(deser);
}
let schema = if required_fields.is_empty() {
quote! {
::serde_json::json!({
"type": "object",
"properties": { #(#schema_entries,)* }
})
}
} else {
quote! {
::serde_json::json!({
"type": "object",
"properties": { #(#schema_entries,)* },
"required": [#(#required_fields,)*]
})
}
};
let expanded = quote! {
#[::async_trait::async_trait]
impl ::cortex_engine::tool::Tool for #struct_name {
fn name(&self) -> &str { #tool_name }
fn description(&self) -> String { #tool_desc.into() }
fn parameters(&self) -> ::serde_json::Value { #schema }
async fn execute(
&self,
__args: ::serde_json::Value,
) -> Result<::serde_json::Value, ::cortex_engine::tool::ToolError> {
#(#deser_code)*
self.__tool_execute_impl(#(#field_idents,)*).await
}
}
};
expanded.into()
}
fn err_compile(msg: &str) -> TokenStream {
syn::Error::new(proc_macro2::Span::call_site(), msg)
.to_compile_error()
.into()
}
fn parse_tool_attrs(input: &DeriveInput) -> (String, String) {
let default_name = camel_to_snake(&input.ident.to_string());
let mut tool_name: Option<String> = None;
let mut tool_desc: Option<String> = None;
for attr in &input.attrs {
if !attr.path().is_ident("tool") { continue; }
if let Ok(metas) = attr.parse_args_with(
Punctuated::<Meta, Comma>::parse_terminated,
) {
for meta in &metas {
if let Meta::NameValue(kv) = meta {
let key = kv.path.get_ident().map(|i| i.to_string());
let value = match &kv.value {
Expr::Lit(lit) => match &lit.lit {
Lit::Str(s) => s.value(),
_ => continue,
},
_ => continue,
};
match key.as_deref() {
Some("name") => tool_name = Some(value),
Some("desc") => tool_desc = Some(value),
_ => {}
}
}
}
}
}
let name = tool_name.unwrap_or(default_name);
let desc = tool_desc.unwrap_or_else(|| format!("工具: {}", input.ident));
(name, desc)
}
fn camel_to_snake(s: &str) -> String {
let mut r = String::with_capacity(s.len() + 4);
for (i, c) in s.chars().enumerate() {
if c.is_uppercase() {
if i > 0 { r.push('_'); }
for l in c.to_lowercase() { r.push(l); }
} else { r.push(c); }
}
r
}
fn field_doc(field: &Field) -> Option<String> {
let docs = get_doc_comments(&field.attrs);
if docs.is_empty() { None } else { Some(docs) }
}
fn get_doc_comments(attrs: &[syn::Attribute]) -> String {
let mut parts = Vec::new();
for attr in attrs {
if !attr.path().is_ident("doc") { continue; }
if let Ok(Meta::NameValue(kv)) = attr.parse_args::<Meta>() {
if let Expr::Lit(el) = &kv.value {
if let Lit::Str(s) = &el.lit {
let val = s.value().trim().to_string();
if !val.is_empty() {
parts.push(val);
}
}
}
}
}
parts.join(" ")
}
fn type_info(ty: &Type, field_name: &str, desc: &str) -> (TokenStream2, TokenStream2, bool) {
let ident = syn::Ident::new(field_name, proc_macro2::Span::call_site());
let key_str = field_name.to_string();
let type_str = quote!(#ty).to_string();
let desc = desc.to_string();
match type_str.as_str() {
"String" | "std :: string :: String" => (
quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
quote! {
let #ident: String = __args.get(#key_str)
.and_then(|v| v.as_str()).unwrap_or("").to_string();
},
true,
),
"i64" => (
quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
quote! {
let #ident: i64 = __args.get(#key_str)
.and_then(|v| v.as_i64()).unwrap_or(0);
},
true,
),
"i32" => (
quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
quote! {
let #ident: i32 = __args.get(#key_str)
.and_then(|v| v.as_i64()).unwrap_or(0) as i32;
},
true,
),
"f64" => (
quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
quote! {
let #ident: f64 = __args.get(#key_str)
.and_then(|v| v.as_f64()).unwrap_or(0.0);
},
true,
),
"f32" => (
quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
quote! {
let #ident: f32 = __args.get(#key_str)
.and_then(|v| v.as_f64()).unwrap_or(0.0) as f32;
},
true,
),
"bool" => (
quote! { ::serde_json::json!({"type": "boolean", "description": #desc}) },
quote! {
let #ident: bool = __args.get(#key_str)
.and_then(|v| v.as_bool()).unwrap_or(false);
},
true,
),
"Option < String >" | "Option < std :: string :: String >" => (
quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
quote! {
let #ident: Option<String> = __args.get(#key_str)
.and_then(|v| v.as_str()).map(String::from);
},
false,
),
"Option < i64 >" => (
quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
quote! {
let #ident: Option<i64> = __args.get(#key_str)
.and_then(|v| v.as_i64());
},
false,
),
"Option < f64 >" => (
quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
quote! {
let #ident: Option<f64> = __args.get(#key_str)
.and_then(|v| v.as_f64());
},
false,
),
"Option < bool >" => (
quote! { ::serde_json::json!({"type": "boolean", "description": #desc}) },
quote! {
let #ident: Option<bool> = __args.get(#key_str)
.and_then(|v| v.as_bool());
},
false,
),
"Vec < String >" | "Vec < std :: string :: String >" => (
quote! { ::serde_json::json!({"type": "array", "items": {"type": "string"}, "description": #desc}) },
quote! {
let #ident: Vec<String> = __args.get(#key_str)
.and_then(|v| v.as_array())
.map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
.unwrap_or_default();
},
true,
),
"Vec < i64 >" => (
quote! { ::serde_json::json!({"type": "array", "items": {"type": "integer"}, "description": #desc}) },
quote! {
let #ident: Vec<i64> = __args.get(#key_str)
.and_then(|v| v.as_array())
.map(|a| a.iter().filter_map(|v| v.as_i64()).collect())
.unwrap_or_default();
},
true,
),
"Vec < f64 >" => (
quote! { ::serde_json::json!({"type": "array", "items": {"type": "number"}, "description": #desc}) },
quote! {
let #ident: Vec<f64> = __args.get(#key_str)
.and_then(|v| v.as_array())
.map(|a| a.iter().filter_map(|v| v.as_f64()).collect())
.unwrap_or_default();
},
true,
),
_ => (
quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
quote! {
let #ident = ::serde_json::from_value::<_>(
__args.get(#key_str).cloned()
.unwrap_or(::serde_json::Value::Null),
).map_err(|e| ::cortex_engine::tool::ToolError::new(
format!("参数 '{}' 格式错误: {}", #key_str, e)
))?;
},
true,
),
}
}