xai_rust/models/

tool.rs

//! Tool definitions for function calling and server-side tools.

use serde::{Deserialize, Serialize};

/// A tool that can be used by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Tool {
    /// A custom function tool.
    Function {
        /// The function definition.
        function: FunctionDefinition,
    },
    /// Web search tool.
    WebSearch {
        /// Search filters.
        #[serde(skip_serializing_if = "Option::is_none", flatten)]
        filters: Option<WebSearchFilters>,
        /// Enable image understanding during search.
        #[serde(skip_serializing_if = "Option::is_none")]
        enable_image_understanding: Option<bool>,
    },
    /// X (Twitter) search tool.
    XSearch {
        /// Allowed X handles.
        #[serde(skip_serializing_if = "Option::is_none")]
        allowed_x_handles: Option<Vec<String>>,
        /// Excluded X handles.
        #[serde(skip_serializing_if = "Option::is_none")]
        excluded_x_handles: Option<Vec<String>>,
        /// Start date for search (YYYY-MM-DD).
        #[serde(skip_serializing_if = "Option::is_none")]
        from_date: Option<String>,
        /// End date for search (YYYY-MM-DD).
        #[serde(skip_serializing_if = "Option::is_none")]
        to_date: Option<String>,
        /// Enable image understanding.
        #[serde(skip_serializing_if = "Option::is_none")]
        enable_image_understanding: Option<bool>,
        /// Enable video understanding.
        #[serde(skip_serializing_if = "Option::is_none")]
        enable_video_understanding: Option<bool>,
    },
    /// Code interpreter tool.
    CodeInterpreter {},
    /// Collections search tool.
    CollectionsSearch {
        /// Collection IDs to search.
        #[serde(skip_serializing_if = "Option::is_none")]
        collection_ids: Option<Vec<String>>,
    },
    /// Remote MCP server tool.
    Mcp {
        /// MCP server configuration.
        server: McpServer,
        /// Specific tools to enable.
        #[serde(skip_serializing_if = "Option::is_none")]
        allowed_tools: Option<Vec<String>>,
    },
}

impl Tool {
    /// Create a function tool.
    pub fn function(
        name: impl Into<String>,
        description: impl Into<String>,
        parameters: serde_json::Value,
    ) -> Self {
        Self::Function {
            function: FunctionDefinition {
                name: name.into(),
                description: Some(description.into()),
                parameters,
                strict: None,
            },
        }
    }

    /// Create a web search tool with default settings.
    pub fn web_search() -> Self {
        Self::WebSearch {
            filters: None,
            enable_image_understanding: None,
        }
    }

    /// Create a web search tool with filters.
    pub fn web_search_filtered(filters: WebSearchFilters) -> Self {
        Self::WebSearch {
            filters: Some(filters),
            enable_image_understanding: None,
        }
    }

    /// Create an X search tool with default settings.
    pub fn x_search() -> Self {
        Self::XSearch {
            allowed_x_handles: None,
            excluded_x_handles: None,
            from_date: None,
            to_date: None,
            enable_image_understanding: None,
            enable_video_understanding: None,
        }
    }

    /// Create a code interpreter tool.
    pub fn code_interpreter() -> Self {
        Self::CodeInterpreter {}
    }

    /// Create a collections search tool.
    pub fn collections_search(collection_ids: Vec<String>) -> Self {
        Self::CollectionsSearch {
            collection_ids: Some(collection_ids),
        }
    }

    /// Create an MCP tool.
    pub fn mcp(server_url: impl Into<String>) -> Self {
        Self::Mcp {
            server: McpServer {
                url: server_url.into(),
                headers: None,
            },
            allowed_tools: None,
        }
    }
}

/// Function definition for custom tools.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionDefinition {
    /// The name of the function.
    pub name: String,
    /// Description of what the function does.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    /// JSON Schema for the function parameters.
    pub parameters: serde_json::Value,
    /// Whether to strictly enforce the schema.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub strict: Option<bool>,
}

impl FunctionDefinition {
    /// Create a new function definition.
    pub fn new(name: impl Into<String>, parameters: serde_json::Value) -> Self {
        Self {
            name: name.into(),
            description: None,
            parameters,
            strict: None,
        }
    }

    /// Set the description.
    pub fn with_description(mut self, description: impl Into<String>) -> Self {
        self.description = Some(description.into());
        self
    }

    /// Enable strict mode.
    pub fn strict(mut self) -> Self {
        self.strict = Some(true);
        self
    }
}

/// Web search filters.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct WebSearchFilters {
    /// Only search these domains.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub allowed_domains: Option<Vec<String>>,
    /// Exclude these domains from search.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub excluded_domains: Option<Vec<String>>,
}

impl WebSearchFilters {
    /// Create filters with allowed domains.
    pub fn allow_domains(domains: Vec<String>) -> Self {
        Self {
            allowed_domains: Some(domains),
            excluded_domains: None,
        }
    }

    /// Create filters with excluded domains.
    pub fn exclude_domains(domains: Vec<String>) -> Self {
        Self {
            allowed_domains: None,
            excluded_domains: Some(domains),
        }
    }
}

/// MCP server configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpServer {
    /// The URL of the MCP server.
    pub url: String,
    /// Custom headers for authentication.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub headers: Option<std::collections::HashMap<String, String>>,
}

/// MCP tool configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpTool {
    /// MCP server configuration.
    pub server: McpServer,
    /// Allowed tools from the server.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub allowed_tools: Option<Vec<String>>,
}

/// A tool call made by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    /// Unique ID for this tool call.
    pub id: String,
    /// Type of tool call (e.g., `"function"`).
    #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
    pub call_type: Option<String>,
    /// Function call details.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub function: Option<FunctionCall>,
}

/// Function call details.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionCall {
    /// Name of the function to call.
    pub name: String,
    /// Arguments as a JSON string.
    pub arguments: String,
}

impl FunctionCall {
    /// Parse the arguments as JSON.
    pub fn parse_arguments<T: serde::de::DeserializeOwned>(&self) -> Result<T, serde_json::Error> {
        serde_json::from_str(&self.arguments)
    }
}

/// Tool choice specification.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ToolChoice {
    /// Let the model decide.
    Auto(ToolChoiceAuto),
    /// Force a specific tool.
    Specific(ToolChoiceSpecific),
}

/// Auto tool choice modes.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ToolChoiceAuto {
    /// Model decides whether to use tools.
    Auto,
    /// Model must use at least one tool.
    Required,
    /// Model should not use tools.
    None,
}

/// Specific tool choice.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolChoiceSpecific {
    /// Type of tool.
    #[serde(rename = "type")]
    pub tool_type: String,
    /// Function specification (for function tools).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub function: Option<ToolChoiceFunction>,
}

/// Function specification for tool choice.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolChoiceFunction {
    /// Name of the function to use.
    pub name: String,
}

impl ToolChoice {
    /// Let the model decide.
    pub fn auto() -> Self {
        Self::Auto(ToolChoiceAuto::Auto)
    }

    /// Require tool use.
    pub fn required() -> Self {
        Self::Auto(ToolChoiceAuto::Required)
    }

    /// Disable tools.
    pub fn none() -> Self {
        Self::Auto(ToolChoiceAuto::None)
    }

    /// Force a specific function.
    pub fn function(name: impl Into<String>) -> Self {
        Self::Specific(ToolChoiceSpecific {
            tool_type: "function".to_string(),
            function: Some(ToolChoiceFunction { name: name.into() }),
        })
    }
}

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

    // ── Tool enum serde roundtrips ─────────────────────────────────────

    #[test]
    fn tool_function_roundtrip() {
        let t = Tool::function(
            "get_weather",
            "Get weather for a location",
            json!({"type": "object", "properties": {"city": {"type": "string"}}}),
        );
        let json_val = serde_json::to_value(&t).unwrap();
        assert_eq!(json_val["type"], "function");
        assert_eq!(json_val["function"]["name"], "get_weather");
        assert_eq!(
            json_val["function"]["description"],
            "Get weather for a location"
        );

        let back: Tool = serde_json::from_value(json_val).unwrap();
        if let Tool::Function { function } = back {
            assert_eq!(function.name, "get_weather");
        } else {
            panic!("Expected Function variant");
        }
    }

    #[test]
    fn tool_web_search_roundtrip() {
        let t = Tool::web_search();
        let json_val = serde_json::to_value(&t).unwrap();
        assert_eq!(json_val["type"], "web_search");

        let back: Tool = serde_json::from_value(json_val).unwrap();
        assert!(matches!(back, Tool::WebSearch { .. }));
    }

    #[test]
    fn tool_web_search_filtered_roundtrip() {
        let t =
            Tool::web_search_filtered(WebSearchFilters::allow_domains(vec!["docs.rs".to_string()]));
        let json_val = serde_json::to_value(&t).unwrap();
        assert_eq!(json_val["type"], "web_search");
        assert_eq!(json_val["allowed_domains"], json!(["docs.rs"]));

        let back: Tool = serde_json::from_value(json_val).unwrap();
        if let Tool::WebSearch { filters, .. } = back {
            let filters = filters.unwrap();
            assert_eq!(filters.allowed_domains.unwrap(), vec!["docs.rs"]);
        } else {
            panic!("Expected WebSearch variant");
        }
    }

    #[test]
    fn tool_x_search_roundtrip() {
        let t = Tool::x_search();
        let json_val = serde_json::to_value(&t).unwrap();
        assert_eq!(json_val["type"], "x_search");

        let back: Tool = serde_json::from_value(json_val).unwrap();
        assert!(matches!(back, Tool::XSearch { .. }));
    }

    #[test]
    fn tool_code_interpreter_roundtrip() {
        let t = Tool::code_interpreter();
        let json_val = serde_json::to_value(&t).unwrap();
        assert_eq!(json_val["type"], "code_interpreter");

        let back: Tool = serde_json::from_value(json_val).unwrap();
        assert!(matches!(back, Tool::CodeInterpreter {}));
    }

    #[test]
    fn tool_collections_search_roundtrip() {
        let t = Tool::collections_search(vec!["col-1".to_string(), "col-2".to_string()]);
        let json_val = serde_json::to_value(&t).unwrap();
        assert_eq!(json_val["type"], "collections_search");
        assert_eq!(json_val["collection_ids"], json!(["col-1", "col-2"]));

        let back: Tool = serde_json::from_value(json_val).unwrap();
        if let Tool::CollectionsSearch { collection_ids } = back {
            assert_eq!(collection_ids.unwrap(), vec!["col-1", "col-2"]);
        } else {
            panic!("Expected CollectionsSearch variant");
        }
    }

    #[test]
    fn tool_mcp_roundtrip() {
        let t = Tool::mcp("https://mcp.example.com");
        let json_val = serde_json::to_value(&t).unwrap();
        assert_eq!(json_val["type"], "mcp");
        assert_eq!(json_val["server"]["url"], "https://mcp.example.com");

        let back: Tool = serde_json::from_value(json_val).unwrap();
        if let Tool::Mcp { server, .. } = back {
            assert_eq!(server.url, "https://mcp.example.com");
        } else {
            panic!("Expected Mcp variant");
        }
    }

    // ── ToolCall serde ─────────────────────────────────────────────────

    #[test]
    fn tool_call_with_call_type_roundtrip() {
        let tc = ToolCall {
            id: "call_abc".to_string(),
            call_type: Some("function".to_string()),
            function: Some(FunctionCall {
                name: "get_weather".to_string(),
                arguments: r#"{"city":"NYC"}"#.to_string(),
            }),
        };
        let json_val = serde_json::to_value(&tc).unwrap();
        assert_eq!(json_val["id"], "call_abc");
        assert_eq!(json_val["type"], "function");
        assert_eq!(json_val["function"]["name"], "get_weather");

        let back: ToolCall = serde_json::from_value(json_val).unwrap();
        assert_eq!(back.id, "call_abc");
        assert_eq!(back.call_type.as_deref(), Some("function"));
        assert_eq!(back.function.as_ref().unwrap().name, "get_weather");
    }

    #[test]
    fn tool_call_without_call_type_roundtrip() {
        // call_type is Option<String>; verify it defaults to None when absent
        let json_val = json!({
            "id": "call_xyz",
            "function": {
                "name": "do_stuff",
                "arguments": "{}"
            }
        });
        let tc: ToolCall = serde_json::from_value(json_val).unwrap();
        assert_eq!(tc.id, "call_xyz");
        assert!(tc.call_type.is_none());
    }

    #[test]
    fn tool_call_none_type_skipped_on_serialize() {
        let tc = ToolCall {
            id: "call_1".to_string(),
            call_type: None,
            function: None,
        };
        let json_val = serde_json::to_value(&tc).unwrap();
        // "type" key should be absent because of skip_serializing_if
        assert!(json_val.get("type").is_none());
    }

    // ── FunctionCall parse_arguments ───────────────────────────────────

    #[test]
    fn function_call_parse_arguments() {
        let fc = FunctionCall {
            name: "test".to_string(),
            arguments: r#"{"x": 42}"#.to_string(),
        };
        let parsed: serde_json::Value = fc.parse_arguments().unwrap();
        assert_eq!(parsed["x"], 42);
    }

    #[test]
    fn function_call_parse_arguments_error() {
        let fc = FunctionCall {
            name: "test".to_string(),
            arguments: "not json".to_string(),
        };
        assert!(fc.parse_arguments::<serde_json::Value>().is_err());
    }

    // ── ToolChoice serde roundtrips ────────────────────────────────────

    #[test]
    fn tool_choice_auto_roundtrip() {
        let choice = ToolChoice::auto();
        let json_val = serde_json::to_value(&choice).unwrap();
        assert_eq!(json_val, json!("auto"));
    }

    #[test]
    fn tool_choice_required_roundtrip() {
        let choice = ToolChoice::required();
        let json_val = serde_json::to_value(&choice).unwrap();
        assert_eq!(json_val, json!("required"));
    }

    #[test]
    fn tool_choice_none_roundtrip() {
        let choice = ToolChoice::none();
        let json_val = serde_json::to_value(&choice).unwrap();
        assert_eq!(json_val, json!("none"));
    }

    #[test]
    fn tool_choice_specific_function_roundtrip() {
        let choice = ToolChoice::function("get_weather");
        let json_val = serde_json::to_value(&choice).unwrap();
        assert_eq!(json_val["type"], "function");
        assert_eq!(json_val["function"]["name"], "get_weather");

        let back: ToolChoice = serde_json::from_value(json_val).unwrap();
        if let ToolChoice::Specific(spec) = back {
            assert_eq!(spec.tool_type, "function");
            assert_eq!(spec.function.unwrap().name, "get_weather");
        } else {
            panic!("Expected Specific variant");
        }
    }

    // ── FunctionDefinition builder ─────────────────────────────────────

    #[test]
    fn function_definition_builder() {
        let fd = FunctionDefinition::new("test", json!({}))
            .with_description("A test function")
            .strict();
        assert_eq!(fd.name, "test");
        assert_eq!(fd.description.as_deref(), Some("A test function"));
        assert_eq!(fd.strict, Some(true));
    }

    #[test]
    fn function_definition_roundtrip() {
        let fd = FunctionDefinition {
            name: "search".to_string(),
            description: Some("Search the web".to_string()),
            parameters: json!({"type": "object"}),
            strict: Some(true),
        };
        let json_val = serde_json::to_value(&fd).unwrap();
        let back: FunctionDefinition = serde_json::from_value(json_val).unwrap();
        assert_eq!(back.name, "search");
        assert_eq!(back.strict, Some(true));
    }
}