omni_llm_kit/tool/
tool_registry.rs

1use crate::{Tool, ToolDyn};
2use derive_more::{Deref, DerefMut};
3use parking_lot::RwLock;
4use rustc_hash::FxHashMap;
5use std::sync::Arc;
6
7#[derive(Default)]
8struct ToolRegistryState {
9    tools: FxHashMap<Arc<str>, Arc<dyn ToolDyn>>,
10}
11#[derive(Default)]
12pub struct ToolRegistry {
13    state: RwLock<ToolRegistryState>,
14}
15
16impl ToolRegistry {
17    pub fn new() -> Self {
18        Self {
19            state: RwLock::new(ToolRegistryState {
20                tools: FxHashMap::default(),
21            }),
22        }
23    }
24    pub fn tool(&self, name: &str) -> Option<Arc<dyn ToolDyn>> {
25        self.state.read().tools.get(name).cloned()
26    }
27    pub fn register_tool(&self, tool: impl ToolDyn + 'static) {
28        eprintln!("registering tool {:?}", tool.name());
29        let mut state = self.state.write();
30        let tool_name: Arc<str> = tool.name().into();
31        state.tools.insert(tool_name, Arc::new(tool));
32    }
33
34    pub fn tools(&self) -> Vec<Arc<dyn ToolDyn>> {
35        self.state.read().tools.values().cloned().collect()
36    }
37}
omni_llm_kit/tool/tool_registry.rs

omni_llm_kit/tool/
tool_registry.rs
