Skip to main content

stynx_code_tools/application/
registry.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use stynx_code_types::Tool;
5use serde_json::{Value, json};
6
7pub struct ToolRegistry {
8    tools: HashMap<String, Arc<dyn Tool>>,
9}
10
11impl ToolRegistry {
12    pub fn new() -> Self {
13        Self {
14            tools: HashMap::new(),
15        }
16    }
17
18    pub fn register(&mut self, tool: Arc<dyn Tool>) {
19        self.tools.insert(tool.name().to_string(), tool);
20    }
21
22    pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
23        self.tools.get(name)
24    }
25
26    pub fn get_by_name_or_alias(&self, name: &str) -> Option<&Arc<dyn Tool>> {
27        if let Some(tool) = self.tools.get(name) {
28            return Some(tool);
29        }
30        self.tools.values().find(|t| t.aliases().contains(&name))
31    }
32
33    fn make_tool_def(t: &dyn Tool) -> Value {
34        let mut def = json!({
35            "name": t.name(),
36            "description": t.description(),
37            "input_schema": t.input_schema(),
38        });
39        if t.strict() {
40            def.as_object_mut().unwrap().insert("strict".to_string(), json!(true));
41        }
42        def
43    }
44
45    pub fn tool_definitions(&self) -> Vec<Value> {
46        self.tools
47            .values()
48            .filter(|t| !t.should_defer())
49            .map(|t| Self::make_tool_def(t.as_ref()))
50            .collect()
51    }
52
53    pub fn deferred_tool_definitions(&self) -> Vec<Value> {
54        self.tools
55            .values()
56            .filter(|t| t.should_defer())
57            .map(|t| {
58                json!({
59                    "name": t.name(),
60                    "description": t.description(),
61                })
62            })
63            .collect()
64    }
65
66    pub fn tool_definitions_filtered<F>(&self, predicate: F) -> Vec<Value>
67    where
68        F: Fn(&dyn Tool) -> bool,
69    {
70        self.tools
71            .values()
72            .filter(|t| predicate(t.as_ref()))
73            .map(|t| Self::make_tool_def(t.as_ref()))
74            .collect()
75    }
76
77    pub fn tool_names(&self) -> Vec<String> {
78        let mut names: Vec<String> = self.tools.keys().cloned().collect();
79        names.sort();
80        names
81    }
82
83    pub fn clone_excluding(&self, exclude: &[&str]) -> Self {
84        let tools = self
85            .tools
86            .iter()
87            .filter(|(name, _)| !exclude.contains(&name.as_str()))
88            .map(|(k, v)| (k.clone(), v.clone()))
89            .collect();
90        Self { tools }
91    }
92
93    pub fn keep_only(&self, keep: &[&str]) -> Self {
94        let tools = self
95            .tools
96            .iter()
97            .filter(|(name, _)| keep.contains(&name.as_str()))
98            .map(|(k, v)| (k.clone(), v.clone()))
99            .collect();
100        Self { tools }
101    }
102}
103
104impl Default for ToolRegistry {
105    fn default() -> Self {
106        Self::new()
107    }
108}