use anyhow::Result;
use async_trait::async_trait;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
#[async_trait]
pub trait LlmTool<C>: Send + Sync {
fn name(&self) -> &'static str;
fn description(&self) -> &'static str;
fn parameters(&self) -> Value;
fn normalize_args(&self, args: &Value) -> Value {
args.clone()
}
async fn call(&self, args: Value, context: &C) -> Result<Value>;
}
pub struct ToolRegistry<C> {
tools: HashMap<&'static str, Arc<dyn LlmTool<C>>>,
}
impl<C> ToolRegistry<C> {
pub fn new() -> Self {
Self {
tools: HashMap::new(),
}
}
pub fn register(&mut self, tool: impl LlmTool<C> + 'static) {
self.tools.insert(tool.name(), Arc::new(tool));
}
pub fn register_boxed(&mut self, tool: Box<dyn LlmTool<C>>) {
self.tools.insert(tool.name(), Arc::from(tool));
}
pub fn declarations(&self) -> Vec<Value> {
self.tools
.values()
.map(|t| {
serde_json::json!({
"name": t.name(),
"description": t.description(),
"parameters": t.parameters(),
})
})
.collect()
}
pub fn normalize_tool_args(&self, name: &str, args: &Value) -> Value {
if let Some(tool) = self.tools.get(name) {
tool.normalize_args(args)
} else {
args.clone()
}
}
pub async fn call(&self, name: &str, args: Value, context: &C) -> Result<Value> {
let tool = self
.tools
.get(name)
.ok_or_else(|| anyhow::anyhow!("Tool not found in registry: {}", name))?;
tool.call(args, context).await
}
}
impl<C> Default for ToolRegistry<C> {
fn default() -> Self {
Self::new()
}
}