cortexai_tools/

macro_tools.rs

//! Tools defined using the derive macro
//!
//! This module demonstrates how to create tools using the `#[derive(Tool)]` macro
//! for minimal boilerplate.

use cortexai_macros::Tool;

// --- Function-based tools using #[tool] attribute macro ---

#[cortexai_macros::tool(description = "Say hello to someone")]
async fn say_hello(
    #[param(description = "Name to greet", required)] name: String,
) -> Result<serde_json::Value, String> {
    Ok(serde_json::json!({ "greeting": format!("Hello, {}!", name) }))
}

#[cortexai_macros::tool(description = "Get weather for a city")]
async fn get_weather(
    #[param(description = "City name", required)] city: String,
    #[param(description = "Unit system")] unit: Option<String>,
) -> Result<serde_json::Value, String> {
    let u = unit.unwrap_or_else(|| "metric".to_string());
    Ok(serde_json::json!({ "city": city, "unit": u, "temp": 72 }))
}

#[cortexai_macros::tool(description = "Compute sum of numbers")]
async fn compute_sum(
    #[param(description = "First number", required)] a: i64,
    #[param(description = "Second number", required)] b: i64,
    #[param(description = "Include absolute value")] absolute: Option<bool>,
) -> Result<serde_json::Value, String> {
    let sum = a + b;
    let result = if absolute.unwrap_or(false) { sum.abs() } else { sum };
    Ok(serde_json::json!({ "result": result }))
}

#[cortexai_macros::tool(description = "Look up a word")]
async fn lookup_word(
    #[param(name = "query", description = "The word to look up", required)] word: String,
) -> Result<serde_json::Value, String> {
    Ok(serde_json::json!({ "word": word, "found": true }))
}

/// A simple echo tool that returns the input message
#[derive(Tool, Default)]
#[tool(name = "echo", description = "Echo back the input message")]
pub struct EchoTool {
    #[tool(param, required, description = "The message to echo back")]
    pub message: String,

    #[tool(param, description = "Number of times to repeat the message")]
    pub repeat: Option<u32>,
}

impl EchoTool {
    pub fn new() -> Self {
        Self::default()
    }

    /// The implementation method that the macro calls
    pub async fn run(
        &self,
        message: String,
        repeat: Option<u32>,
    ) -> Result<serde_json::Value, String> {
        let times = repeat.unwrap_or(1);
        let result: Vec<String> = (0..times).map(|_| message.clone()).collect();

        Ok(serde_json::json!({
            "echoed": if times == 1 { serde_json::json!(message) } else { serde_json::json!(result) },
            "repeat_count": times
        }))
    }
}

/// A greeting tool that generates personalized greetings
#[derive(Tool, Default)]
#[tool(
    name = "greet",
    description = "Generate a personalized greeting message"
)]
pub struct GreetTool {
    #[tool(param, required, description = "Name of the person to greet")]
    pub name: String,

    #[tool(param, description = "Language for the greeting (en, pt, es)")]
    pub language: Option<String>,

    #[tool(param, description = "Whether to use formal greeting")]
    pub formal: Option<bool>,
}

impl GreetTool {
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn run(
        &self,
        name: String,
        language: Option<String>,
        formal: Option<bool>,
    ) -> Result<serde_json::Value, String> {
        let lang = language.unwrap_or_else(|| "en".to_string());
        let is_formal = formal.unwrap_or(false);

        let greeting = match (lang.as_str(), is_formal) {
            ("pt", true) => format!("Prezado(a) {}, é um prazer cumprimentá-lo(a).", name),
            ("pt", false) => format!("Olá, {}! Tudo bem?", name),
            ("es", true) => format!("Estimado(a) {}, es un placer saludarle.", name),
            ("es", false) => format!("¡Hola, {}! ¿Qué tal?", name),
            (_, true) => format!("Dear {}, it is a pleasure to greet you.", name),
            (_, false) => format!("Hello, {}! How are you?", name),
        };

        Ok(serde_json::json!({
            "greeting": greeting,
            "language": lang,
            "formal": is_formal
        }))
    }
}

/// A text transformation tool
#[derive(Tool, Default)]
#[tool(
    name = "transform_text",
    description = "Transform text with various operations"
)]
pub struct TransformTextTool {
    #[tool(param, required, description = "The text to transform")]
    pub text: String,

    #[tool(
        param,
        required,
        description = "Operation: uppercase, lowercase, reverse, title"
    )]
    pub operation: String,
}

impl TransformTextTool {
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn run(&self, text: String, operation: String) -> Result<serde_json::Value, String> {
        let result = match operation.to_lowercase().as_str() {
            "uppercase" => text.to_uppercase(),
            "lowercase" => text.to_lowercase(),
            "reverse" => text.chars().rev().collect(),
            "title" => text
                .split_whitespace()
                .map(|word| {
                    let mut chars = word.chars();
                    match chars.next() {
                        None => String::new(),
                        Some(first) => first
                            .to_uppercase()
                            .chain(chars.flat_map(|c| c.to_lowercase()))
                            .collect(),
                    }
                })
                .collect::<Vec<_>>()
                .join(" "),
            _ => return Err(format!("Unknown operation: {}", operation)),
        };

        Ok(serde_json::json!({
            "original": text,
            "transformed": result,
            "operation": operation
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cortexai_core::tool::{ExecutionContext, Tool};
    use cortexai_core::types::AgentId;
    use std::collections::HashMap;

    /// Test tool with Vec<String> parameter
    #[derive(Tool, Default)]
    #[tool(name = "tag_tool", description = "A tool with vector parameter")]
    struct TagTool {
        #[tool(param, required, description = "List of tags")]
        tags: Vec<String>,
    }

    impl TagTool {
        async fn run(&self, tags: Vec<String>) -> Result<serde_json::Value, String> {
            Ok(serde_json::json!({ "tags": tags }))
        }
    }

    /// Test tool with Option<Vec<u32>> parameter
    #[derive(Tool, Default)]
    #[tool(name = "score_tool", description = "A tool with optional vec parameter")]
    struct ScoreTool {
        #[tool(param, required, description = "Name of scorer")]
        name: String,

        #[tool(param, description = "Optional list of scores")]
        scores: Option<Vec<u32>>,
    }

    impl ScoreTool {
        async fn run(&self, name: String, scores: Option<Vec<u32>>) -> Result<serde_json::Value, String> {
            Ok(serde_json::json!({ "name": name, "scores": scores }))
        }
    }

    /// Test tool with HashMap<String, String> parameter
    #[derive(Tool, Default)]
    #[tool(name = "metadata_tool", description = "A tool with hashmap parameter")]
    struct MetadataTool {
        #[tool(param, required, description = "Key-value metadata")]
        metadata: HashMap<String, String>,
    }

    impl MetadataTool {
        async fn run(&self, metadata: HashMap<String, String>) -> Result<serde_json::Value, String> {
            Ok(serde_json::json!({ "metadata": metadata }))
        }
    }

    fn ctx() -> ExecutionContext {
        ExecutionContext::new(AgentId::new("test"))
    }

    #[test]
    fn test_echo_tool_schema() {
        let tool = EchoTool::new();
        let schema = tool.schema();

        assert_eq!(schema.name, "echo");
        assert!(schema.description.contains("Echo"));
        assert!(schema.parameters["properties"]["message"].is_object());
        assert!(schema.parameters["required"]
            .as_array()
            .unwrap()
            .contains(&serde_json::json!("message")));
    }

    #[tokio::test]
    async fn test_echo_tool_execute() {
        let tool = EchoTool::new();

        let result = tool
            .execute(
                &ctx(),
                serde_json::json!({
                    "message": "Hello, World!"
                }),
            )
            .await
            .unwrap();

        assert_eq!(result["echoed"], "Hello, World!");
        assert_eq!(result["repeat_count"], 1);
    }

    #[tokio::test]
    async fn test_echo_tool_repeat() {
        let tool = EchoTool::new();

        let result = tool
            .execute(
                &ctx(),
                serde_json::json!({
                    "message": "Hi",
                    "repeat": 3
                }),
            )
            .await
            .unwrap();

        assert_eq!(result["echoed"].as_array().unwrap().len(), 3);
        assert_eq!(result["repeat_count"], 3);
    }

    #[test]
    fn test_greet_tool_schema() {
        let tool = GreetTool::new();
        let schema = tool.schema();

        assert_eq!(schema.name, "greet");
        assert!(schema.parameters["properties"]["name"].is_object());
        assert!(schema.parameters["properties"]["language"].is_object());
        assert!(schema.parameters["properties"]["formal"].is_object());
    }

    #[tokio::test]
    async fn test_greet_tool_execute() {
        let tool = GreetTool::new();

        let result = tool
            .execute(
                &ctx(),
                serde_json::json!({
                    "name": "Lucas",
                    "language": "pt",
                    "formal": false
                }),
            )
            .await
            .unwrap();

        assert!(result["greeting"].as_str().unwrap().contains("Olá"));
        assert!(result["greeting"].as_str().unwrap().contains("Lucas"));
    }

    #[tokio::test]
    async fn test_transform_tool_uppercase() {
        let tool = TransformTextTool::new();

        let result = tool
            .execute(
                &ctx(),
                serde_json::json!({
                    "text": "hello world",
                    "operation": "uppercase"
                }),
            )
            .await
            .unwrap();

        assert_eq!(result["transformed"], "HELLO WORLD");
    }

    #[tokio::test]
    async fn test_transform_tool_reverse() {
        let tool = TransformTextTool::new();

        let result = tool
            .execute(
                &ctx(),
                serde_json::json!({
                    "text": "abc",
                    "operation": "reverse"
                }),
            )
            .await
            .unwrap();

        assert_eq!(result["transformed"], "cba");
    }

    // ---- Function-based #[tool] macro tests ----

    #[test]
    fn test_fn_tool_struct_name_and_schema() {
        let tool = SayHelloTool::default();
        let schema = tool.schema();

        assert_eq!(schema.name, "say_hello");
        assert_eq!(schema.description, "Say hello to someone");
        assert!(schema.parameters["properties"]["name"].is_object());
        assert_eq!(
            schema.parameters["properties"]["name"]["description"],
            "Name to greet"
        );
        assert!(schema.parameters["required"]
            .as_array()
            .unwrap()
            .contains(&serde_json::json!("name")));
    }

    #[tokio::test]
    async fn test_fn_tool_execute_basic() {
        let tool = SayHelloTool::default();

        let result = tool
            .execute(
                &ctx(),
                serde_json::json!({ "name": "World" }),
            )
            .await
            .unwrap();

        assert_eq!(result["greeting"], "Hello, World!");
    }

    #[test]
    fn test_fn_tool_option_param_is_optional_in_schema() {
        let tool = GetWeatherTool::default();
        let schema = tool.schema();

        assert_eq!(schema.name, "get_weather");
        assert_eq!(schema.description, "Get weather for a city");

        // city is required, unit is not
        let required = schema.parameters["required"].as_array().unwrap();
        assert!(required.contains(&serde_json::json!("city")));
        assert!(!required.contains(&serde_json::json!("unit")));

        // Both appear in properties
        assert!(schema.parameters["properties"]["city"].is_object());
        assert!(schema.parameters["properties"]["unit"].is_object());
    }

    #[tokio::test]
    async fn test_fn_tool_option_param_execute_without_optional() {
        let tool = GetWeatherTool::default();

        let result = tool
            .execute(&ctx(), serde_json::json!({ "city": "London" }))
            .await
            .unwrap();

        assert_eq!(result["city"], "London");
        assert_eq!(result["unit"], "metric"); // default
    }

    #[tokio::test]
    async fn test_fn_tool_option_param_execute_with_optional() {
        let tool = GetWeatherTool::default();

        let result = tool
            .execute(&ctx(), serde_json::json!({ "city": "NYC", "unit": "imperial" }))
            .await
            .unwrap();

        assert_eq!(result["city"], "NYC");
        assert_eq!(result["unit"], "imperial");
    }

    #[test]
    fn test_fn_tool_multiple_params_schema() {
        let tool = ComputeSumTool::default();
        let schema = tool.schema();

        assert_eq!(schema.name, "compute_sum");

        let required = schema.parameters["required"].as_array().unwrap();
        assert!(required.contains(&serde_json::json!("a")));
        assert!(required.contains(&serde_json::json!("b")));
        assert!(!required.contains(&serde_json::json!("absolute")));

        assert_eq!(schema.parameters["properties"]["a"]["type"], "integer");
        assert_eq!(schema.parameters["properties"]["b"]["type"], "integer");
        assert_eq!(schema.parameters["properties"]["absolute"]["type"], "boolean");
    }

    #[tokio::test]
    async fn test_fn_tool_multiple_params_execute() {
        let tool = ComputeSumTool::default();

        let result = tool
            .execute(&ctx(), serde_json::json!({ "a": 3, "b": 5 }))
            .await
            .unwrap();

        assert_eq!(result["result"], 8);
    }

    #[tokio::test]
    async fn test_fn_tool_multiple_params_with_optional() {
        let tool = ComputeSumTool::default();

        let result = tool
            .execute(&ctx(), serde_json::json!({ "a": -3, "b": 1, "absolute": true }))
            .await
            .unwrap();

        assert_eq!(result["result"], 2);
    }

    #[test]
    fn test_fn_tool_param_name_override() {
        let tool = LookupWordTool::default();
        let schema = tool.schema();

        // The schema should use "query" not "word"
        assert!(schema.parameters["properties"]["query"].is_object());
        assert!(schema.parameters["properties"]["word"].is_null());
        assert!(schema.parameters["required"]
            .as_array()
            .unwrap()
            .contains(&serde_json::json!("query")));
    }

    #[tokio::test]
    async fn test_fn_tool_param_name_override_execute() {
        let tool = LookupWordTool::default();

        let result = tool
            .execute(&ctx(), serde_json::json!({ "query": "hello" }))
            .await
            .unwrap();

        assert_eq!(result["word"], "hello");
    }

    #[tokio::test]
    async fn test_fn_tool_missing_required_param() {
        let tool = GetWeatherTool::default();

        // city is required but not provided
        let result = tool.execute(&ctx(), serde_json::json!({})).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_missing_required_param() {
        let tool = EchoTool::new();

        let result = tool.execute(&ctx(), serde_json::json!({})).await;

        assert!(result.is_err());
    }

    #[test]
    fn test_vec_string_param_generates_array_schema() {
        let tool = TagTool::default();
        let schema = tool.schema();
        let tags_prop = &schema.parameters["properties"]["tags"];

        assert_eq!(tags_prop["type"], "array");
        assert_eq!(tags_prop["items"]["type"], "string");
    }

    #[test]
    fn test_option_vec_u32_param_is_optional_array() {
        let tool = ScoreTool::default();
        let schema = tool.schema();
        let scores_prop = &schema.parameters["properties"]["scores"];

        // Should be array type with integer items
        assert_eq!(scores_prop["type"], "array");
        assert_eq!(scores_prop["items"]["type"], "integer");

        // Should NOT be in the required list
        let required = schema.parameters["required"].as_array().unwrap();
        assert!(
            !required.contains(&serde_json::json!("scores")),
            "Option<Vec<u32>> should not be required"
        );
    }

    #[test]
    fn test_hashmap_param_generates_object_schema() {
        let tool = MetadataTool::default();
        let schema = tool.schema();
        let meta_prop = &schema.parameters["properties"]["metadata"];

        assert_eq!(meta_prop["type"], "object");
        assert_eq!(meta_prop["additionalProperties"]["type"], "string");
    }

    #[test]
    fn test_generated_schema_is_valid_json() {
        let tool = ScoreTool::default();
        let schema = tool.schema();

        // The parameters should be a valid JSON object
        assert!(schema.parameters.is_object());
        assert!(schema.parameters["properties"].is_object());
        assert!(schema.parameters["required"].is_array());
    }
}