ic-rig 0.2.0

//! Tool trait and type-erased tool dispatch.
//!
//! A [`Tool`] is a typed async function the model can invoke. Define the schema
//! via [`ToolDefinition`] and implement [`Tool::call`] with your logic.
//!
//! [`ToolSet`] holds a collection of type-erased tools keyed by name. The
//! [`Agent`](crate::agent::Agent) uses it to dispatch model-requested tool calls.
//!
//! # Example
//!
//! ```rust
//! use irig::tool::{Tool, ToolDefinition, ToolSet};
//! use serde::Deserialize;
//! use serde_json::json;
//!
//! #[derive(Deserialize)]
//! struct AddArgs { x: f64, y: f64 }
//!
//! struct Adder;
//!
//! impl Tool for Adder {
//!     const NAME: &'static str = "add";
//!     type Error = std::convert::Infallible;
//!     type Args  = AddArgs;
//!     type Output = f64;
//!
//!     fn definition(&self) -> ToolDefinition {
//!         ToolDefinition {
//!             name: "add".into(),
//!             description: "Add two numbers.".into(),
//!             parameters: json!({
//!                 "type": "object",
//!                 "properties": {
//!                     "x": { "type": "number" },
//!                     "y": { "type": "number" }
//!                 },
//!                 "required": ["x", "y"]
//!             }),
//!         }
//!     }
//!
//!     async fn call(&self, args: AddArgs) -> Result<f64, std::convert::Infallible> {
//!         Ok(args.x + args.y)
//!     }
//! }
//!
//! let mut tools = ToolSet::new();
//! tools.add(Adder);
//! ```

use crate::wasm_compat::BoxFuture;
use serde::{Serialize, de::DeserializeOwned};
use std::collections::HashMap;
use thiserror::Error;

// ── ToolDefinition ────────────────────────────────────────────────────────────

/// Describes a tool to the LLM.
///
/// The `parameters` field must be a valid JSON Schema object describing the
/// arguments the model should provide when calling this tool.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    /// JSON Schema for the tool's arguments.
    pub parameters: serde_json::Value,
}

// ── Tool trait ────────────────────────────────────────────────────────────────

/// A typed, async tool that an [`Agent`](crate::agent::Agent) can call.
///
/// Implement this trait to create tools. Register them with a [`ToolSet`].
pub trait Tool {
    /// Must match `ToolDefinition::name` exactly.
    const NAME: &'static str;

    /// Error type returned if the tool fails.
    type Error: std::error::Error + 'static;

    /// Deserialised argument type. Must match the `parameters` schema.
    type Args: DeserializeOwned;

    /// Serialisable output type.
    type Output: Serialize;

    /// Schema + description sent to the model before each completion call.
    fn definition(&self) -> ToolDefinition;

    /// Execute the tool with the deserialised arguments.
    fn call(&self, args: Self::Args) -> impl std::future::Future<Output = Result<Self::Output, Self::Error>>;
}

// ── ToolError ─────────────────────────────────────────────────────────────────

/// Errors that can occur when dispatching a tool call.
#[derive(Debug, Error)]
pub enum ToolError {
    /// The model requested a tool that is not registered.
    #[error("Unknown tool: {0}")]
    NotFound(String),

    /// The model-supplied arguments could not be deserialised.
    #[error("Failed to deserialise tool arguments: {0}")]
    ArgsParse(#[from] serde_json::Error),

    /// The tool itself returned an error.
    #[error("Tool execution failed: {0}")]
    Execution(String),
}

// ── Type-erased tool ──────────────────────────────────────────────────────────

/// Object-safe wrapper around a concrete [`Tool`].
///
/// You don't need to use this directly — [`ToolSet::add`] handles it.
trait ErasedTool {
    fn definition(&self) -> ToolDefinition;

    fn call<'a>(
        &'a self,
        args: serde_json::Value,
    ) -> BoxFuture<'a, Result<serde_json::Value, ToolError>>;
}

/// Concrete wrapper that bridges `Tool` to `ErasedTool`.
struct ToolWrapper<T>(T);

impl<T: Tool + 'static> ErasedTool for ToolWrapper<T> {
    fn definition(&self) -> ToolDefinition {
        self.0.definition()
    }

    fn call<'a>(
        &'a self,
        raw: serde_json::Value,
    ) -> BoxFuture<'a, Result<serde_json::Value, ToolError>> {
        Box::pin(async move {
            let args: T::Args =
                serde_json::from_value(raw).map_err(ToolError::ArgsParse)?;
            let output = self
                .0
                .call(args)
                .await
                .map_err(|e| ToolError::Execution(e.to_string()))?;
            serde_json::to_value(output).map_err(ToolError::ArgsParse)
        })
    }
}

// ── ToolSet ───────────────────────────────────────────────────────────────────

/// A collection of tools keyed by name.
///
/// Register tools with [`add`](ToolSet::add), then pass the set to
/// [`AgentBuilder::tool`](crate::agent::AgentBuilder::tool).
type BoxedTool = Box<dyn ErasedTool>;

pub struct ToolSet {
    tools: HashMap<String, BoxedTool>,
}

impl ToolSet {
    pub fn new() -> Self {
        Self { tools: HashMap::new() }
    }

    /// Register a tool. `T::NAME` is used as the lookup key.
    pub fn add<T: Tool + 'static>(&mut self, tool: T) -> &mut Self {
        self.tools.insert(T::NAME.to_owned(), Box::new(ToolWrapper(tool)));
        self
    }

    /// Schema list to include in each [`CompletionRequest`](crate::completion::CompletionRequest).
    pub fn definitions(&self) -> Vec<ToolDefinition> {
        self.tools.values().map(|t| t.definition()).collect()
    }

    /// Dispatch a tool call by name.
    pub async fn call(
        &self,
        name: &str,
        args: serde_json::Value,
    ) -> Result<serde_json::Value, ToolError> {
        let tool = self.tools.get(name).ok_or_else(|| ToolError::NotFound(name.to_owned()))?;
        tool.call(args).await
    }

    pub fn is_empty(&self) -> bool {
        self.tools.is_empty()
    }
}

impl Default for ToolSet {
    fn default() -> Self {
        Self::new()
    }
}