openfunctions-rs 0.1.0

//! Registry for managing and accessing tools and agents.

use crate::core::{Agent, Tool};
use anyhow::Result;
use std::collections::HashMap;
use std::path::PathBuf;
use tracing::{info, warn};

/// A registry for discovering, storing, and retrieving tools and agents.
///
/// The `Registry` is responsible for loading tools and agents from the filesystem
/// and providing a central place to access them by name.
#[derive(Debug, Default)]
pub struct Registry {
    tools: HashMap<String, Tool>,
    agents: HashMap<String, Agent>,
}

impl Registry {
    /// Creates a new, empty `Registry`.
    pub fn new() -> Self {
        Self::default()
    }

    /// Loads all tools from the specified directories.
    ///
    /// This method scans the provided directories for supported script files
    /// (e.g., `.sh`, `.js`, `.py`) and loads them as `Tool`s.
    pub async fn load_tools(&mut self, tool_dirs: &[PathBuf]) -> Result<()> {
        info!(dirs = ?tool_dirs, "Loading tools from directories");
        for dir in tool_dirs {
            if !dir.exists() {
                warn!(dir = %dir.display(), "Tool directory not found, skipping");
                continue;
            }
            if dir.is_dir() {
                let mut entries = tokio::fs::read_dir(dir).await?;
                while let Some(entry) = entries.next_entry().await? {
                    let path = entry.path();
                    if path.is_file() {
                        if let Some(ext) = path.extension() {
                            if matches!(ext.to_str(), Some("sh" | "js" | "py")) {
                                match Tool::from_file(&path).await {
                                    Ok(tool) => {
                                        info!(name = %tool.name, "Loaded tool");
                                        self.tools.insert(tool.name.clone(), tool);
                                    }
                                    Err(e) => {
                                        warn!(path = %path.display(), error = ?e, "Failed to load tool");
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Loads all agents from the specified directories.
    ///
    /// This method scans the provided directories for agent definitions and
    /// loads them into the registry.
    pub async fn load_agents(&mut self, agent_dirs: &[PathBuf]) -> Result<()> {
        info!(dirs = ?agent_dirs, "Loading agents from directories");
        for dir in agent_dirs {
            if !dir.exists() {
                warn!(dir = %dir.display(), "Agent directory not found, skipping");
                continue;
            }
            if dir.is_dir() {
                let mut entries = tokio::fs::read_dir(dir).await?;
                while let Some(entry) = entries.next_entry().await? {
                    let path = entry.path();
                    if path.is_dir() {
                        match Agent::from_directory(&path).await {
                            Ok(agent) => {
                                info!(name = %agent.name, "Loaded agent");
                                self.agents.insert(agent.name.clone(), agent);
                            }
                            Err(e) => {
                                warn!(path = %path.display(), error = ?e, "Failed to load agent");
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Retrieves a reference to a tool by its name.
    pub fn get_tool(&self, name: &str) -> Option<&Tool> {
        self.tools.get(name)
    }

    /// Retrieves a reference to an agent by its name.
    pub fn get_agent(&self, name: &str) -> Option<&Agent> {
        self.agents.get(name)
    }

    /// Returns a list of names of all registered tools.
    pub fn list_tools(&self) -> Vec<String> {
        self.tools.keys().cloned().collect()
    }

    /// Returns a list of names of all registered agents.
    pub fn list_agents(&self) -> Vec<String> {
        self.agents.keys().cloned().collect()
    }
}