PicoBot/src/tools/registry.rs

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use crate::domain::tools::{Tool, ToolFunction};

use super::traits::Tool as ToolTrait;

pub struct ToolRegistry {
    tools: RwLock<HashMap<String, Arc<dyn ToolTrait>>>,
}

impl ToolRegistry {
    pub fn new() -> Self {
        Self {
            tools: RwLock::new(HashMap::new()),
        }
    }

    pub fn register<T: ToolTrait + 'static>(&self, tool: T) {
        self.tools
            .write()
            .expect("ToolRegistry lock poisoned")
            .insert(tool.name().to_string(), Arc::new(tool));
    }

    pub fn get(&self, name: &str) -> Option<Arc<dyn ToolTrait>> {
        self.tools
            .read()
            .expect("ToolRegistry lock poisoned")
            .get(name)
            .cloned()
    }

    /// Get all registered tools.
    /// Used for concurrent tool execution when we need to look up tools by name.
    pub fn get_all(&self) -> Vec<Arc<dyn ToolTrait>> {
        self.tools
            .read()
            .expect("ToolRegistry lock poisoned")
            .values()
            .cloned()
            .collect()
    }

    pub fn get_definitions(&self) -> Vec<Tool> {
        self.tools
            .read()
            .expect("ToolRegistry lock poisoned")
            .values()
            .map(|tool| Tool {
                tool_type: "function".to_string(),
                function: ToolFunction {
                    name: tool.name().to_string(),
                    description: tool.description().to_string(),
                    parameters: tool.parameters_schema(),
                },
            })
            .collect()
    }

    pub fn has_tools(&self) -> bool {
        !self.tools
            .read()
            .expect("ToolRegistry lock poisoned")
            .is_empty()
    }

    pub fn tool_names(&self) -> Vec<String> {
        self.tools
            .read()
            .expect("ToolRegistry lock poisoned")
            .keys()
            .cloned()
            .collect()
    }

    /// 创建一个排除指定工具的新 registry 副本
    pub fn without(&self, exclude: &[&str]) -> Self {
        let exclude_set: std::collections::HashSet<&str> = exclude.iter().copied().collect();
        let tools = self.tools.read().expect("ToolRegistry lock poisoned");
        let filtered: HashMap<String, Arc<dyn ToolTrait>> = tools
            .iter()
            .filter(|(name, _)| !exclude_set.contains(name.as_str()))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        let new_registry = ToolRegistry::new();
        *new_registry.tools.write().expect("ToolRegistry lock poisoned") = filtered;
        new_registry
    }
}

impl Default for ToolRegistry {
    fn default() -> Self {
        Self::new()
    }
}