use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use log::warn;
pub type ToolFunction = Box<dyn Fn(Option<HashMap<String, Value>>) -> Result<Value, String> + Send + Sync>;
#[derive(Debug, Clone)]
pub enum DuplicateBehavior {
Warn,
Error,
Replace,
Ignore,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ToolAnnotations {
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(rename = "readOnlyHint", skip_serializing_if = "Option::is_none")]
pub read_only_hint: Option<bool>,
#[serde(rename = "destructiveHint", skip_serializing_if = "Option::is_none")]
pub destructive_hint: Option<bool>,
#[serde(rename = "idempotentHint", skip_serializing_if = "Option::is_none")]
pub idempotent_hint: Option<bool>,
#[serde(rename = "openWorldHint", skip_serializing_if = "Option::is_none")]
pub open_world_hint: Option<bool>,
}
#[derive(Serialize, Deserialize)]
pub struct FunctionTool {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
pub description: String,
#[serde(rename = "inputSchema", skip_serializing_if = "Option::is_none")]
pub input_schema: Option<Value>,
#[serde(rename = "outputSchema", skip_serializing_if = "Option::is_none")]
pub output_schema: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub annotations: Option<ToolAnnotations>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tags: Option<Vec<String>>,
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<Value>,
#[serde(skip)]
function: Option<Arc<ToolFunction>>,
}
impl Clone for FunctionTool {
fn clone(&self) -> Self {
Self {
name: self.name.clone(),
title: self.title.clone(),
description: self.description.clone(),
input_schema: self.input_schema.clone(),
output_schema: self.output_schema.clone(),
annotations: self.annotations.clone(),
tags: self.tags.clone(),
meta: self.meta.clone(),
function: None, }
}
}
impl std::fmt::Debug for FunctionTool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FunctionTool")
.field("name", &self.name)
.field("title", &self.title)
.field("description", &self.description)
.field("input_schema", &self.input_schema)
.field("output_schema", &self.output_schema)
.field("annotations", &self.annotations)
.field("tags", &self.tags)
.field("meta", &self.meta)
.finish()
}
}
impl FunctionTool {
#[allow(clippy::too_many_arguments)]
#[allow(dead_code)]
pub fn from_function<F>(
function: F,
name: Option<String>,
title: Option<String>,
description: Option<String>,
input_schema: Option<Value>,
output_schema: Option<Value>,
annotations: Option<ToolAnnotations>,
tags: Option<Vec<String>>,
meta: Option<Value>,
) -> Self
where
F: Fn(Option<HashMap<String, Value>>) -> Result<Value, String> + Send + Sync + 'static,
{
Self {
function: Some(Arc::new(Box::new(function))),
name: name.unwrap_or_else(|| "unnamed_tool".to_string()),
title,
description: description.unwrap_or_default(),
input_schema,
output_schema,
annotations,
tags,
meta,
}
}
#[allow(dead_code)]
pub fn call(&self, args: Option<HashMap<String, Value>>) -> Result<Value, String> {
if let Some(ref function) = self.function {
function(args)
} else {
Err("Tool function not available".to_string())
}
}
}
#[derive(Debug, Clone)]
pub struct ToolManager {
tools: HashMap<String, FunctionTool>,
duplicate_behavior: DuplicateBehavior,
}
impl ToolManager {
pub fn new() -> Self {
Self {
tools: HashMap::new(),
duplicate_behavior: DuplicateBehavior::Warn,
}
}
pub fn with_behavior(duplicate_behavior: DuplicateBehavior) -> Self {
Self {
tools: HashMap::new(),
duplicate_behavior,
}
}
}
impl Default for ToolManager {
fn default() -> Self {
Self::new()
}
}
impl ToolManager {
#[allow(dead_code)]
pub fn add_tool(&mut self, tool: FunctionTool) {
if self.tools.contains_key(&tool.name) {
match self.duplicate_behavior {
DuplicateBehavior::Warn => {
warn!("Tool '{}' already exists, replacing", tool.name);
self.tools.insert(tool.name.clone(), tool);
}
DuplicateBehavior::Error => {
panic!("Tool '{}' already exists", tool.name);
}
DuplicateBehavior::Replace => {
self.tools.insert(tool.name.clone(), tool);
}
DuplicateBehavior::Ignore => {
}
}
} else {
self.tools.insert(tool.name.clone(), tool);
}
}
#[allow(dead_code)]
pub fn get_tool(&self, name: &str) -> Option<&FunctionTool> {
self.tools.get(name)
}
#[allow(dead_code)]
pub fn list_tools(&self) -> Vec<&FunctionTool> {
self.tools.values().collect()
}
#[allow(dead_code)]
pub fn call_tool(&self, name: &str, args: Option<HashMap<String, Value>>) -> Result<Value, String> {
if let Some(tool) = self.get_tool(name) {
tool.call(args)
} else {
Err(format!("Tool '{}' not found", name))
}
}
}