aacode-rs 1.8.15

// Copyright (c) 2026 xiefujin <490021684@qq.com>
// Licensed under GPL-3.0, see LICENSE file for full license terms.

//! Tool schema definitions + conversion to OpenAI/Anthropic native tools.
//!
//! Ported from Python `utils/tool_schemas.py` + `utils/tool_adapter.py`.

use serde_json::{json, Value};

/// JSON-schema primitive type for a parameter.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParamType {
    String,
    Integer,
    Number,
    Boolean,
    Array,
    Object,
}

impl ParamType {
    pub fn json_name(&self) -> &'static str {
        match self {
            ParamType::String => "string",
            ParamType::Integer => "integer",
            ParamType::Number => "number",
            ParamType::Boolean => "boolean",
            ParamType::Array => "array",
            ParamType::Object => "object",
        }
    }

    /// Whether a serde_json Value matches this declared type (used for validation).
    pub fn matches(&self, v: &Value) -> bool {
        match self {
            ParamType::String => v.is_string(),
            ParamType::Integer => v.is_i64() || v.is_u64(),
            ParamType::Number => v.is_number(),
            ParamType::Boolean => v.is_boolean(),
            ParamType::Array => v.is_array(),
            ParamType::Object => v.is_object(),
        }
    }

    /// Best-effort coerce a value to this type. LLMs frequently send numbers /
    /// booleans as strings (e.g. `"max_output": "3000"`); non-coercible values
    /// are returned unchanged so validation still reports a clear type error.
    pub fn coerce(&self, v: &Value) -> Value {
        match (self, v) {
            (ParamType::Integer, Value::String(s)) => s
                .trim()
                .parse::<i64>()
                .map(|n| json!(n))
                .unwrap_or_else(|_| v.clone()),
            (ParamType::Number, Value::String(s)) => {
                if let Ok(n) = s.trim().parse::<i64>() {
                    json!(n)
                } else if let Ok(f) = s.trim().parse::<f64>() {
                    json!(f)
                } else {
                    v.clone()
                }
            }
            (ParamType::Boolean, Value::String(s)) => {
                match s.trim().to_ascii_lowercase().as_str() {
                    "true" | "1" | "yes" => json!(true),
                    "false" | "0" | "no" => json!(false),
                    _ => v.clone(),
                }
            }
            _ => v.clone(),
        }
    }
}

/// A single tool parameter definition.
#[derive(Debug, Clone)]
pub struct ToolParameter {
    pub name: &'static str,
    pub ty: ParamType,
    pub required: bool,
    pub description: &'static str,
    /// Alternative names that get normalized to `name`.
    pub aliases: &'static [&'static str],
}

impl ToolParameter {
    pub const fn new(
        name: &'static str,
        ty: ParamType,
        required: bool,
        description: &'static str,
        aliases: &'static [&'static str],
    ) -> Self {
        ToolParameter {
            name,
            ty,
            required,
            description,
            aliases,
        }
    }
}

/// A tool's schema: name, description, parameters.
#[derive(Debug, Clone)]
pub struct ToolSchema {
    pub name: &'static str,
    pub description: &'static str,
    pub parameters: Vec<ToolParameter>,
}

impl ToolSchema {
    pub fn new(
        name: &'static str,
        description: &'static str,
        parameters: Vec<ToolParameter>,
    ) -> Self {
        ToolSchema {
            name,
            description,
            parameters,
        }
    }

    /// Build the JSON-schema `properties` object and `required` array.
    fn json_schema(&self) -> (Value, Vec<Value>) {
        let mut props = serde_json::Map::new();
        let mut required = Vec::new();
        for p in &self.parameters {
            props.insert(
                p.name.to_string(),
                json!({"type": p.ty.json_name(), "description": p.description}),
            );
            if p.required {
                required.push(Value::String(p.name.to_string()));
            }
        }
        (Value::Object(props), required)
    }

    /// OpenAI native tool object.
    pub fn to_openai_tool(&self) -> Value {
        let (props, required) = self.json_schema();
        json!({
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": {
                    "type": "object",
                    "properties": props,
                    "required": required,
                }
            }
        })
    }

    /// Anthropic native tool object.
    pub fn to_anthropic_tool(&self) -> Value {
        let (props, required) = self.json_schema();
        json!({
            "name": self.name,
            "description": self.description,
            "input_schema": {
                "type": "object",
                "properties": props,
                "required": required,
            }
        })
    }

    /// Normalize input params: rewrite known aliases to canonical names.
    pub fn normalize_params(&self, input: &Value) -> Value {
        let obj = match input.as_object() {
            Some(o) => o,
            None => return input.clone(),
        };
        // alias -> canonical
        let mut alias_map: std::collections::HashMap<&str, &str> = std::collections::HashMap::new();
        for p in &self.parameters {
            alias_map.insert(p.name, p.name);
            for a in p.aliases {
                alias_map.insert(a, p.name);
            }
        }
        let mut out = serde_json::Map::new();
        for (k, v) in obj {
            let canonical = alias_map.get(k.as_str()).copied().unwrap_or(k.as_str());
            // Coerce string scalars to the declared numeric/boolean type.
            let coerced = self
                .parameters
                .iter()
                .find(|p| p.name == canonical)
                .map(|p| p.ty.coerce(v))
                .unwrap_or_else(|| v.clone());
            out.insert(canonical.to_string(), coerced);
        }
        Value::Object(out)
    }

    /// Validate (after normalization): required present + type match.
    /// Returns Ok(()) or Err(human message).
    pub fn validate(&self, normalized: &Value) -> std::result::Result<(), String> {
        let obj = normalized.as_object();
        // missing required
        let mut missing = Vec::new();
        for p in &self.parameters {
            if p.required {
                let present = obj.map(|o| o.contains_key(p.name)).unwrap_or(false);
                if !present {
                    missing.push(p.name);
                }
            }
        }
        if !missing.is_empty() {
            let details: Vec<String> = missing
                .iter()
                .map(|m| {
                    let p = self.parameters.iter().find(|p| &p.name == m).unwrap();
                    format!("  • {} ({}): {}", p.name, p.ty.json_name(), p.description)
                })
                .collect();
            return Err(format!(
                "❌ Missing required parameters: {}\n\n📋 Parameter details:\n{}",
                missing.join(", "),
                details.join("\n")
            ));
        }
        // type check present params
        if let Some(o) = obj {
            for p in &self.parameters {
                if let Some(v) = o.get(p.name) {
                    if !v.is_null() && !p.ty.matches(v) {
                        return Err(format!(
                            "Parameter '{}' expected type {}, got a different type",
                            p.name,
                            p.ty.json_name()
                        ));
                    }
                }
            }
        }
        Ok(())
    }

    /// Human-readable documentation for the tool (used in not-found/error hints).
    pub fn documentation(&self) -> String {
        let mut doc = format!(
            "## {}\n\n{}\n\n### Parameters\n",
            self.name, self.description
        );
        for p in &self.parameters {
            let req = if p.required { "required" } else { "optional" };
            let aliases = if p.aliases.is_empty() {
                String::new()
            } else {
                format!(" (aliases: {})", p.aliases.join(", "))
            };
            doc.push_str(&format!(
                "- {}{} ({}, {}): {}\n",
                p.name,
                aliases,
                p.ty.json_name(),
                req,
                p.description
            ));
        }
        doc
    }
}

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

    fn sample() -> ToolSchema {
        ToolSchema::new(
            "run_shell",
            "run a shell command",
            vec![
                ToolParameter::new(
                    "command",
                    ParamType::String,
                    true,
                    "the command",
                    &["cmd", "shell", "script", "exec"],
                ),
                ToolParameter::new(
                    "timeout",
                    ParamType::Integer,
                    false,
                    "seconds",
                    &["time_limit"],
                ),
            ],
        )
    }

    #[test]
    fn openai_tool_shape() {
        let t = sample().to_openai_tool();
        assert_eq!(t["type"], "function");
        assert_eq!(t["function"]["name"], "run_shell");
        assert_eq!(
            t["function"]["parameters"]["properties"]["command"]["type"],
            "string"
        );
        assert_eq!(t["function"]["parameters"]["required"][0], "command");
    }

    #[test]
    fn anthropic_tool_shape() {
        let t = sample().to_anthropic_tool();
        assert_eq!(t["name"], "run_shell");
        assert_eq!(
            t["input_schema"]["properties"]["timeout"]["type"],
            "integer"
        );
    }

    #[test]
    fn alias_normalization() {
        let s = sample();
        let n = s.normalize_params(&json!({"cmd": "ls", "time_limit": 5}));
        assert_eq!(n["command"], "ls");
        assert_eq!(n["timeout"], 5);
    }

    #[test]
    fn validate_missing_required() {
        let s = sample();
        let n = s.normalize_params(&json!({"timeout": 5}));
        let err = s.validate(&n).unwrap_err();
        assert!(err.contains("command"));
    }

    #[test]
    fn validate_type_error() {
        let s = sample();
        let n = s.normalize_params(&json!({"command": 123}));
        assert!(s.validate(&n).is_err());
    }

    #[test]
    fn validate_ok() {
        let s = sample();
        let n = s.normalize_params(&json!({"command": "ls", "timeout": 5}));
        assert!(s.validate(&n).is_ok());
    }

    #[test]
    fn normalize_coerces_string_scalars_to_declared_types() {
        let s = ToolSchema::new(
            "t",
            "d",
            vec![
                ToolParameter::new("n", ParamType::Integer, false, "", &[]),
                ToolParameter::new("f", ParamType::Number, false, "", &[]),
                ToolParameter::new("b", ParamType::Boolean, false, "", &[]),
            ],
        );
        let out = s.normalize_params(&json!({"n": "3000", "f": "1.5", "b": "true"}));
        assert_eq!(out["n"], json!(3000));
        assert_eq!(out["f"], json!(1.5));
        assert_eq!(out["b"], json!(true));
        assert!(s.validate(&out).is_ok(), "coerced params must validate");
        // Non-coercible strings are left as-is (validation then reports a clear error).
        let out = s.normalize_params(&json!({"n": "abc"}));
        assert_eq!(out["n"], json!("abc"));
        assert!(s.validate(&out).is_err());
    }

    #[test]
    fn coerce_handles_whitespace_and_word_booleans() {
        assert_eq!(ParamType::Integer.coerce(&json!(" 42 ")), json!(42));
        assert_eq!(ParamType::Number.coerce(&json!("42")), json!(42));
        assert_eq!(ParamType::Number.coerce(&json!("2.5")), json!(2.5));
        assert_eq!(ParamType::Boolean.coerce(&json!(" YES ")), json!(true));
        assert_eq!(ParamType::Boolean.coerce(&json!("no")), json!(false));
        assert_eq!(ParamType::Boolean.coerce(&json!("1")), json!(true));
        // Non-coercible values are returned unchanged.
        assert_eq!(ParamType::Boolean.coerce(&json!("maybe")), json!("maybe"));
        assert_eq!(ParamType::Integer.coerce(&json!(3.5)), json!(3.5));
    }

    #[test]
    fn normalize_passthrough_for_non_objects() {
        let s = sample();
        assert_eq!(s.normalize_params(&json!([1, 2])), json!([1, 2]));
        assert_eq!(s.normalize_params(&json!("x")), json!("x"));
        assert_eq!(s.normalize_params(&json!(null)), json!(null));
    }

    #[test]
    fn validate_null_optional_is_ok() {
        let s = sample();
        let n = s.normalize_params(&json!({"command": "ls", "timeout": null}));
        assert!(s.validate(&n).is_ok());
    }

    #[test]
    fn documentation_lists_params_and_aliases() {
        let doc = sample().documentation();
        assert!(doc.contains("run_shell"));
        assert!(doc.contains("command"));
        assert!(doc.contains("aliases: cmd"));
        assert!(doc.contains("required"));
    }
}