mistralai-client-rs/src/v1/tool.rs

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::{any::Any, fmt::Debug};

// -----------------------------------------------------------------------------
// Definitions

#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub struct ToolCall {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub r#type: Option<String>,
    pub function: ToolCallFunction,
}

#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub struct ToolCallFunction {
    pub name: String,
    pub arguments: String,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Tool {
    pub r#type: ToolType,
    pub function: ToolFunction,
}
impl Tool {
    /// Create a tool with a JSON Schema parameters object.
    pub fn new(
        function_name: String,
        function_description: String,
        parameters: serde_json::Value,
    ) -> Self {
        Self {
            r#type: ToolType::Function,
            function: ToolFunction {
                name: function_name,
                description: function_description,
                parameters,
            },
        }
    }
}

// -----------------------------------------------------------------------------
// Request

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ToolFunction {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub enum ToolType {
    #[serde(rename = "function")]
    Function,
}

/// An enum representing how functions should be called.
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub enum ToolChoice {
    /// The model is forced to call a function.
    #[serde(rename = "any")]
    Any,
    /// The model can choose to either generate a message or call a function.
    #[serde(rename = "auto")]
    Auto,
    /// The model won't call a function and will generate a message instead.
    #[serde(rename = "none")]
    None,
    /// The model must call at least one tool.
    #[serde(rename = "required")]
    Required,
}

// -----------------------------------------------------------------------------
// Custom

#[async_trait]
pub trait Function: Send {
    async fn execute(&self, arguments: String) -> Box<dyn Any + Send>;
}

impl Debug for dyn Function {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Function()")
    }
}