#[cfg(feature = "derive")]
pub use aither_derive::tool;
use alloc::borrow::Cow;
use serde_json::Value;
use crate::Result;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use alloc::{boxed::Box, collections::BTreeMap};
use core::any::Any;
use core::fmt::{Debug, Display};
use core::{future::Future, pin::Pin};
pub use mime::Mime;
use schemars::{JsonSchema, Schema, schema_for};
use serde::{Serialize, de::DeserializeOwned};
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "serde", serde(tag = "kind", rename_all = "snake_case"))]
pub enum ToolResult {
Done,
Text {
text: String,
},
Tsv {
text: String,
},
Json {
value: Value,
},
Binary {
mime: String,
content: Vec<u8>,
},
Error {
message: String,
},
}
impl ToolResult {
#[must_use]
pub fn text(s: impl Into<String>) -> Self {
Self::Text { text: s.into() }
}
#[must_use]
pub fn tsv(s: impl Into<String>) -> Self {
Self::Tsv { text: s.into() }
}
pub fn json<T: Serialize>(value: &T) -> Result<Self> {
Ok(Self::Json {
value: serde_json::to_value(value)?,
})
}
#[must_use]
pub const fn json_value(value: Value) -> Self {
Self::Json { value }
}
#[must_use]
pub fn image(data: Vec<u8>, media_type: &str) -> Self {
Self::Binary {
mime: parse_media_type_or_octet_stream(media_type),
content: data,
}
}
#[must_use]
pub fn binary(data: Vec<u8>) -> Self {
Self::Binary {
mime: mime::APPLICATION_OCTET_STREAM.essence_str().to_string(),
content: data,
}
}
#[must_use]
pub fn error(message: impl Into<String>) -> Self {
Self::Error {
message: message.into(),
}
}
#[must_use]
pub const fn is_done(&self) -> bool {
matches!(self, Self::Done)
}
#[must_use]
pub const fn is_error(&self) -> bool {
matches!(self, Self::Error { .. })
}
#[must_use]
pub fn as_text(&self) -> Option<&str> {
match self {
Self::Text { text } | Self::Tsv { text } => Some(text),
Self::Error { message } => Some(message),
Self::Done | Self::Json { .. } | Self::Binary { .. } => None,
}
}
#[must_use]
pub fn error_message(&self) -> Option<&str> {
match self {
Self::Error { message } => Some(message),
Self::Done
| Self::Text { .. }
| Self::Tsv { .. }
| Self::Json { .. }
| Self::Binary { .. } => None,
}
}
pub fn render_for_model(&self) -> Result<String> {
match self {
Self::Done => Ok(String::new()),
Self::Text { text } | Self::Tsv { text } => Ok(text.clone()),
Self::Json { value } => Ok(serde_json::to_string(value)?),
Self::Binary { mime, content } => {
let mut rendered = String::new();
rendered.push_str("[binary tool result: ");
rendered.push_str(mime);
rendered.push_str(", ");
rendered.push_str(content.len().to_string().as_str());
rendered.push_str(" bytes]");
Ok(rendered)
}
Self::Error { message } => Ok(message.clone()),
}
}
pub fn render_for_cli(&self) -> Result<String> {
match self {
Self::Done => Ok(String::new()),
Self::Text { text } | Self::Tsv { text } => Ok(text.clone()),
Self::Json { value } => Ok(serde_json::to_string_pretty(value)?),
Self::Binary { mime, content } => {
let mut rendered = String::new();
rendered.push_str("[binary tool result: ");
rendered.push_str(mime);
rendered.push_str(", ");
rendered.push_str(content.len().to_string().as_str());
rendered.push_str(" bytes]");
Ok(rendered)
}
Self::Error { message } => Ok(message.clone()),
}
}
#[must_use]
pub fn mime(&self) -> Option<Mime> {
match self {
Self::Binary { mime, .. } => mime.parse().ok(),
Self::Done
| Self::Text { .. }
| Self::Tsv { .. }
| Self::Json { .. }
| Self::Error { .. } => None,
}
}
#[must_use]
pub fn content(&self) -> Option<&[u8]> {
match self {
Self::Binary { content, .. } => Some(content),
Self::Done
| Self::Text { .. }
| Self::Tsv { .. }
| Self::Json { .. }
| Self::Error { .. } => None,
}
}
}
pub trait IntoToolResult {
fn into_tool_result(self) -> Result<ToolResult>;
}
impl IntoToolResult for ToolResult {
fn into_tool_result(self) -> Result<ToolResult> {
Ok(self)
}
}
impl IntoToolResult for () {
fn into_tool_result(self) -> Result<ToolResult> {
Ok(ToolResult::Done)
}
}
impl IntoToolResult for String {
fn into_tool_result(self) -> Result<ToolResult> {
Ok(ToolResult::text(self))
}
}
impl IntoToolResult for &str {
fn into_tool_result(self) -> Result<ToolResult> {
Ok(ToolResult::text(self))
}
}
impl IntoToolResult for Cow<'_, str> {
fn into_tool_result(self) -> Result<ToolResult> {
Ok(ToolResult::text(self.into_owned()))
}
}
impl IntoToolResult for Value {
fn into_tool_result(self) -> Result<ToolResult> {
Ok(ToolResult::json_value(self))
}
}
impl<T> IntoToolResult for Option<T>
where
T: IntoToolResult,
{
fn into_tool_result(self) -> Result<ToolResult> {
self.map_or_else(|| Ok(ToolResult::Done), IntoToolResult::into_tool_result)
}
}
impl<T, E> IntoToolResult for core::result::Result<T, E>
where
T: Serialize,
E: Display,
{
fn into_tool_result(self) -> Result<ToolResult> {
match self {
Ok(value) => serialize_success_value(&value),
Err(error) => Ok(ToolResult::error(error.to_string())),
}
}
}
fn parse_media_type_or_octet_stream(media_type: &str) -> String {
media_type
.parse::<Mime>()
.unwrap_or(mime::APPLICATION_OCTET_STREAM)
.essence_str()
.to_string()
}
fn serialize_success_value<T: Serialize>(value: &T) -> Result<ToolResult> {
let value = serde_json::to_value(value)?;
if let Some(tsv) = json_value_to_tsv(&value) {
return Ok(ToolResult::tsv(tsv));
}
match value {
Value::String(text) => Ok(ToolResult::text(text)),
other => Ok(ToolResult::json_value(other)),
}
}
#[must_use]
pub fn json_value_to_tsv(value: &Value) -> Option<String> {
let rows = match value {
Value::Array(arr) if !arr.is_empty() => arr
.iter()
.map(|value| flatten_json_value(value, ""))
.collect::<Vec<_>>(),
Value::Object(_) => alloc::vec![flatten_json_value(value, "")],
Value::Array(_) | Value::String(_) | Value::Number(_) | Value::Bool(_) | Value::Null => {
return None;
}
};
if rows.is_empty() {
return None;
}
let mut columns: Vec<String> = Vec::new();
let mut seen: alloc::collections::BTreeSet<String> = alloc::collections::BTreeSet::new();
for row in &rows {
for (key, _) in row {
if seen.insert(key.clone()) {
columns.push(key.clone());
}
}
}
if columns.is_empty() {
return None;
}
let mut tsv = String::new();
for (index, column) in columns.iter().enumerate() {
if index > 0 {
tsv.push('\t');
}
tsv.push_str(&escape_tsv_field(column));
}
tsv.push('\n');
for row in &rows {
let row_map: alloc::collections::BTreeMap<&str, &str> = row
.iter()
.map(|(key, value)| (key.as_str(), value.as_str()))
.collect::<alloc::collections::BTreeMap<&str, &str>>();
for (index, column) in columns.iter().enumerate() {
if index > 0 {
tsv.push('\t');
}
if let Some(value) = row_map.get(column.as_str()) {
tsv.push_str(&escape_tsv_field(value));
}
}
tsv.push('\n');
}
Some(tsv)
}
fn flatten_json_value(value: &Value, prefix: &str) -> Vec<(String, String)> {
let mut flattened = Vec::new();
match value {
Value::Object(map) => {
for (key, child) in map {
let full_key = if prefix.is_empty() {
key.clone()
} else {
format!("{prefix}.{key}")
};
flattened.extend(flatten_json_value(child, &full_key));
}
}
Value::Array(_) => {
let serialized = serde_json::to_string(value).unwrap_or_default();
flattened.push((prefix.to_string(), serialized));
}
Value::String(text) => {
flattened.push((prefix.to_string(), text.clone()));
}
Value::Number(number) => {
flattened.push((prefix.to_string(), number.to_string()));
}
Value::Bool(boolean) => {
flattened.push((prefix.to_string(), boolean.to_string()));
}
Value::Null => {
flattened.push((prefix.to_string(), String::new()));
}
}
flattened
}
fn escape_tsv_field(value: &str) -> String {
value.replace(['\t', '\n', '\r'], " ")
}
pub trait Tool: Send + Sync {
fn name(&self) -> Cow<'static, str>;
fn description(&self) -> Cow<'static, str> {
description_from_schema::<Self::Arguments>().unwrap_or_default()
}
type Arguments: Send + JsonSchema + DeserializeOwned;
type Res: IntoToolResult + Send;
fn call(&self, arguments: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send;
}
pub fn json<T: Serialize>(value: &T) -> Result<String> {
let value = serde_json::to_value(value)?;
Ok(value
.as_str()
.map_or_else(|| format!("{value:#}"), ToString::to_string))
}
trait ToolImpl: Send + Sync + Any {
fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>>;
fn definition(&self) -> &ToolDefinition;
fn as_any(&self) -> &dyn Any;
fn as_any_mut(&mut self) -> &mut dyn Any;
}
struct DynToolImpl<F>
where
F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>> + Send + Sync,
{
definition: ToolDefinition,
handler: F,
}
impl<F> ToolImpl for DynToolImpl<F>
where
F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>> + Send + Sync + 'static,
{
fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>> {
(self.handler)(args)
}
fn definition(&self) -> &ToolDefinition {
&self.definition
}
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
}
fn schema_is_object(value: &Value) -> bool {
matches!(value.get("type").and_then(Value::as_str), Some("object"))
|| value.get("properties").is_some()
|| value.get("oneOf").is_some()
|| value.get("anyOf").is_some()
|| value.get("$defs").is_some()
}
fn is_object<T: JsonSchema>() -> bool {
schema_is_object(&schema_for!(T).to_value())
}
fn arguments_schema<T: JsonSchema>() -> Schema {
if is_object::<T>() {
schema_for!(T)
} else {
schema_for!(ToolArgument<T>)
}
}
fn description_from_schema<T: JsonSchema>() -> Option<Cow<'static, str>> {
schema_for!(T)
.to_value()
.get("description")
.and_then(Value::as_str)
.filter(|text| !text.trim().is_empty())
.map(|text| Cow::Owned(text.to_string()))
}
struct RegisteredTool<T: Tool> {
tool: T,
definition: ToolDefinition,
args_are_object: bool,
}
impl<T: Tool> RegisteredTool<T> {
fn new(tool: T) -> Self {
let definition = ToolDefinition::new(&tool);
let args_are_object = is_object::<T::Arguments>();
Self {
tool,
definition,
args_are_object,
}
}
}
impl<T: Tool + 'static> ToolImpl for RegisteredTool<T> {
fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>> {
let result = if self.args_are_object {
serde_json::from_str::<T::Arguments>(args)
} else {
serde_json::from_str::<ToolArgument<T::Arguments>>(args).map(|wrapper| wrapper.value)
};
let Ok(arguments) = result else {
let name = self.definition.name().to_string();
let schema_str =
serde_json::to_string_pretty(&self.definition.arguments_openai_schema())
.unwrap_or_else(|_| "{}".to_string());
return Box::pin(async move {
Err(anyhow::Error::msg(format!(
"Invalid arguments for tool '{name}'. Expected schema:\n{schema_str}"
)))
});
};
Box::pin(async move { Tool::call(&self.tool, arguments).await?.into_tool_result() })
}
fn definition(&self) -> &ToolDefinition {
&self.definition
}
fn as_any(&self) -> &dyn Any {
&self.tool
}
fn as_any_mut(&mut self) -> &mut dyn Any {
&mut self.tool
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InvalidSchema {
name: Cow<'static, str>,
}
impl InvalidSchema {
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
impl Display for InvalidSchema {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(
f,
"tool '{}' has an argument schema that is neither an object nor a boolean",
self.name
)
}
}
impl core::error::Error for InvalidSchema {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RegisterError {
DuplicateName(Cow<'static, str>),
EmptyDescription(Cow<'static, str>),
}
impl Display for RegisterError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::DuplicateName(name) => {
write!(f, "a tool named '{name}' is already registered")
}
Self::EmptyDescription(name) => write!(
f,
"tool '{name}' has an empty description; add a rustdoc comment to its \
Arguments type or implement Tool::description"
),
}
}
}
impl core::error::Error for RegisterError {}
pub struct Tools {
tools: BTreeMap<Cow<'static, str>, Box<dyn ToolImpl>>,
}
impl Debug for Tools {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Tools")
.field("tools", &self.tools.keys().collect::<Vec<_>>())
.finish()
}
}
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ToolDefinition {
name: Cow<'static, str>,
description: Cow<'static, str>,
arguments: Schema,
}
impl ToolDefinition {
#[must_use]
pub fn new<T: Tool>(tool: &T) -> Self {
Self {
name: tool.name(),
description: tool.description(),
arguments: arguments_schema::<T::Arguments>(),
}
}
pub fn from_parts(
name: Cow<'static, str>,
description: Cow<'static, str>,
schema: Value,
) -> core::result::Result<Self, InvalidSchema> {
let arguments: Schema = schema
.try_into()
.map_err(|_| InvalidSchema { name: name.clone() })?;
Ok(Self {
name,
description,
arguments,
})
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn description(&self) -> &str {
&self.description
}
#[must_use]
pub fn arguments_openai_schema(&self) -> serde_json::Value {
let mut inner = self.arguments.clone().to_value();
clean_schema(&mut inner);
inner
}
}
#[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
struct ToolArgument<T> {
value: T,
}
fn clean_schema(value: &mut Value) {
let defs = extract_defs(value);
resolve_and_clean(value, &defs);
if let Value::Object(map) = value {
map.remove("description");
if map.contains_key("properties") && !map.contains_key("type") {
map.insert("type".to_string(), Value::String("object".to_string()));
}
}
}
fn extract_defs(value: &Value) -> serde_json::Map<String, Value> {
if let Value::Object(map) = value
&& let Some(Value::Object(defs)) = map.get("$defs").or_else(|| map.get("definitions"))
{
return defs.clone();
}
serde_json::Map::new()
}
#[allow(clippy::too_many_lines)]
fn resolve_and_clean(value: &mut Value, defs: &serde_json::Map<String, Value>) {
resolve_and_clean_inner(value, defs, false);
}
#[allow(clippy::too_many_lines)]
fn resolve_and_clean_inner(
value: &mut Value,
defs: &serde_json::Map<String, Value>,
inside_properties: bool,
) {
match value {
Value::Object(map) => {
if let Some(Value::String(ref_path)) = map.remove("$ref")
&& let Some(Value::Object(resolved_map)) = resolve_ref(&ref_path, defs)
{
let existing_description = map.remove("description");
for (k, v) in resolved_map {
map.entry(k).or_insert(v);
}
if let Some(desc) = existing_description {
map.insert("description".to_string(), desc);
}
}
if let Some(const_val) = map.remove("const") {
map.insert("enum".to_string(), Value::Array(alloc::vec![const_val]));
}
if let Some(Value::Array(variants)) =
map.remove("oneOf").or_else(|| map.remove("anyOf"))
{
let is_simple_enum = variants.iter().all(|v| {
if let Value::Object(vm) = v {
(vm.contains_key("const") || vm.contains_key("enum"))
&& !vm.contains_key("properties")
} else {
false
}
});
if is_simple_enum {
let mut enum_values: alloc::vec::Vec<Value> = alloc::vec::Vec::new();
let mut variant_type: Option<String> = None;
for variant in &variants {
if let Value::Object(vm) = variant {
if let Some(const_val) = vm.get("const")
&& !enum_values.contains(const_val)
{
enum_values.push(const_val.clone());
}
if let Some(Value::Array(arr)) = vm.get("enum") {
for val in arr {
if !enum_values.contains(val) {
enum_values.push(val.clone());
}
}
}
if variant_type.is_none()
&& let Some(Value::String(t)) = vm.get("type")
{
variant_type = Some(t.clone());
}
}
}
if !enum_values.is_empty() {
map.insert("enum".to_string(), Value::Array(enum_values));
if let Some(t) = variant_type {
map.insert("type".to_string(), Value::String(t));
}
}
} else {
let mut all_properties = serde_json::Map::new();
for variant in variants {
if let Value::Object(variant_map) = variant
&& let Some(Value::Object(props)) = variant_map.get("properties")
{
for (key, val) in props {
let new_values: Option<alloc::vec::Vec<Value>> =
if let Value::Object(val_obj) = val {
if let Some(Value::Array(arr)) = val_obj.get("enum") {
Some(arr.clone())
} else {
val_obj
.get("const")
.map(|const_val| alloc::vec![const_val.clone()])
}
} else {
None
};
if all_properties.contains_key(key) {
if let Some(values) = new_values
&& let Some(Value::Object(existing_obj)) =
all_properties.get_mut(key)
&& let Some(Value::Array(existing_enum)) =
existing_obj.get_mut("enum")
{
for e in values {
if !existing_enum.contains(&e) {
existing_enum.push(e);
}
}
}
} else {
let mut val_clone = val.clone();
if let Value::Object(obj) = &mut val_clone
&& let Some(const_val) = obj.remove("const")
{
obj.insert(
"enum".to_string(),
Value::Array(alloc::vec![const_val]),
);
}
all_properties.insert(key.clone(), val_clone);
}
}
}
}
if !all_properties.is_empty() {
map.insert("type".to_string(), Value::String("object".to_string()));
map.insert("properties".to_string(), Value::Object(all_properties));
}
}
}
if !inside_properties {
let allowed = [
"type",
"description",
"properties",
"required",
"items",
"enum",
"nullable",
];
map.retain(|k, _| allowed.contains(&k.as_str()));
}
if let Some(Value::Array(types)) = map.get("type") {
let non_null: Vec<&Value> = types
.iter()
.filter(|t| !matches!(t, Value::String(s) if s == "null"))
.collect();
if non_null.len() == 1 {
map.insert("type".to_string(), non_null[0].clone());
}
}
for (key, v) in map.iter_mut() {
let child_inside_props = key == "properties";
resolve_and_clean_inner(v, defs, child_inside_props);
}
}
Value::Array(arr) => {
for v in arr {
resolve_and_clean_inner(v, defs, false);
}
}
_ => {}
}
}
fn resolve_ref(ref_path: &str, defs: &serde_json::Map<String, Value>) -> Option<Value> {
let name = ref_path
.strip_prefix("#/$defs/")
.or_else(|| ref_path.strip_prefix("#/definitions/"))?;
defs.get(name).cloned()
}
impl Default for Tools {
fn default() -> Self {
Self::new()
}
}
impl Tools {
#[must_use]
pub const fn new() -> Self {
Self {
tools: BTreeMap::new(),
}
}
#[must_use]
pub fn get<T>(&self) -> Option<&T>
where
T: Tool + 'static,
{
self.tools
.values()
.find_map(|tool| tool.as_any().downcast_ref::<T>())
}
#[must_use]
pub fn get_mut<T>(&mut self) -> Option<&mut T>
where
T: Tool + 'static,
{
self.tools
.values_mut()
.find_map(|tool| tool.as_any_mut().downcast_mut::<T>())
}
#[must_use]
pub fn definitions(&self) -> Vec<ToolDefinition> {
self.tools
.values()
.map(|tool| tool.definition().clone())
.collect()
}
pub fn register<T: Tool + 'static>(
&mut self,
tool: T,
) -> core::result::Result<(), RegisterError> {
self.insert(Box::new(RegisteredTool::new(tool)))
}
pub fn register_dyn<F>(
&mut self,
definition: ToolDefinition,
handler: F,
) -> core::result::Result<(), RegisterError>
where
F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>>
+ Send
+ Sync
+ 'static,
{
self.insert(Box::new(DynToolImpl {
definition,
handler,
}))
}
fn insert(&mut self, tool: Box<dyn ToolImpl>) -> core::result::Result<(), RegisterError> {
let name = tool.definition().name.clone();
if self.tools.contains_key(&name) {
return Err(RegisterError::DuplicateName(name));
}
if tool.definition().description().trim().is_empty() {
return Err(RegisterError::EmptyDescription(name));
}
self.tools.insert(name, tool);
Ok(())
}
pub fn unregister(&mut self, name: &str) {
self.tools.remove(name);
}
pub async fn call(&self, name: &str, args: &str) -> Result<ToolResult> {
if let Some(tool) = self.tools.get(name) {
tool.call(args).await
} else {
Err(anyhow::Error::msg(format!("Tool '{name}' not found")))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::{format, string::ToString, vec};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
#[derive(JsonSchema, Deserialize, Debug, PartialEq)]
struct CalculatorArgs {
operation: String,
a: f64,
b: f64,
}
struct Calculator;
impl Tool for Calculator {
fn name(&self) -> Cow<'static, str> {
"calculator".into()
}
type Arguments = CalculatorArgs;
type Res = ToolResult;
fn call(&self, args: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send {
core::future::ready(match args.operation.as_str() {
"add" => Ok(ToolResult::text((args.a + args.b).to_string())),
"subtract" => Ok(ToolResult::text((args.a - args.b).to_string())),
"multiply" => Ok(ToolResult::text((args.a * args.b).to_string())),
"divide" => {
if args.b == 0.0 {
Err(anyhow::Error::msg("Division by zero"))
} else {
Ok(ToolResult::text((args.a / args.b).to_string()))
}
}
_ => Err(anyhow::Error::msg(format!(
"Unknown operation: {}",
args.operation
))),
})
}
}
#[derive(JsonSchema, Deserialize)]
struct GreetArgs {
name: String,
}
struct Greeter;
impl Tool for Greeter {
fn name(&self) -> Cow<'static, str> {
"greeter".into()
}
type Arguments = GreetArgs;
type Res = ToolResult;
fn call(&self, args: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send {
core::future::ready(Ok(ToolResult::text(format!("Hello, {}!", args.name))))
}
}
#[test]
fn from_parts_accepts_object_and_boolean_schemas() {
for schema in [
serde_json::json!({"type": "object"}),
serde_json::json!(true),
] {
assert!(
ToolDefinition::from_parts("t".into(), "does a thing".into(), schema.clone())
.is_ok(),
"{schema} should be accepted"
);
}
}
#[test]
fn from_parts_rejects_non_schema_values() {
for schema in [
serde_json::json!("a string"),
serde_json::json!([1, 2, 3]),
serde_json::json!(7),
serde_json::json!(null),
] {
let result = ToolDefinition::from_parts("weird".into(), "d".into(), schema.clone());
let Err(err) = result else {
panic!("{schema} should be rejected");
};
assert_eq!(err.name(), "weird");
}
}
#[test]
fn json_utility() {
let value = serde_json::json!({
"name": "test",
"value": 42
});
let json_str = json(&value).expect("a JSON value always serializes");
assert!(json_str.contains("\"name\": \"test\""));
assert!(json_str.contains("\"value\": 42"));
}
#[test]
fn tool_definition_creation() {
let calculator = Calculator;
let definition = ToolDefinition::new(&calculator);
assert_eq!(definition.name, "calculator");
assert_eq!(
definition.description,
"Performs basic mathematical operations."
);
}
#[test]
fn tools_creation() {
let tools = Tools::new();
assert_eq!(tools.definitions().len(), 0);
}
#[test]
fn tools_default() {
let tools = Tools::default();
assert_eq!(tools.definitions().len(), 0);
}
#[tokio::test]
async fn tools_register_and_call() {
let mut tools = Tools::new();
tools.register(Calculator).expect("calculator registers");
let definitions = tools.definitions();
assert_eq!(definitions.len(), 1);
assert_eq!(definitions[0].name, "calculator");
let result = tools
.call("calculator", r#"{"operation": "add", "a": 5, "b": 3}"#)
.await;
assert!(result.is_ok());
assert_eq!(result.unwrap().as_text(), Some("8"));
}
#[tokio::test]
async fn calculator_operations() {
let mut tools = Tools::new();
tools.register(Calculator).expect("calculator registers");
let result = tools
.call("calculator", r#"{"operation": "add", "a": 10, "b": 5}"#)
.await;
assert_eq!(result.unwrap().as_text(), Some("15"));
let result = tools
.call(
"calculator",
r#"{"operation": "subtract", "a": 10, "b": 3}"#,
)
.await;
assert_eq!(result.unwrap().as_text(), Some("7"));
let result = tools
.call("calculator", r#"{"operation": "multiply", "a": 4, "b": 3}"#)
.await;
assert_eq!(result.unwrap().as_text(), Some("12"));
let result = tools
.call("calculator", r#"{"operation": "divide", "a": 15, "b": 3}"#)
.await;
assert_eq!(result.unwrap().as_text(), Some("5"));
}
#[tokio::test]
async fn calculator_division_by_zero() {
let mut tools = Tools::new();
tools.register(Calculator).expect("calculator registers");
let result = tools
.call("calculator", r#"{"operation": "divide", "a": 10, "b": 0}"#)
.await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("Division by zero"));
}
#[tokio::test]
async fn calculator_unknown_operation() {
let mut tools = Tools::new();
tools.register(Calculator).expect("calculator registers");
let result = tools
.call("calculator", r#"{"operation": "modulo", "a": 10, "b": 3}"#)
.await;
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("Unknown operation")
);
}
#[tokio::test]
async fn multiple_tools() {
let mut tools = Tools::new();
tools.register(Calculator).expect("calculator registers");
tools.register(Greeter).expect("greeter registers");
let definitions = tools.definitions();
assert_eq!(definitions.len(), 2);
let calc_def = definitions.iter().find(|d| d.name == "calculator").unwrap();
let greet_def = definitions.iter().find(|d| d.name == "greeter").unwrap();
assert_eq!(
calc_def.description,
"Performs basic mathematical operations."
);
assert_eq!(greet_def.description, "Greets a person by name.");
let calc_result = tools
.call("calculator", r#"{"operation": "add", "a": 2, "b": 3}"#)
.await;
assert_eq!(calc_result.unwrap().as_text(), Some("5"));
let greet_result = tools.call("greeter", r#"{"name": "Alice"}"#).await;
assert_eq!(greet_result.unwrap().as_text(), Some("Hello, Alice!"));
}
#[derive(Debug, Serialize)]
struct TableRow {
name: &'static str,
count: u32,
}
#[derive(Debug, Serialize)]
struct NestedTableRow {
user: TableRow,
ok: bool,
}
#[derive(Debug)]
struct ToolFailure(&'static str);
impl core::fmt::Display for ToolFailure {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.write_str(self.0)
}
}
impl core::error::Error for ToolFailure {}
#[test]
fn into_tool_result_string_is_plain_text() {
assert_eq!(
String::from("hello").into_tool_result().unwrap(),
ToolResult::text("hello")
);
}
#[test]
fn into_tool_result_str_is_plain_text() {
assert_eq!(
"hello".into_tool_result().unwrap(),
ToolResult::text("hello")
);
}
#[test]
fn into_tool_result_option_none_is_done() {
let result = Option::<String>::None.into_tool_result().unwrap();
assert_eq!(result, ToolResult::Done);
}
#[test]
fn into_tool_result_option_some_delegates() {
let result = Some("hello").into_tool_result().unwrap();
assert_eq!(result, ToolResult::text("hello"));
}
#[test]
fn into_tool_result_result_ok_string_is_text() {
let result = core::result::Result::<String, ToolFailure>::Ok(String::from("hello"))
.into_tool_result()
.unwrap();
assert_eq!(result, ToolResult::text("hello"));
}
#[test]
fn into_tool_result_result_ok_object_is_tsv() {
let result = core::result::Result::<TableRow, ToolFailure>::Ok(TableRow {
name: "alpha",
count: 3,
})
.into_tool_result()
.unwrap();
assert_eq!(result, ToolResult::tsv("count\tname\n3\talpha\n"));
}
#[test]
fn into_tool_result_result_ok_array_of_objects_is_tsv() {
let result = core::result::Result::<Vec<TableRow>, ToolFailure>::Ok(vec![
TableRow {
name: "alpha",
count: 3,
},
TableRow {
name: "beta",
count: 5,
},
])
.into_tool_result()
.unwrap();
assert_eq!(result, ToolResult::tsv("count\tname\n3\talpha\n5\tbeta\n"));
}
#[test]
fn into_tool_result_result_ok_scalar_is_json() {
let result = core::result::Result::<bool, ToolFailure>::Ok(true)
.into_tool_result()
.unwrap();
assert_eq!(result, ToolResult::json_value(Value::Bool(true)));
}
#[test]
fn into_tool_result_result_err_is_typed_error() {
let result = core::result::Result::<TableRow, ToolFailure>::Err(ToolFailure("boom"))
.into_tool_result()
.unwrap();
assert_eq!(result, ToolResult::error("boom"));
assert!(result.is_error());
assert_eq!(result.error_message(), Some("boom"));
}
#[test]
fn json_value_to_tsv_flattens_nested_objects() {
let value = serde_json::to_value(NestedTableRow {
user: TableRow {
name: "alpha",
count: 3,
},
ok: true,
})
.unwrap();
assert_eq!(
json_value_to_tsv(&value),
Some("ok\tuser.count\tuser.name\ntrue\t3\talpha\n".to_string())
);
}
#[tokio::test]
async fn tool_not_found() {
let tools = Tools::new();
let result = tools.call("nonexistent", "{}").await;
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("Tool 'nonexistent' not found")
);
}
#[tokio::test]
async fn invalid_json() {
let mut tools = Tools::new();
tools.register(Calculator).expect("calculator registers");
let result = tools.call("calculator", "invalid json").await;
assert!(result.is_err());
}
#[test]
fn tools_unregister() {
let mut tools = Tools::new();
tools.register(Calculator).expect("calculator registers");
tools.register(Greeter).expect("greeter registers");
assert_eq!(tools.definitions().len(), 2);
tools.unregister("calculator");
assert_eq!(tools.definitions().len(), 1);
let remaining = &tools.definitions()[0];
assert_eq!(remaining.name, "greeter");
tools.unregister("greeter");
assert_eq!(tools.definitions().len(), 0);
}
#[test]
fn tools_debug() {
let mut tools = Tools::new();
tools.register(Calculator).expect("calculator registers");
tools.register(Greeter).expect("greeter registers");
let debug_str = format!("{tools:?}");
assert!(debug_str.contains("Tools"));
assert!(debug_str.contains("calculator"));
assert!(debug_str.contains("greeter"));
}
#[test]
fn tool_definition_debug() {
let calculator = Calculator;
let definition = ToolDefinition::new(&calculator);
let debug_str = format!("{definition:?}");
assert!(debug_str.contains("ToolDefinition"));
assert!(debug_str.contains("calculator"));
assert!(debug_str.contains("Performs basic mathematical operations"));
}
#[test]
fn tool_definition_clone() {
let calculator = Calculator;
let original = ToolDefinition::new(&calculator);
let cloned = original.clone();
assert_eq!(original.name, cloned.name);
assert_eq!(original.description, cloned.description);
}
#[test]
fn schema_preserves_enum() {
#[derive(JsonSchema, Deserialize)]
#[serde(rename_all = "snake_case")]
enum Status {
Pending,
InProgress,
Completed,
}
#[allow(dead_code)]
#[derive(JsonSchema, Deserialize)]
struct Item {
status: Status,
}
#[allow(dead_code)]
#[derive(JsonSchema, Deserialize)]
struct Args {
items: Vec<Item>,
}
struct TestTool;
impl Tool for TestTool {
fn name(&self) -> Cow<'static, str> {
"test".into()
}
type Arguments = Args;
type Res = ToolResult;
fn call(
&self,
_args: Self::Arguments,
) -> impl Future<Output = Result<Self::Res>> + Send {
core::future::ready(Ok(ToolResult::text("ok")))
}
}
let tool = TestTool;
let def = ToolDefinition::new(&tool);
let schema = def.arguments_openai_schema();
let schema_obj = schema.as_object().expect("schema should be object");
let properties = schema_obj
.get("properties")
.expect("should have properties")
.as_object()
.unwrap();
let items = properties
.get("items")
.expect("should have items")
.as_object()
.unwrap();
let item_props = items
.get("items")
.expect("items should have items schema")
.as_object()
.unwrap();
let item_properties = item_props
.get("properties")
.expect("item should have properties")
.as_object()
.unwrap();
let status = item_properties
.get("status")
.expect("should have status")
.as_object()
.unwrap();
assert!(
status.contains_key("enum"),
"Status should have enum field. Full schema: {}",
serde_json::to_string_pretty(&schema).unwrap()
);
}
#[test]
fn schema_ref_resolution() {
let raw_schema = serde_json::json!({
"type": "object",
"properties": {
"status": {
"$ref": "#/$defs/Status"
}
},
"$defs": {
"Status": {
"type": "string",
"enum": ["pending", "in_progress", "completed"]
}
}
});
let mut schema = raw_schema;
clean_schema(&mut schema);
let props = schema.get("properties").unwrap().as_object().unwrap();
let status = props.get("status").unwrap().as_object().unwrap();
assert!(
status.contains_key("enum"),
"Status should have enum after ref resolution. Got: {}",
serde_json::to_string_pretty(&schema).unwrap()
);
}
#[test]
fn schema_nested_ref_in_array() {
let raw_schema = serde_json::json!({
"type": "object",
"properties": {
"todos": {
"type": "array",
"items": {
"$ref": "#/$defs/TodoItem"
}
}
},
"$defs": {
"TodoItem": {
"type": "object",
"properties": {
"content": { "type": "string" },
"status": { "$ref": "#/$defs/TodoStatus" }
},
"required": ["content", "status"]
},
"TodoStatus": {
"type": "string",
"enum": ["pending", "in_progress", "completed"]
}
}
});
let mut schema = raw_schema;
clean_schema(&mut schema);
let props = schema.get("properties").unwrap().as_object().unwrap();
let todos = props.get("todos").unwrap().as_object().unwrap();
let items = todos.get("items").unwrap().as_object().unwrap();
let item_props = items.get("properties").unwrap().as_object().unwrap();
let status = item_props.get("status").unwrap().as_object().unwrap();
assert!(
status.contains_key("enum"),
"Nested status should have enum. Full schema: {}",
serde_json::to_string_pretty(&schema).unwrap()
);
}
}