use std::collections::HashSet;
use proc_macro2::{Span, TokenStream};
use quote::{quote, ToTokens, TokenStreamExt};
use regex::Regex;
use syn::spanned::Spanned;
use syn::{parse_macro_input, GenericArgument, ItemImpl, PathArguments, Signature};
use syn::{FnArg, Ident, Pat, Type};
fn create_tool_schema_const_indentifier(struct_name: &str) -> Ident {
Ident::new(
&format!("_{}_SCHEMA", struct_name.to_uppercase(),),
Span::call_site(),
)
}
struct FunctionDefintion {
is_async: bool,
name: Ident,
name_str: String,
parameters: Vec<Parameter>,
return_type: ReturnType,
description: Option<String>,
}
impl FunctionDefintion {
fn create_schema_const_indentifier(&self, struct_name: &str) -> Ident {
Ident::new(
&format!(
"_{}_{}_PARMETER_SCHEMA",
struct_name.to_uppercase(),
self.name_str.to_uppercase()
),
Span::call_site(),
)
}
}
struct Parameter {
name: Ident,
name_str: String,
param_type: syn::Type,
description: Option<String>,
}
enum ReturnType {
Result(ResultReturnType),
Other(OtherReturnType),
}
struct ResultReturnType {
okay: Type,
error: Type,
}
struct OtherReturnType {
other: Type,
}
#[proc_macro_attribute]
pub fn tool(
_attr: proc_macro::TokenStream,
item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
let mut input = parse_macro_input!(item as ItemImpl);
let struct_name = match &*input.self_ty {
Type::Path(type_path) => &type_path.path.segments.last().unwrap().ident,
_ => panic!("Invalid impl type"),
};
let generics = &input.generics;
let struct_name_str = struct_name.to_token_stream().to_string();
let methods: Vec<_> = input
.items
.clone()
.into_iter()
.filter_map(|item| {
if let syn::ImplItem::Fn(method) = item {
let attrs = &method.attrs;
for attr in attrs.iter() {
let path = attr.path();
if path.is_ident("tool_part") {
return Some(method);
}
}
}
None
})
.collect();
input
.items
.iter_mut()
.for_each(|item| {
if let syn::ImplItem::Fn(method) = item {
method.attrs.retain(|attr|{
!attr.path().is_ident("tool_part")
});
}
});
let mut function_definitions = Vec::new();
for method in methods {
let syn::ImplItemFn {
attrs,
vis: _,
defaultness: _,
sig,
block: _,
} = method;
let mut function_definition = match extract_function_defintion(sig) {
Ok(okay) => okay,
Err(error) => return error.into_compile_error().into(),
};
match extract_description(&mut function_definition, attrs) {
Ok(_) => {}
Err(error) => return error.into_compile_error().into(),
}
function_definitions.push(function_definition);
}
if function_definitions.is_empty() {
return syn::Error::new_spanned(
struct_name,
"No functions found in this tool. Please add functions to the tool with the `#[tool_part]` attribute.",
)
.into_compile_error()
.into();
}
let function_schema = create_tool_json_schema(&struct_name_str, &mut function_definitions);
let parameter_json_schema = function_definitions.iter_mut().map(|function_definition| {
create_function_parameter_json_schema(&struct_name_str, function_definition)
}).fold(TokenStream::new(), |mut acc, item| { acc.append_all(item); acc });
let impl_traits = impl_traits(&struct_name, &struct_name_str, generics, &function_definitions);
let expanded = quote! {
#input
#function_schema
#parameter_json_schema
#impl_traits
};
proc_macro::TokenStream::from(expanded)
}
struct CommonReturnTypes<'a> {
result_err: HashSet<&'a Type>,
result_ok_and_regular: HashSet<&'a Type>,
}
impl<'a> CommonReturnTypes<'a> {
pub fn new() -> Self {
Self {
result_err: HashSet::new(),
result_ok_and_regular: HashSet::new(),
}
}
}
fn impl_traits(struct_name: &syn::Ident, struct_name_str: &str, generics: &syn::Generics, function_definitions: &Vec<FunctionDefintion>) -> TokenStream {
let mut common_return_types = CommonReturnTypes::new();
for function_definition in function_definitions.iter() {
match &function_definition.return_type {
ReturnType::Result(result_return_type) => {
common_return_types.result_err.insert(&result_return_type.error);
common_return_types.result_ok_and_regular.insert(&result_return_type.okay);
}
ReturnType::Other(other_return_type) => {
common_return_types.result_ok_and_regular.insert(&other_return_type.other);
}
}
}
let mut common_err_type: Option<Type> = None;
let all_are_results_with_same_err_type = common_return_types.result_err.len() == 1;
if all_are_results_with_same_err_type {
let first = *common_return_types.result_err.iter().next().unwrap();
common_err_type = Some(first.clone());
}
let mut common_ok_type: Option<Type> = None;
let all_have_same_ok_type = common_return_types.result_ok_and_regular.len() == 1;
if all_have_same_ok_type {
let first = *common_return_types.result_ok_and_regular.iter().next().unwrap();
common_ok_type = Some(first.clone());
}
let all_functions_are_regular = common_return_types.result_err.len() == 0; let impls_needed = determine_impls_needed(common_ok_type, common_err_type, all_functions_are_regular);
let mut all_impl_tokens = TokenStream::new();
let box_any_type = quote! {
Box<dyn std::any::Any>
};
let box_error_type = quote! {
Box<dyn std::error::Error>
};
let infallible_type = quote! {
std::convert::Infallible
};
for impl_needed in impls_needed {
let tokens = match impl_needed {
ImplTypes::BoxAndBox => impl_trait(struct_name, struct_name_str, generics,function_definitions, true, true, &box_any_type, &box_error_type),
ImplTypes::BoxAndSpecific(err_type) => impl_trait(struct_name, struct_name_str, generics,function_definitions, true, false, &box_any_type, &err_type.to_token_stream()),
ImplTypes::SpecificAndBox(ok_type) => impl_trait(struct_name, struct_name_str, generics,function_definitions, false, true, &ok_type.to_token_stream(), &box_error_type),
ImplTypes::SpecificAndSpecific(ok_type, err_type) => impl_trait(struct_name, struct_name_str, generics,function_definitions, false, false, &ok_type.to_token_stream(), &err_type.to_token_stream()),
ImplTypes::BoxAndInfallible => impl_trait(struct_name, struct_name_str, generics,function_definitions, true, false, &box_any_type, &infallible_type),
ImplTypes::SpecificAndInfallible(ok_type) => impl_trait(struct_name, struct_name_str, generics,function_definitions, false, false, &ok_type.to_token_stream(), &infallible_type),
};
all_impl_tokens.append_all(tokens);
}
all_impl_tokens
}
enum ImplTypes {
BoxAndBox,
BoxAndSpecific(Type),
SpecificAndBox(Type),
SpecificAndSpecific(Type, Type),
BoxAndInfallible,
SpecificAndInfallible(Type),
}
fn determine_impls_needed(common_ok_type: Option<Type>, common_err_type: Option<Type>, all_functions_are_regular: bool) -> Vec<ImplTypes> {
let mut vecs = match (common_ok_type.clone(), common_err_type.clone()) {
(None, None) => vec![],
(None, Some(err_type)) => vec![ImplTypes::BoxAndSpecific(err_type)],
(Some(ok_type), None) => vec![ImplTypes::SpecificAndBox(ok_type)],
(Some(ok_type), Some(err_type)) => vec![ImplTypes::BoxAndSpecific(err_type.clone()), ImplTypes::SpecificAndBox(ok_type.clone()), ImplTypes::SpecificAndSpecific(ok_type, err_type)],
};
if all_functions_are_regular {
assert!(common_err_type.is_none(), "If there are no result functions, there should be no error type");
vecs.push(ImplTypes::BoxAndInfallible);
if let Some(common_ok_type) = common_ok_type {
vecs.push(ImplTypes::SpecificAndInfallible(common_ok_type));
}
}
vecs.push(ImplTypes::BoxAndBox);
vecs
}
fn impl_trait(struct_name: &syn::Ident, struct_name_str:&str, generics: &syn::Generics, function_definitions: &Vec<FunctionDefintion>, ok_needs_box: bool, err_needs_box: bool, ok_type: &TokenStream, err_type: &TokenStream) -> TokenStream {
let function_names = function_definitions.iter().map(|function_definition| {
&function_definition.name_str
});
let run_arms = function_definitions.iter().map(|function_definition| {
let function_parameter_statements = function_definition.parameters.iter().map(|parameter|{
let Parameter {
name,
name_str,
param_type,
description: _,
} = parameter;
let serde_message = format!("Parameter `{}` does not follow schema", name_str);
let missing_message = format!("Missing `{}` parameter", name_str);
let deserialize= match param_type {
Type::Reference(type_reference) => match &*type_reference.elem {
Type::Path(type_path) => {
if type_path.path.get_ident().is_some_and(|item| &*item.to_string() == "str") {
Some(quote! {
let #name: &str = &*serde_json::from_value::<String>(#name).map_err(|_| llmtoolbox::FunctionCallError::parsing(#serde_message.to_owned()))?;
})
}
else {
Some(quote! {
let #name: #param_type = &serde_json::from_value::<#type_path>(#name).map_err(|_| llmtoolbox::FunctionCallError::parsing(#serde_message.to_owned()))?;
})
}
},
_ => None,
},
_ => None,
}.unwrap_or(quote! {
let #name: #param_type = serde_json::from_value::<#param_type>(#name).map_err(|_| llmtoolbox::FunctionCallError::parsing(#serde_message.to_owned()))?;
});
quote! {
let #name = parameters.remove(#name_str).ok_or_else(|| llmtoolbox::FunctionCallError::parsing(#missing_message.to_owned()))?;
#deserialize
}
});
let return_statement =
make_return_statement(function_definition, ok_needs_box, err_needs_box);
let function_name_str = &function_definition.name_str;
quote! {
#function_name_str => {
#(#function_parameter_statements)*
#return_statement
}
}
}).fold(TokenStream::new(), |mut acc, item| { acc.append_all(item); acc });
let schema = create_tool_schema_const_indentifier(struct_name_str);
let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
quote! {
impl #impl_generics llmtoolbox::Tool<#ok_type, #err_type> for #struct_name #ty_generics #where_clause {
fn function_names(&self) -> &[&'static str] {
&[
#(#function_names),*
]
}
fn schema(&self) -> &'static serde_json::Map<String, serde_json::Value> {
#schema.as_object().unwrap()
}
fn call_function<'life0, 'life1, 'async_trait>(
&'life0 self,
name: &'life1 str,
parameters: serde_json::Map<String, serde_json::Value>,
) -> ::core::pin::Pin<
Box<
dyn ::core::future::Future<
Output = Result<
Result<#ok_type, #err_type>,
llmtoolbox::FunctionCallError,
>,
> + ::core::marker::Send
+ 'async_trait,
>,
>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
Box::pin(async move {
if let ::core::option::Option::Some(__ret) = ::core::option::Option::None::<
Result<
Result<#ok_type, #err_type>,
llmtoolbox::FunctionCallError,
>,
> {
#[allow(unreachable_code)]
return __ret;
}
let __self = self;
let mut parameters = parameters;
let __ret: Result<
Result<#ok_type, #err_type>,
llmtoolbox::FunctionCallError,
> = {
match &*name {
#run_arms
_ => return Err(llmtoolbox::FunctionCallError::function_not_found(name.to_owned())),
}
};
#[allow(unreachable_code)]
__ret
})
}
}
}
}
fn make_return_statement(function_definition: &FunctionDefintion, ok_needs_box: bool, err_needs_box: bool) -> TokenStream {
let async_part;
if function_definition.is_async {
async_part = quote! {
.await
}
}
else {
async_part = quote! {}
}
let function_parameters = function_definition.parameters.iter().map(|parameter| {
¶meter.name
});
let function_name = &function_definition.name;
match function_definition.return_type {
ReturnType::Result(_) => {
if ok_needs_box {
if err_needs_box {
quote! {
return Ok(match self.#function_name(#(#function_parameters),*)#async_part {
Ok(value) => Ok(Box::new(value) as Box<dyn std::any::Any>),
Err(value) => Err(Box::new(value) as Box<dyn std::error::Error>),
});
}
}
else {
quote! {
return Ok(self.#function_name(#(#function_parameters),*)#async_part.map(|value| Box::new(value) as Box<dyn std::any::Any>));
}
}
}
else {
if err_needs_box {
quote! {
return Ok(self.#function_name(#(#function_parameters),*)#async_part.map_err(|error| Box::new(error) as Box<dyn std::error::Error>));
}
}
else {
quote! {
return Ok(self.#function_name(#(#function_parameters),*)#async_part);
}
}
}
},
ReturnType::Other(_) => {
if ok_needs_box {
quote! {
return Ok(Ok(Box::new(self.#function_name(#(#function_parameters),*)#async_part)));
}
}
else {
quote! {
return Ok(Ok(self.#function_name(#(#function_parameters),*)#async_part));
}
}
}
}
}
fn extract_function_defintion(signature: Signature) -> syn::Result<FunctionDefintion> {
let inputs = &signature.inputs;
let parameters = inputs
.iter()
.filter_map(|arg| {
if let FnArg::Typed(arg) = arg {
if let Pat::Ident(pat_ident) = &*arg.pat {
let name_str = pat_ident.ident.to_string();
let name = pat_ident.ident.clone();
let type_ = *arg.ty.clone();
Some(Parameter {
name,
name_str,
param_type: type_,
description: None,
})
} else {
None
}
} else {
None
}
})
.collect::<Vec<_>>();
let return_type = match signature.output {
syn::ReturnType::Default => {
return Err(syn::Error::new_spanned(
signature,
"Currently, tool functions must have a return type, even if it is just `()`.",
))
}
syn::ReturnType::Type(_, return_type) => *return_type,
};
let return_type = (|| {
match &return_type {
Type::Path(type_path) => {
let segments = &type_path.path.segments;
if segments.len() != 1 {
return ReturnType::Other(OtherReturnType { other: return_type });
}
let segment = segments.last().unwrap();
if let PathArguments::AngleBracketed(angle_bracketed_args) = &segment.arguments {
let mut generics = angle_bracketed_args.args.iter();
if let (Some(GenericArgument::Type(okay)), Some(GenericArgument::Type(error))) =
(generics.next(), generics.next())
{
return ReturnType::Result(ResultReturnType {
okay: okay.clone(),
error: error.clone(),
});
}
}
}
_ => {}
}
return ReturnType::Other(OtherReturnType { other: return_type });
})();
let is_async = signature.asyncness.is_some();
let name = signature.ident;
let name_str = name.to_string();
Ok(FunctionDefintion {
is_async,
name,
name_str,
parameters,
return_type,
description: None,
})
}
fn extract_description(
function_definition: &mut FunctionDefintion,
attrs: Vec<syn::Attribute>,
) -> syn::Result<()> {
let FunctionDefintion {
is_async: _,
name,
name_str,
parameters,
return_type: _,
description,
} = function_definition;
let re = Regex::new(r".*?`(?<name>.*?)`\s*-\s*(?<description>.*)$").unwrap();
for attr in attrs.iter() {
match &attr.meta {
syn::Meta::NameValue(name_value) => match &name_value.value {
syn::Expr::Lit(lit) => match &lit.lit {
syn::Lit::Str(str) => {
let haystack = str.value();
let arg_caps = match re.captures(&haystack) {
Some(caps) => caps,
None => {
if let Some(description) = description {
description.push_str(&*format!("{}\n", &str.value().trim()));
} else {
let _ = description.insert(str.value().trim().to_string());
}
continue;
}
};
let name = arg_caps["name"].to_string();
let desc = arg_caps["description"].to_string();
if let Some(param) = parameters.iter_mut().find(|p| p.name_str == name) {
param.description = Some(desc);
} else {
return Err(syn::Error::new_spanned(
attr,
format!("parameter `{}` not found in function definition", name),
));
}
}
_ => {}
},
_ => {}
},
_ => {}
}
}
for parameter in parameters {
if parameter.description.is_none() {
return Err(syn::Error::new_spanned(
parameter.name.clone(),
format!("missing description for parameter `{}`. Descriptions are doc comments the form of:\n\
/// `parameter_name` - This is the description for the parameter.", parameter.name_str),
));
}
}
if function_definition.description.is_none() {
return Err(syn::Error::new_spanned(
name.clone(),
format!("missing description for function `{}`", name_str),
));
}
Ok(())
}
fn rust_type_to_known_json_schema_type(ty: &Type) -> Option<&'static str> {
match ty {
Type::Path(type_path) => {
if let Some(segment) = type_path.path.segments.last() {
return match segment.ident.to_string().as_str() {
"String" | "str" => Some("string"),
"i8" | "i16" | "i32" | "i64" | "isize" => Some("integer"),
"u8" | "u16" | "u32" | "u64" | "usize" => Some("integer"), "u128" | "i128" => Some("integer"), "f32" | "f64" => Some("number"),
"bool" => Some("boolean"),
_ => None,
};
} else {
None
}
}
Type::Reference(type_ref) => rust_type_to_known_json_schema_type(&type_ref.elem),
_ => None,
}
}
fn create_tool_json_schema(
struct_name: &str,
function_definitions: &Vec<FunctionDefintion>,
) -> proc_macro2::TokenStream {
let mut function_schemas = Vec::new();
for function_definition in function_definitions {
let id = function_definition.create_schema_const_indentifier(struct_name);
let description = &function_definition.description;
let name = &function_definition.name;
function_schemas.push(quote! {
serde_json::json!(
{
"type": "object",
"description": stringify!(#description),
"properties": {
"function_name": {
"const": stringify!(#name),
},
"parameters": *#id
},
"required": ["function_name", "parameters"]
}
)
});
}
let id = create_tool_schema_const_indentifier(struct_name);
quote! {
const #id: std::cell::LazyCell<&'static serde_json::Value> = std::cell::LazyCell::new(|| {
Box::leak(Box::new(serde_json::json!(
{
"$schema": "http://json-schema.org/draft-07/schema#",
"oneOf": [
#(#function_schemas),*
]
}
)))
});
}
}
fn create_function_parameter_json_schema(
struct_name: &str,
function_definition: &mut FunctionDefintion,
) -> proc_macro2::TokenStream {
let parameters = &function_definition.parameters;
let mut known_properties = Vec::new();
let mut known_required_property_name = Vec::new();
let mut computed_required_property_name = Vec::new();
let mut computed_properties_outer_definitions = Vec::new();
let mut computed_properties = Vec::new();
let mut num_of_computed_properties = 0;
for parameter in parameters {
let name = ¶meter.name_str;
let description = ¶meter.description;
let param_type = ¶meter.param_type;
let json_schema_type = rust_type_to_known_json_schema_type(¶meter.param_type);
if let Some(param_type) = json_schema_type {
known_properties.push(quote! {
#name: {
"type": #param_type,
"description": #description
}
});
known_required_property_name.push(quote! {
#name
});
} else {
num_of_computed_properties +=1;
let id = Ident::new(
&format!("computed{num_of_computed_properties}"),
json_schema_type.span(),
);
computed_properties_outer_definitions.push(quote! {
let #id = (|| {
let schema_settings = schemars::generate::SchemaSettings::draft07();
let schema = schemars::SchemaGenerator::new(schema_settings).into_root_schema_for::<#param_type>();
let mut schema = schema.to_value();
llmtoolbox::clean_up_schema(&mut schema);
match schema {
serde_json::Value::Object(ref mut map) => {
map.insert("description".to_string(), serde_json::Value::String(#description.to_string()));
},
_ => panic!("schema should always generate a map type.")
}
return schema;
})();
});
computed_properties.push(quote! {
#name: #id
});
computed_required_property_name.push(quote! {
#name
});
}
}
let id = function_definition.create_schema_const_indentifier(struct_name);
quote! {
const #id: std::cell::LazyCell<serde_json::Value> = std::cell::LazyCell::new(|| {
#(#computed_properties_outer_definitions)*
serde_json::json!(
{
"type": "object",
"required": [
#(#known_required_property_name),*
#(#computed_required_property_name),*
],
"properties": {
#(#known_properties),*
#(#computed_properties),*
},
}
)
});
}
}