jules_core/tool/

mod.rs

//! Tool calling support and abstractions.

use crate::errors::ToolError;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;

/// A tool parameter schema.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolParameter {
    /// The parameter type (e.g., "string", "number", "boolean", "object").
    #[serde(rename = "type")]
    pub param_type: String,
    /// An optional description of the parameter.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
}

/// The parameters expected by a tool.
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct ToolParameters {
    /// The map of parameter names to their schemas.
    pub properties: HashMap<String, ToolParameter>,
    /// The list of required parameters.
    pub required: Vec<String>,
}

/// A generic trait for a tool that can be called by an LLM model.
pub trait Tool: Send + Sync {
    /// Returns the name of the tool.
    fn name(&self) -> &str;

    /// Returns a description of what the tool does.
    fn description(&self) -> &str;

    /// Returns the parameters expected by this tool.
    fn parameters(&self) -> ToolParameters;

    /// Calls the tool with the given arguments.
    /// The arguments are typically a JSON string generated by the LLM.
    fn call(&self, args: &str) -> impl Future<Output = Result<String, ToolError>> + Send;
}

/// A wrapper trait for tools that is object-safe (`dyn Tool`).
pub trait DynTool: Send + Sync {
    /// Returns the name of the tool.
    fn name(&self) -> &str;

    /// Returns a description of what the tool does.
    fn description(&self) -> &str;

    /// Returns the parameters expected by this tool.
    fn parameters(&self) -> ToolParameters;

    /// Calls the tool with the given arguments.
    fn call_dyn<'a>(
        &'a self,
        args: &'a str,
    ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>>;
}

impl<T: Tool> DynTool for T {
    fn name(&self) -> &str {
        Tool::name(self)
    }

    fn description(&self) -> &str {
        Tool::description(self)
    }

    fn parameters(&self) -> ToolParameters {
        Tool::parameters(self)
    }

    fn call_dyn<'a>(
        &'a self,
        args: &'a str,
    ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
        Box::pin(Tool::call(self, args))
    }
}

/// A registry that stores and manages tools.
pub struct ToolRegistry {
    tools: HashMap<String, Box<dyn DynTool>>,
}

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

impl ToolRegistry {
    /// Creates a new, empty `ToolRegistry`.
    #[must_use]
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    /// Registers a new tool in the registry.
    pub fn register<T: Tool + 'static>(&mut self, tool: T) {
        self.tools.insert(tool.name().to_string(), Box::new(tool));
    }

    /// Retrieves a tool by name from the registry, if it exists.
    #[must_use]
    pub fn get(&self, name: &str) -> Option<&dyn DynTool> {
        self.tools.get(name).map(std::convert::AsRef::as_ref)
    }

    /// Returns a list of all tools in the registry.
    #[must_use]
    pub fn list(&self) -> Vec<&dyn DynTool> {
        self.tools
            .values()
            .map(std::convert::AsRef::as_ref)
            .collect()
    }
}

/// Information about a requested tool call.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolCallInfo {
    /// The ID of the tool call.
    pub id: String,
    /// The name of the tool to call.
    pub name: String,
    /// The arguments to pass to the tool, usually as a JSON string.
    pub arguments: String,
}

/// A parsed tool call requested by the model.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolCall {
    /// Information about the tool call.
    pub function: ToolCallInfo,
    /// The type of tool call (e.g., "function").
    #[serde(rename = "type")]
    pub tool_type: String,
}

/// The result of a tool call to be returned to the conversation.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolResult {
    /// The ID of the original tool call.
    pub tool_call_id: String,
    /// The result of the tool call, usually as a string.
    pub content: String,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::errors::ToolError;

    struct EchoTool;

    impl Tool for EchoTool {
        fn name(&self) -> &'static str {
            "echo"
        }

        fn description(&self) -> &'static str {
            "Echoes the input."
        }

        fn parameters(&self) -> ToolParameters {
            let mut properties = HashMap::new();
            properties.insert(
                "input".to_string(),
                ToolParameter {
                    param_type: "string".to_string(),
                    description: Some("The text to echo".to_string()),
                },
            );
            ToolParameters {
                properties,
                required: vec!["input".to_string()],
            }
        }

        fn call(&self, args: &str) -> impl Future<Output = Result<String, ToolError>> + Send {
            let args_owned = args.to_string();
            async move { Ok(args_owned) }
        }
    }

    #[test]
    fn test_tool_registry() {
        let mut registry = ToolRegistry::new();
        registry.register(EchoTool);

        assert!(registry.get("echo").is_some());
        assert!(registry.get("missing").is_none());
        assert_eq!(registry.list().len(), 1);
    }

    #[tokio::test]
    async fn test_tool_execution() {
        let tool = EchoTool;
        let dyn_tool: &dyn DynTool = &tool;

        let result = dyn_tool.call_dyn("hello world").await.unwrap();
        assert_eq!(result, "hello world");
    }
}