openfunctions-rs 0.1.0

//! Agent management and representation.
//!
//! This module defines the `Agent` struct, which represents an AI agent,
//! its configuration, and its associated tools.

use crate::core::Tool;
use crate::models::AgentDefinition;
use anyhow::Result;
use std::collections::HashMap;
use std::path::{Path, PathBuf};

/// Represents an AI agent, including its definition, tools, and configuration.
///
/// An agent is a fundamental concept in this framework, encapsulating the logic
/// and capabilities needed to perform tasks.
#[derive(Debug, Clone)]
pub struct Agent {
    /// A unique name for the agent.
    pub name: String,

    /// The file system path to the agent's directory.
    pub path: PathBuf,

    /// The agent's definition, loaded from its `index.yaml` file.
    pub definition: AgentDefinition,

    /// A list of tools that are specific to this agent.
    pub tools: Vec<Tool>,

    /// A list of names of shared tools that this agent can use.
    pub shared_tools: Vec<String>,

    /// A map of variables and their values for this agent.
    pub variables: HashMap<String, String>,
}

impl Agent {
    /// Loads an agent from a specified directory.
    ///
    /// This function reads the agent's definition from `index.yaml`, loads its
    /// agent-specific tools from a `tools` subdirectory, and reads the list of
    /// shared tools from `tools.txt`.
    pub async fn from_directory<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref().to_path_buf();
        let name = path
            .file_name()
            .ok_or_else(|| anyhow::anyhow!("Invalid agent directory"))?
            .to_string_lossy()
            .to_string();

        let index_path = path.join("index.yaml");
        let index_content = tokio::fs::read_to_string(&index_path).await?;
        let definition: AgentDefinition = serde_yaml::from_str(&index_content)?;

        // Load agent-specific tools from a 'tools' subdirectory
        let mut tools = Vec::new();
        let agent_tools_dir = path.join("tools");
        if agent_tools_dir.is_dir() {
            let mut entries = tokio::fs::read_dir(agent_tools_dir).await?;
            while let Some(entry) = entries.next_entry().await? {
                let tool_path = entry.path();
                if tool_path.is_file() {
                    match Tool::from_file(&tool_path).await {
                        Ok(tool) => tools.push(tool),
                        Err(e) => tracing::warn!(
                            "Failed to load agent-specific tool {}: {}",
                            tool_path.display(),
                            e
                        ),
                    }
                }
            }
        }

        let shared_tools = if path.join("tools.txt").exists() {
            let content = tokio::fs::read_to_string(path.join("tools.txt")).await?;
            content
                .lines()
                .filter(|line| !line.trim().is_empty() && !line.starts_with('#'))
                .map(|s| s.to_string())
                .collect()
        } else {
            Vec::new()
        };

        let mut variables = HashMap::new();
        for var in &definition.variables {
            if let Some(default) = &var.default {
                variables.insert(var.name.clone(), default.clone());
            }
        }

        Ok(Self {
            name,
            path,
            definition,
            tools,
            shared_tools,
            variables,
        })
    }

    /// Retrieves all tools available to this agent, including its own tools and shared ones.
    ///
    /// # Arguments
    ///
    /// * `registry` - A reference to the `Registry` where shared tools are stored.
    pub fn get_all_tools<'a>(
        &'a self,
        registry: &'a crate::core::Registry,
    ) -> Result<Vec<&'a Tool>> {
        let mut all_tools: Vec<&'a Tool> = self.tools.iter().collect();

        for tool_name in &self.shared_tools {
            if let Some(tool) = registry.get_tool(tool_name) {
                all_tools.push(tool);
            } else {
                tracing::warn!(
                    "Shared tool '{}' for agent '{}' not found in registry.",
                    tool_name,
                    self.name
                );
            }
        }

        Ok(all_tools)
    }
}