gitdataai/lib/ai/tool/register.rs
2026-05-30 01:38:40 +08:00

66 lines
1.3 KiB
Rust

use crate::tool::tools::FunctionCall;
use std::collections::HashMap;
use std::sync::Arc;
#[derive(Clone)]
pub struct ToolRegister<C>
where
C: Clone + Send + Sync + 'static,
{
pub tools: Vec<Arc<dyn FunctionCall<Context = C>>>,
index: HashMap<String, usize>,
}
impl<C> ToolRegister<C>
where
C: Clone + Send + Sync + 'static,
{
pub fn new() -> Self {
ToolRegister {
tools: Vec::new(),
index: HashMap::new(),
}
}
pub fn register<T>(&mut self, tool: T)
where
T: FunctionCall<Context = C> + 'static,
{
let idx = self.tools.len();
self.index.insert(tool.name().to_string(), idx);
self.tools.push(Arc::new(tool));
}
pub fn with_tool<T>(mut self, tool: T) -> Self
where
T: FunctionCall<Context = C> + 'static,
{
self.register(tool);
self
}
pub fn get(
&self,
name: &str,
) -> Option<Arc<dyn FunctionCall<Context = C>>> {
self.index.get(name).map(|&idx| self.tools[idx].clone())
}
pub fn is_empty(&self) -> bool {
self.tools.is_empty()
}
pub fn len(&self) -> usize {
self.tools.len()
}
}
impl<C> Default for ToolRegister<C>
where
C: Clone + Send + Sync + 'static,
{
fn default() -> Self {
Self::new()
}
}