Skip to main content

cortex_derive/
lib.rs

1//! Cortex derive macros — `#[derive(Tool)]`
2
3use proc_macro::TokenStream;
4use proc_macro2::TokenStream as TokenStream2;
5use quote::quote;
6use syn::{
7    parse_macro_input, Data, DeriveInput, Expr, Field, Fields, Lit, Meta, Type,
8};
9use syn::punctuated::Punctuated;
10use syn::token::Comma;
11
12/// 为 struct 生成 `Tool` trait 实现。
13///
14/// 自动生成 `name()`、`description()`、`parameters()` 和 `execute()`。
15/// `execute()` 将参数反序列化后调用 `__tool_execute_impl(...)`。
16#[proc_macro_derive(Tool, attributes(tool))]
17pub fn derive_tool(input: TokenStream) -> TokenStream {
18    let input = parse_macro_input!(input as DeriveInput);
19
20    let fields = match &input.data {
21        Data::Struct(data) => match &data.fields {
22            Fields::Named(named) => &named.named,
23            _ => return err_compile("Tool 只支持具名字段的 struct"),
24        },
25        _ => return err_compile("Tool 只能 derive 在 struct 上"),
26    };
27
28    let struct_name = &input.ident;
29    let (tool_name, tool_desc) = parse_tool_attrs(&input);
30
31    let mut field_idents: Vec<TokenStream2> = Vec::new();
32    let mut schema_entries = Vec::new();
33    let mut deser_code = Vec::new();
34    let mut required_fields: Vec<String> = Vec::new();
35
36    for field in fields.iter() {
37        let fname = field.ident.as_ref().unwrap();
38        let fname_str = fname.to_string();
39        let desc = field_doc(field).unwrap_or_else(|| format!("字段 {}", fname_str));
40        let (schema_json, deser, required) = type_info(&field.ty, &fname_str, &desc);
41
42        if required {
43            required_fields.push(fname_str.clone());
44        }
45        field_idents.push(quote! { #fname });
46        schema_entries.push(quote! {
47            #fname_str: #schema_json
48        });
49        deser_code.push(deser);
50    }
51
52    let schema = if required_fields.is_empty() {
53        quote! {
54            ::serde_json::json!({
55                "type": "object",
56                "properties": { #(#schema_entries,)* }
57            })
58        }
59    } else {
60        quote! {
61            ::serde_json::json!({
62                "type": "object",
63                "properties": { #(#schema_entries,)* },
64                "required": [#(#required_fields,)*]
65            })
66        }
67    };
68
69    let expanded = quote! {
70        #[::async_trait::async_trait]
71        impl ::cortex_engine::tool::Tool for #struct_name {
72            fn name(&self) -> &str { #tool_name }
73            fn description(&self) -> String { #tool_desc.into() }
74            fn parameters(&self) -> ::serde_json::Value { #schema }
75            async fn execute(
76                &self,
77                __args: ::serde_json::Value,
78            ) -> Result<::serde_json::Value, ::cortex_engine::tool::ToolError> {
79                #(#deser_code)*
80                self.__tool_execute_impl(#(#field_idents,)*).await
81            }
82        }
83    };
84
85    expanded.into()
86}
87
88fn err_compile(msg: &str) -> TokenStream {
89    syn::Error::new(proc_macro2::Span::call_site(), msg)
90        .to_compile_error()
91        .into()
92}
93
94/// 解析 #[tool(name = "...", desc = "...")]
95fn parse_tool_attrs(input: &DeriveInput) -> (String, String) {
96    let default_name = camel_to_snake(&input.ident.to_string());
97    let mut tool_name: Option<String> = None;
98    let mut tool_desc: Option<String> = None;
99
100    for attr in &input.attrs {
101        if !attr.path().is_ident("tool") { continue; }
102        if let Ok(metas) = attr.parse_args_with(
103            Punctuated::<Meta, Comma>::parse_terminated,
104        ) {
105            for meta in &metas {
106                if let Meta::NameValue(kv) = meta {
107                    let key = kv.path.get_ident().map(|i| i.to_string());
108                    let value = match &kv.value {
109                        Expr::Lit(lit) => match &lit.lit {
110                            Lit::Str(s) => s.value(),
111                            _ => continue,
112                        },
113                        _ => continue,
114                    };
115                    match key.as_deref() {
116                        Some("name") => tool_name = Some(value),
117                        Some("desc") => tool_desc = Some(value),
118                        _ => {}
119                    }
120                }
121            }
122        }
123    }
124
125    let name = tool_name.unwrap_or(default_name);
126    let desc = tool_desc.unwrap_or_else(|| format!("工具: {}", input.ident));
127    (name, desc)
128}
129
130fn camel_to_snake(s: &str) -> String {
131    let mut r = String::with_capacity(s.len() + 4);
132    for (i, c) in s.chars().enumerate() {
133        if c.is_uppercase() {
134            if i > 0 { r.push('_'); }
135            for l in c.to_lowercase() { r.push(l); }
136        } else { r.push(c); }
137    }
138    r
139}
140
141/// 提取 struct 字段的 doc comment（支持 `/// text` 和 `#[doc = "text"]`）
142fn field_doc(field: &Field) -> Option<String> {
143    let docs = get_doc_comments(&field.attrs);
144    if docs.is_empty() { None } else { Some(docs) }
145}
146
147/// 从属性列表中提取所有 `#[doc = "..."]` 并合并
148fn get_doc_comments(attrs: &[syn::Attribute]) -> String {
149    let mut parts = Vec::new();
150    for attr in attrs {
151        if !attr.path().is_ident("doc") { continue; }
152        // #[doc = "text"] → 解析赋值语句
153        // tokens: `= "text"`
154        if let Ok(Meta::NameValue(kv)) = attr.parse_args::<Meta>() {
155            if let Expr::Lit(el) = &kv.value {
156                if let Lit::Str(s) = &el.lit {
157                    let val = s.value().trim().to_string();
158                    if !val.is_empty() {
159                        parts.push(val);
160                    }
161                }
162            }
163        }
164    }
165    parts.join(" ")
166}
167
168/// 获取类型对应的 JSON Schema（含 desc）+ 反序列化代码 + 是否必填
169fn type_info(ty: &Type, field_name: &str, desc: &str) -> (TokenStream2, TokenStream2, bool) {
170    let ident = syn::Ident::new(field_name, proc_macro2::Span::call_site());
171    let key_str = field_name.to_string();
172    let type_str = quote!(#ty).to_string();
173    let desc = desc.to_string();
174
175    match type_str.as_str() {
176        "String" | "std :: string :: String" => (
177            quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
178            quote! {
179                let #ident: String = __args.get(#key_str)
180                    .and_then(|v| v.as_str()).unwrap_or("").to_string();
181            },
182            true,
183        ),
184        "i64" => (
185            quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
186            quote! {
187                let #ident: i64 = __args.get(#key_str)
188                    .and_then(|v| v.as_i64()).unwrap_or(0);
189            },
190            true,
191        ),
192        "i32" => (
193            quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
194            quote! {
195                let #ident: i32 = __args.get(#key_str)
196                    .and_then(|v| v.as_i64()).unwrap_or(0) as i32;
197            },
198            true,
199        ),
200        "f64" => (
201            quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
202            quote! {
203                let #ident: f64 = __args.get(#key_str)
204                    .and_then(|v| v.as_f64()).unwrap_or(0.0);
205            },
206            true,
207        ),
208        "f32" => (
209            quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
210            quote! {
211                let #ident: f32 = __args.get(#key_str)
212                    .and_then(|v| v.as_f64()).unwrap_or(0.0) as f32;
213            },
214            true,
215        ),
216        "bool" => (
217            quote! { ::serde_json::json!({"type": "boolean", "description": #desc}) },
218            quote! {
219                let #ident: bool = __args.get(#key_str)
220                    .and_then(|v| v.as_bool()).unwrap_or(false);
221            },
222            true,
223        ),
224        // ---------- Option 包装 ----------
225        "Option < String >" | "Option < std :: string :: String >" => (
226            quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
227            quote! {
228                let #ident: Option<String> = __args.get(#key_str)
229                    .and_then(|v| v.as_str()).map(String::from);
230            },
231            false,
232        ),
233        "Option < i64 >" => (
234            quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
235            quote! {
236                let #ident: Option<i64> = __args.get(#key_str)
237                    .and_then(|v| v.as_i64());
238            },
239            false,
240        ),
241        "Option < f64 >" => (
242            quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
243            quote! {
244                let #ident: Option<f64> = __args.get(#key_str)
245                    .and_then(|v| v.as_f64());
246            },
247            false,
248        ),
249        "Option < bool >" => (
250            quote! { ::serde_json::json!({"type": "boolean", "description": #desc}) },
251            quote! {
252                let #ident: Option<bool> = __args.get(#key_str)
253                    .and_then(|v| v.as_bool());
254            },
255            false,
256        ),
257        // ---------- Vec 包装 ----------
258        "Vec < String >" | "Vec < std :: string :: String >" => (
259            quote! { ::serde_json::json!({"type": "array", "items": {"type": "string"}, "description": #desc}) },
260            quote! {
261                let #ident: Vec<String> = __args.get(#key_str)
262                    .and_then(|v| v.as_array())
263                    .map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
264                    .unwrap_or_default();
265            },
266            true,
267        ),
268        "Vec < i64 >" => (
269            quote! { ::serde_json::json!({"type": "array", "items": {"type": "integer"}, "description": #desc}) },
270            quote! {
271                let #ident: Vec<i64> = __args.get(#key_str)
272                    .and_then(|v| v.as_array())
273                    .map(|a| a.iter().filter_map(|v| v.as_i64()).collect())
274                    .unwrap_or_default();
275            },
276            true,
277        ),
278        "Vec < f64 >" => (
279            quote! { ::serde_json::json!({"type": "array", "items": {"type": "number"}, "description": #desc}) },
280            quote! {
281                let #ident: Vec<f64> = __args.get(#key_str)
282                    .and_then(|v| v.as_array())
283                    .map(|a| a.iter().filter_map(|v| v.as_f64()).collect())
284                    .unwrap_or_default();
285            },
286            true,
287        ),
288        // ---------- 通配：serde_json::from_value ----------
289        _ => (
290            quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
291            quote! {
292                let #ident = ::serde_json::from_value::<_>(
293                    __args.get(#key_str).cloned()
294                        .unwrap_or(::serde_json::Value::Null),
295                ).map_err(|e| ::cortex_engine::tool::ToolError::new(
296                    format!("参数 '{}' 格式错误: {}", #key_str, e)
297                ))?;
298            },
299            true,
300        ),
301    }
302}