openfunctions-rs 0.1.0

//! Data models for tools, agents, and execution results.
//!
//! This module defines the primary data structures used throughout the application.
//! These models are designed to be serializable and deserializable, and they form
//! the basis for configuration, tool definition, and communication between
//! components.

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// The definition of a tool, parsed from comments in its source code.
///
/// This structure captures everything needed to understand and use a tool,
/// including its purpose, parameters, and dependencies.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct ToolDefinition {
    /// A description of the tool's purpose and functionality.
    pub description: String,

    /// The parameters that the tool accepts.
    pub parameters: Vec<ParameterDefinition>,

    /// Environment variables that are required for the tool's execution.
    pub env_vars: Vec<EnvVarDefinition>,

    /// A list of external command-line tools that this tool depends on.
    pub required_tools: Vec<String>,

    /// A map for any additional, unstructured metadata about the tool.
    pub metadata: HashMap<String, serde_json::Value>,
}

impl ToolDefinition {
    /// Converts a `ToolDefinition` into a `FunctionDeclaration` for use with an LLM.
    ///
    /// This is a crucial step for function calling, as it transforms the internal
    /// tool representation into a format that the AI model can understand.
    pub fn to_function_declaration(
        &self,
        name: &str,
    ) -> crate::Result<crate::core::function::FunctionDeclaration> {
        use crate::core::function::{FunctionDeclaration, Parameter};

        let parameters = self
            .parameters
            .iter()
            .map(|p| {
                let param_type = match &p.param_type {
                    ParameterType::String => {
                        if let Some(ref enum_values) = p.enum_values {
                            crate::core::function::ParameterType::Enum(enum_values.clone())
                        } else {
                            crate::core::function::ParameterType::String
                        }
                    }
                    ParameterType::Integer => crate::core::function::ParameterType::Integer,
                    ParameterType::Number => crate::core::function::ParameterType::Number,
                    ParameterType::Boolean => crate::core::function::ParameterType::Boolean,
                    ParameterType::Array => crate::core::function::ParameterType::Array,
                    ParameterType::Object => crate::core::function::ParameterType::Object,
                    ParameterType::Enum(values) => {
                        crate::core::function::ParameterType::Enum(values.clone())
                    }
                };

                Parameter {
                    name: p.name.clone(),
                    param_type,
                    description: p.description.clone(),
                    required: p.required,
                    default: p.default.clone(),
                }
            })
            .collect();

        Ok(FunctionDeclaration {
            name: name.to_string(),
            description: self.description.clone(),
            parameters,
        })
    }
}

/// The definition of a single parameter for a tool.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct ParameterDefinition {
    /// The name of the parameter.
    pub name: String,

    /// The data type of the parameter.
    pub param_type: ParameterType,

    /// A description of the parameter's purpose.
    pub description: String,

    /// Whether the parameter is required for the tool to be called.
    pub required: bool,

    /// An optional default value for the parameter.
    pub default: Option<serde_json::Value>,

    /// For `String` or `Enum` types, a list of allowed values.
    pub enum_values: Option<Vec<String>>,
}

/// An enumeration of data types for tool parameters.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ParameterType {
    /// UTF-8 string type
    String,
    /// 64-bit signed integer
    Integer,
    /// 64-bit floating point number
    Number,
    /// Boolean value (true/false)
    Boolean,
    /// Array of values
    Array,
    /// JSON object
    Object,
    /// String with predefined allowed values
    #[serde(rename = "enum")]
    Enum(Vec<String>),
}

/// The definition of an environment variable required by a tool.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct EnvVarDefinition {
    /// The name of the environment variable (e.g., "API_KEY").
    pub name: String,

    /// A description of what the environment variable is used for.
    pub description: String,

    /// Whether the environment variable must be set for the tool to run.
    pub required: bool,

    /// An optional default value for the environment variable.
    pub default: Option<String>,
}

/// The definition of an agent, loaded from its `index.yaml` file.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct AgentDefinition {
    /// The unique name of the agent.
    pub name: String,

    /// A description of the agent's purpose and capabilities.
    pub description: String,

    /// The version of the agent.
    pub version: String,

    /// The system prompt or instructions that guide the agent's behavior.
    pub instructions: String,

    /// Variables that can be configured for the agent.
    #[serde(default)]
    pub variables: Vec<AgentVariable>,

    /// A list of documents to be used for Retrieval-Augmented Generation (RAG).
    #[serde(default)]
    pub documents: Vec<String>,

    /// A list of example prompts to start a conversation with the agent.
    #[serde(default)]
    pub conversation_starters: Vec<String>,

    /// Additional, unstructured metadata for the agent.
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,
}

/// The definition of a configurable variable for an agent.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct AgentVariable {
    /// The name of the variable.
    pub name: String,

    /// A description of the variable's purpose.
    pub description: String,

    /// An optional default value for the variable.
    pub default: Option<String>,
}

/// The result of an execution of a tool function.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionResult {
    /// `true` if the execution was successful, `false` otherwise.
    pub success: bool,

    /// The standard output (stdout) from the execution.
    pub output: String,

    /// The standard error (stderr) from the execution, if any.
    pub error: Option<String>,

    /// The duration of the execution.
    pub duration: std::time::Duration,

    /// Additional, unstructured metadata about the execution.
    pub metadata: HashMap<String, serde_json::Value>,
}

/// Represents a request to call a function, typically from an LLM.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionCall {
    /// The name of the function to call.
    pub name: String,

    /// The arguments for the function, as a JSON value.
    pub arguments: serde_json::Value,

    /// A unique ID for tracking the request.
    pub request_id: String,

    /// The timestamp of when the call was requested.
    pub timestamp: chrono::DateTime<chrono::Utc>,
}

/// Represents the response from a function call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionResponse {
    /// The ID of the request that this response corresponds to.
    pub request_id: String,

    /// The result of the function's execution.
    pub result: ExecutionResult,

    /// The timestamp of when the response was generated.
    pub timestamp: chrono::DateTime<chrono::Utc>,
}