mod prompt_macro;
mod resource_macro;
#[cfg(feature = "md-tmpl")]
mod response_struct_gen;
use convert_case::{Case, Casing};
use proc_macro::TokenStream;
use quote::{format_ident, quote};
#[cfg(feature = "md-tmpl")]
use syn::Ident;
use syn::{ItemFn, LitStr, parse_macro_input};
#[proc_macro_attribute]
pub fn llm_tool(attr: TokenStream, item: TokenStream) -> TokenStream {
let func = parse_macro_input!(item as ItemFn);
let tool_attr = if attr.is_empty() {
None
} else {
match syn::parse::<ToolAttr>(attr) {
Ok(parsed) => Some(parsed),
Err(err) => return err.to_compile_error().into(),
}
};
match tool_impl(&func, tool_attr.as_ref()) {
Ok(tokens) => tokens.into(),
Err(err) => err.to_compile_error().into(),
}
}
#[proc_macro_attribute]
pub fn llm_prompt(attr: TokenStream, item: TokenStream) -> TokenStream {
let func = parse_macro_input!(item as ItemFn);
let tool_attr = if attr.is_empty() {
None
} else {
match syn::parse::<ToolAttr>(attr) {
Ok(parsed) => Some(parsed),
Err(err) => return err.to_compile_error().into(),
}
};
match prompt_macro::prompt_impl(&func, tool_attr.as_ref()) {
Ok(tokens) => tokens.into(),
Err(err) => err.to_compile_error().into(),
}
}
#[proc_macro_attribute]
pub fn llm_resource(attr: TokenStream, item: TokenStream) -> TokenStream {
let func = parse_macro_input!(item as ItemFn);
let res_attr = match syn::parse::<resource_macro::ResourceAttr>(attr) {
Ok(parsed) => parsed,
Err(err) => return err.to_compile_error().into(),
};
match resource_macro::resource_impl(&func, &res_attr) {
Ok(tokens) => tokens.into(),
Err(err) => err.to_compile_error().into(),
}
}
struct ToolAttr {
description_inline: Option<LitStr>,
description_file_path: Option<LitStr>,
response_file_path: Option<LitStr>,
response_inline: Option<LitStr>,
#[cfg(feature = "md-tmpl")]
inline_params: Vec<(Ident, LitStr)>,
#[cfg(feature = "md-tmpl")]
env_vars: Vec<(Ident, syn::Lit)>,
#[cfg(feature = "md-tmpl")]
context_fn: Option<syn::Path>,
has_inline_params: bool,
has_context_fn: bool,
}
pub(crate) const MACRO_LLM_TOOL: &str = "llm_tool";
pub(crate) const MACRO_LLM_PROMPT: &str = "llm_prompt";
pub(crate) const MACRO_LLM_RESOURCE: &str = "llm_resource";
pub(crate) const ATTR_DESCRIPTION: &str = "description";
pub(crate) const ATTR_DESCRIPTION_FILE: &str = "description_file";
pub(crate) const ATTR_RESPONSE_FILE: &str = "response_file";
pub(crate) const ATTR_RESPONSE: &str = "response";
pub(crate) const ATTR_PARAMS: &str = "params";
pub(crate) const ATTR_CONTEXT: &str = "context";
pub(crate) const ATTR_ENV: &str = "env";
pub(crate) const ATTR_DOC: &str = "doc";
pub(crate) const TYPE_OPTION: &str = "Option";
pub(crate) const TYPE_TOOL_CONTEXT: &str = "ToolContext";
pub(crate) const TYPE_STR: &str = "str";
pub(crate) const TYPE_RESULT: &str = "Result";
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub(crate) enum ToolAttrKey {
Description,
DescriptionFile,
ResponseFile,
Response,
Params,
Env,
Context,
}
impl ToolAttrKey {
pub(crate) const ALL: &'static [Self] = &[
Self::Description,
Self::DescriptionFile,
Self::Response,
Self::ResponseFile,
Self::Params,
Self::Env,
Self::Context,
];
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::Description => ATTR_DESCRIPTION,
Self::DescriptionFile => ATTR_DESCRIPTION_FILE,
Self::ResponseFile => ATTR_RESPONSE_FILE,
Self::Response => ATTR_RESPONSE,
Self::Params => ATTR_PARAMS,
Self::Env => ATTR_ENV,
Self::Context => ATTR_CONTEXT,
}
}
pub(crate) fn expected_keys_error(span: proc_macro2::Span) -> syn::Error {
let mut parts: Vec<String> = Self::ALL
.iter()
.map(|k| format!("`{}`", k.as_str()))
.collect();
let last = parts.pop().unwrap_or_default();
let formatted = if parts.is_empty() {
last
} else {
format!("{}, or {last}", parts.join(", "))
};
syn::Error::new(span, format!("expected {formatted}"))
}
}
impl TryFrom<&syn::Ident> for ToolAttrKey {
type Error = syn::Error;
fn try_from(ident: &syn::Ident) -> Result<Self, Self::Error> {
let s = ident.to_string();
for &variant in Self::ALL {
if s == variant.as_str() {
return Ok(variant);
}
}
Err(Self::expected_keys_error(ident.span()))
}
}
#[derive(Default)]
struct ToolAttrBuilder {
description_inline: Option<syn::LitStr>,
description_file_path: Option<syn::LitStr>,
response_file_path: Option<syn::LitStr>,
response_inline: Option<syn::LitStr>,
#[cfg(feature = "md-tmpl")]
inline_params: Vec<(syn::Ident, syn::LitStr)>,
#[cfg(feature = "md-tmpl")]
env_vars: Vec<(syn::Ident, syn::Lit)>,
#[cfg(feature = "md-tmpl")]
context_fn: Option<syn::Path>,
#[cfg(not(feature = "md-tmpl"))]
has_inline_params: bool,
#[cfg(not(feature = "md-tmpl"))]
has_context_fn: bool,
#[cfg(not(feature = "md-tmpl"))]
has_env: bool,
}
impl ToolAttrBuilder {
fn parse_params_attr(&mut self, input: syn::parse::ParseStream) -> syn::Result<()> {
let content;
syn::parenthesized!(content in input);
while !content.is_empty() {
let key: syn::Ident = content.parse()?;
let _: syn::Token![=] = content.parse()?;
let value: syn::LitStr = content.parse()?;
#[cfg(feature = "md-tmpl")]
self.inline_params.push((key, value));
#[cfg(not(feature = "md-tmpl"))]
{
drop(key);
drop(value);
}
if !content.is_empty() {
let _: syn::Token![,] = content.parse()?;
}
}
#[cfg(not(feature = "md-tmpl"))]
{
self.has_inline_params = true;
}
Ok(())
}
fn parse_env_attr(&mut self, input: syn::parse::ParseStream) -> syn::Result<()> {
let content;
syn::parenthesized!(content in input);
while !content.is_empty() {
let key: syn::Ident = content.parse()?;
let _: syn::Token![=] = content.parse()?;
let value: syn::Lit = content.parse()?;
match &value {
syn::Lit::Str(_) | syn::Lit::Int(_) | syn::Lit::Float(_) | syn::Lit::Bool(_) => {}
other => {
return Err(syn::Error::new(
other.span(),
"env values must be string, integer, float, or bool literals",
));
}
}
#[cfg(feature = "md-tmpl")]
self.env_vars.push((key, value));
#[cfg(not(feature = "md-tmpl"))]
{
drop(key);
drop(value);
}
if !content.is_empty() {
let _: syn::Token![,] = content.parse()?;
}
}
#[cfg(not(feature = "md-tmpl"))]
{
self.has_env = true;
}
Ok(())
}
fn parse_single(&mut self, input: syn::parse::ParseStream) -> syn::Result<()> {
let ident: syn::Ident = input.parse()?;
let key = ToolAttrKey::try_from(&ident)?;
match key {
ToolAttrKey::Description => {
let _: syn::Token![=] = input.parse()?;
if self.description_inline.is_some() {
return Err(syn::Error::new(
ident.span(),
format!("duplicate `{}` attribute", key.as_str()),
));
}
self.description_inline = Some(input.parse::<syn::LitStr>()?);
}
ToolAttrKey::DescriptionFile => {
let _: syn::Token![=] = input.parse()?;
if self.description_file_path.is_some() {
return Err(syn::Error::new(
ident.span(),
format!("duplicate `{}` attribute", key.as_str()),
));
}
self.description_file_path = Some(input.parse::<syn::LitStr>()?);
}
ToolAttrKey::ResponseFile => {
let _: syn::Token![=] = input.parse()?;
if self.response_file_path.is_some() {
return Err(syn::Error::new(
ident.span(),
format!("duplicate `{}` attribute", key.as_str()),
));
}
self.response_file_path = Some(input.parse::<syn::LitStr>()?);
}
ToolAttrKey::Response => {
let _: syn::Token![=] = input.parse()?;
if self.response_inline.is_some() {
return Err(syn::Error::new(
ident.span(),
format!("duplicate `{}` attribute", key.as_str()),
));
}
self.response_inline = Some(input.parse::<syn::LitStr>()?);
}
ToolAttrKey::Params => {
self.parse_params_attr(input)?;
}
ToolAttrKey::Env => {
self.parse_env_attr(input)?;
}
ToolAttrKey::Context => {
let _: syn::Token![=] = input.parse()?;
#[cfg(feature = "md-tmpl")]
{
if self.context_fn.is_some() {
return Err(syn::Error::new(
ident.span(),
format!("duplicate `{}` attribute", key.as_str()),
));
}
self.context_fn = Some(input.parse::<syn::Path>()?);
}
#[cfg(not(feature = "md-tmpl"))]
{
let _path: syn::Path = input.parse()?;
if self.has_context_fn {
return Err(syn::Error::new(
ident.span(),
format!("duplicate `{}` attribute", key.as_str()),
));
}
self.has_context_fn = true;
}
}
}
Ok(())
}
}
impl syn::parse::Parse for ToolAttr {
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
let mut builder = ToolAttrBuilder::default();
while !input.is_empty() {
builder.parse_single(input)?;
if !input.is_empty() {
let _: syn::Token![,] = input.parse()?;
}
}
#[cfg(feature = "md-tmpl")]
let has_inline_params = !builder.inline_params.is_empty();
#[cfg(not(feature = "md-tmpl"))]
let has_inline_params = builder.has_inline_params;
#[cfg(feature = "md-tmpl")]
let has_context_fn = builder.context_fn.is_some();
#[cfg(not(feature = "md-tmpl"))]
let has_context_fn = builder.has_context_fn;
validate_tool_attr(&builder)?;
Ok(Self {
description_inline: builder.description_inline,
description_file_path: builder.description_file_path,
response_file_path: builder.response_file_path,
response_inline: builder.response_inline,
#[cfg(feature = "md-tmpl")]
inline_params: builder.inline_params,
#[cfg(feature = "md-tmpl")]
env_vars: builder.env_vars,
#[cfg(feature = "md-tmpl")]
context_fn: builder.context_fn,
has_inline_params,
has_context_fn,
})
}
}
fn validate_tool_attr(builder: &ToolAttrBuilder) -> syn::Result<()> {
if builder.description_inline.is_some() && builder.description_file_path.is_some() {
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`description` and `description_file` are mutually exclusive",
));
}
if builder.response_file_path.is_some() && builder.response_inline.is_some() {
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`response` and `response_file` are mutually exclusive",
));
}
#[cfg(feature = "md-tmpl")]
let has_inline_params = !builder.inline_params.is_empty();
#[cfg(not(feature = "md-tmpl"))]
let has_inline_params = builder.has_inline_params;
#[cfg(feature = "md-tmpl")]
let has_context_fn = builder.context_fn.is_some();
#[cfg(not(feature = "md-tmpl"))]
let has_context_fn = builder.has_context_fn;
#[cfg(feature = "md-tmpl")]
let has_env = !builder.env_vars.is_empty();
#[cfg(not(feature = "md-tmpl"))]
let has_env = builder.has_env;
if has_inline_params && has_context_fn {
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`params(...)` and `context = ...` are mutually exclusive; \
use `params` for compile-time values or `context` for runtime values",
));
}
if has_inline_params
&& builder.description_file_path.is_none()
&& builder.description_inline.is_none()
{
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`params(...)` requires `description_file = \"...\"` or `description = \"...\"`",
));
}
if has_context_fn
&& builder.description_file_path.is_none()
&& builder.description_inline.is_none()
{
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`context = ...` requires `description_file = \"...\"` or `description = \"...\"`",
));
}
if has_env && builder.description_file_path.is_none() && builder.description_inline.is_none() {
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`env(...)` requires `description_file = \"...\"` or `description = \"...\"`",
));
}
#[cfg(not(feature = "md-tmpl"))]
if builder.description_file_path.is_some() {
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`description_file` requires the `md-tmpl` feature of `llm-tool`",
));
}
#[cfg(not(feature = "md-tmpl"))]
if builder.response_file_path.is_some() {
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`response_file` requires the `md-tmpl` feature of `llm-tool`",
));
}
#[cfg(not(feature = "md-tmpl"))]
if builder.response_inline.is_some() {
return Err(syn::Error::new(
proc_macro2::Span::call_site(),
"`response` requires the `md-tmpl` feature of `llm-tool`",
));
}
Ok(())
}
struct ParamInfo {
name: syn::Ident,
ty: Box<syn::Type>,
doc_attrs: Vec<syn::Attribute>,
is_context: bool,
is_mut: bool,
}
enum ReturnInfo {
ResultType {
ok_type: Box<syn::Type>,
err_type: Box<syn::Type>,
},
BareType,
}
fn tool_impl(func: &ItemFn, attr: Option<&ToolAttr>) -> syn::Result<proc_macro2::TokenStream> {
let crate_path = quote! { ::llm_tool };
let fn_name = &func.sig.ident;
reject_generic_signature(func, MACRO_LLM_TOOL)?;
let tool_name_str = fn_name.to_string();
let struct_name = format_ident!("{}", tool_name_str.to_case(Case::Pascal));
let params_name = format_ident!("{}Params", struct_name);
let DescriptionInfo {
static_description,
helper_tokens,
description_method,
dep_tracking,
} = resolve_description(func, attr)?;
let response_info = resolve_response_template(attr, &struct_name, fn_name)?;
let all_params = extract_params(func, MACRO_LLM_TOOL)?;
let ctx_count = all_params.iter().filter(|p| p.is_context).count();
if ctx_count > 1 {
return Err(syn::Error::new_spanned(
&func.sig,
"#[llm_tool] functions can accept at most one ToolContext parameter",
));
}
let ctx_param = all_params.iter().find(|p| p.is_context);
let params: Vec<&ParamInfo> = all_params.iter().filter(|p| !p.is_context).collect();
for param in ¶ms {
if param.doc_attrs.is_empty() {
return Err(syn::Error::new_spanned(
¶m.name,
format!(
"#[llm_tool] parameter `{}` must have a doc comment \
(used as the parameter description in the JSON schema)",
param.name
),
));
}
}
let return_info = parse_return_type(func, MACRO_LLM_TOOL)?;
let param_names: Vec<_> = params.iter().map(|p| &p.name).collect();
let param_descriptions: Vec<String> = params
.iter()
.map(|p| extract_doc_string(&p.doc_attrs))
.collect();
let (param_struct_types, borrow_bindings) = build_param_types_and_borrows(¶ms);
let serde_defaults = build_serde_defaults(¶ms);
let body_tokens = build_body_tokens(func, &return_info, &crate_path, &response_info);
let vis = &func.vis;
let params_doc = format!("Auto-generated parameters for the [`{struct_name}`] tool.");
let struct_doc = format!(
"Auto-generated tool struct. See the `#[llm_tool]`-annotated function `{fn_name}` for the implementation."
);
let ctx_binding = if let Some(cp) = ctx_param {
let ctx_name = &cp.name;
quote! { let #ctx_name = _ctx; }
} else {
quote! {}
};
let mut_tokens: Vec<proc_macro2::TokenStream> = params
.iter()
.map(|p| {
if p.is_mut {
quote! { mut }
} else {
quote! {}
}
})
.collect();
let response_dep_tracking = &response_info.dep_tracking;
let response_helper_tokens = &response_info.helper_tokens;
Ok(quote! {
#dep_tracking
#response_dep_tracking
#helper_tokens
#response_helper_tokens
#[doc = #params_doc]
#[derive(::serde::Deserialize, ::schemars::JsonSchema)]
#vis struct #params_name {
#(
#[schemars(description = #param_descriptions)]
#serde_defaults
pub #param_names: #param_struct_types,
)*
}
#[doc = #struct_doc]
#vis struct #struct_name;
impl #crate_path::RustTool for #struct_name {
type Params = #params_name;
const NAME: &'static str = #tool_name_str;
const DESCRIPTION: &'static str = #static_description;
#description_method
#[allow(unknown_lints, clippy::unused_async_trait_impl)]
async fn call(&self, params: Self::Params, _ctx: &#crate_path::ToolContext) -> ::core::result::Result<#crate_path::ToolOutput, #crate_path::ToolError> {
use #crate_path::__private::SerializeFallback as _;
let #params_name { #( #mut_tokens #param_names, )* } = params;
#( #borrow_bindings )*
#ctx_binding
#body_tokens
}
}
})
}
struct DescriptionInfo {
static_description: String,
helper_tokens: proc_macro2::TokenStream,
description_method: Option<proc_macro2::TokenStream>,
dep_tracking: proc_macro2::TokenStream,
}
pub(crate) mod desc;
pub(crate) mod helpers;
pub(crate) use desc::resolve_description;
pub(crate) use helpers::{
build_body_tokens, build_param_types_and_borrows, build_serde_defaults, extract_doc_string,
extract_params, parse_return_type, reject_generic_signature, resolve_response_template,
};