Skip to main content

matrixcode_core/tools/
registry.rs

1//! Tool Registry
2//!
3//! 工具注册中心 - 仅管理代理工具
4
5use super::ToolDefinition;
6use super::toolproxy::{ProxyMetadata, ProxyTool};
7
8/// 工具注册中心 - 仅管理代理工具(内置工具由 Agent 直接管理)
9pub struct ToolRegistry {
10    /// 代理工具(外部注入的工具)
11    proxy_tools: Vec<ProxyTool>,
12}
13
14impl ToolRegistry {
15    /// 创建空的工具注册中心
16    pub fn new() -> Self {
17        Self {
18            proxy_tools: Vec::new(),
19        }
20    }
21
22    /// 注册代理工具
23    pub fn register(&mut self, tool: ProxyTool) {
24        // 检查工具名称冲突
25        if self
26            .proxy_tools
27            .iter()
28            .any(|t| t.definition().name == tool.definition().name)
29        {
30            log::warn!(
31                "Proxy tool '{}' already exists, will be replaced",
32                tool.definition().name
33            );
34            // 移除旧的
35            self.proxy_tools
36                .retain(|t| t.definition().name != tool.definition().name);
37        }
38
39        self.proxy_tools.push(tool);
40    }
41
42    /// 批量注册代理工具
43    pub fn register_batch(&mut self, tools: Vec<ProxyTool>) {
44        for tool in tools {
45            self.register(tool);
46        }
47    }
48
49    /// 获取所有代理工具定义(用于提交给大模型)
50    pub fn proxy_definitions(&self) -> Vec<ToolDefinition> {
51        self.proxy_tools.iter().map(|t| t.definition()).collect()
52    }
53
54    /// 获取所有代理工具
55    pub fn proxy_tools(&self) -> &[ProxyTool] {
56        &self.proxy_tools
57    }
58
59    /// 查找代理工具
60    pub fn find_proxy(&self, name: &str) -> Option<&ProxyTool> {
61        self.proxy_tools
62            .iter()
63            .find(|t| t.definition().name == name)
64    }
65
66    /// 获取代理工具数量
67    pub fn count(&self) -> usize {
68        self.proxy_tools.len()
69    }
70
71    /// 检查是否是代理工具
72    pub fn is_proxy(&self, name: &str) -> bool {
73        self.proxy_tools.iter().any(|t| t.definition().name == name)
74    }
75
76    /// 获取代理工具元数据
77    pub fn get_metadata(&self, name: &str) -> Option<&ProxyMetadata> {
78        self.proxy_tools
79            .iter()
80            .find(|t| t.definition().name == name)
81            .map(|t| t.metadata())
82    }
83
84    /// 清空所有代理工具
85    pub fn clear(&mut self) {
86        self.proxy_tools.clear();
87    }
88}
89
90impl Default for ToolRegistry {
91    fn default() -> Self {
92        Self::new()
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use serde_json::json;
100
101    #[test]
102    fn test_registry_creation() {
103        let registry = ToolRegistry::new();
104        assert_eq!(registry.count(), 0);
105    }
106
107    #[test]
108    fn test_register_proxy_tool() {
109        let mut registry = ToolRegistry::new();
110
111        let proxy_tool = ProxyTool::new(
112            ToolDefinition {
113                name: "custom_search".to_string(),
114                description: "自定义搜索".to_string(),
115                parameters: json!({"type": "object"}),
116                ..Default::default()
117            },
118            ProxyMetadata {
119                tool_type: "search".to_string(),
120                endpoint: None,
121                timeout_ms: 1000,
122                custom: None,
123            },
124        );
125
126        registry.register(proxy_tool);
127
128        assert_eq!(registry.count(), 1);
129        assert!(registry.is_proxy("custom_search"));
130        assert!(registry.find_proxy("custom_search").is_some());
131    }
132
133    #[test]
134    fn test_proxy_definitions() {
135        let mut registry = ToolRegistry::new();
136
137        registry.register(ProxyTool::new(
138            ToolDefinition {
139                name: "test".to_string(),
140                description: "Test".to_string(),
141                parameters: json!({}),
142                ..Default::default()
143            },
144            ProxyMetadata {
145                tool_type: "test".to_string(),
146                endpoint: None,
147                timeout_ms: 1000,
148                custom: None,
149            },
150        ));
151
152        let definitions = registry.proxy_definitions();
153        assert_eq!(definitions.len(), 1);
154        assert_eq!(definitions[0].name, "test");
155    }
156}