hms/crates/erp-ai/src/agent/registry.rs

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;

use super::tool::AgentTool;

/// Tool 注册表 — 管理所有可用的 Agent Tool
pub struct ToolRegistry {
    tools: HashMap<String, Arc<dyn AgentTool>>,
}

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

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

    pub fn register(&mut self, tool: Arc<dyn AgentTool>) {
        self.tools.insert(tool.name().to_string(), tool);
    }

    pub fn get(&self, name: &str) -> Option<&Arc<dyn AgentTool>> {
        self.tools.get(name)
    }

    pub fn all_tools(&self) -> Vec<&Arc<dyn AgentTool>> {
        self.tools.values().collect()
    }

    /// 根据角色允许的 Tool 列表过滤，返回过滤后的 ToolDefinition
    pub fn tool_definitions_filtered(
        &self,
        allowed_tools: &HashSet<String>,
    ) -> Vec<crate::dto::ToolDefinition> {
        self.tools
            .values()
            .filter(|t| allowed_tools.contains(t.name()))
            .map(|t| crate::dto::ToolDefinition {
                name: t.name().to_string(),
                description: t.description().to_string(),
                parameters: t.parameters_schema(),
            })
            .collect()
    }

    /// 生成传给 LLM 的 ToolDefinition 列表
    pub fn tool_definitions(&self) -> Vec<crate::dto::ToolDefinition> {
        self.tools
            .values()
            .map(|t| crate::dto::ToolDefinition {
                name: t.name().to_string(),
                description: t.description().to_string(),
                parameters: t.parameters_schema(),
            })
            .collect()
    }
}