Skip to main content

temporal_agent_rs/
tool.rs

1//! Tool registry — bridges AutoAgents' [`ToolT`] trait into a name-based
2//! dispatch table used by the `execute_tool` activity.
3//!
4//! Tools authored with AutoAgents' `#[tool]` derive can be registered as-is.
5//! The workflow only references tools by name; actual `ToolT::execute` calls
6//! happen exclusively inside the activity, preserving workflow determinism.
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use autoagents_core::tool::ToolT;
12
13use crate::state::ToolSchema;
14
15/// Immutable map of tool name → boxed tool implementation.
16///
17/// Shared across all in-flight activity invocations on a worker via `Arc`.
18#[derive(Clone, Default)]
19pub struct ToolRegistry {
20    tools: Arc<HashMap<String, Arc<dyn ToolT>>>,
21}
22
23impl ToolRegistry {
24    /// Start building a registry with [`ToolRegistryBuilder`].
25    pub fn builder() -> ToolRegistryBuilder {
26        ToolRegistryBuilder::default()
27    }
28
29    /// Look up a registered tool by its name. Returns `None` if no tool with
30    /// that name was registered with the worker.
31    pub fn get(&self, name: &str) -> Option<&Arc<dyn ToolT>> {
32        self.tools.get(name)
33    }
34
35    /// Iterate the names of every registered tool (no defined order).
36    pub fn names(&self) -> impl Iterator<Item = &str> {
37        self.tools.keys().map(String::as_str)
38    }
39
40    /// Produce JSON-Schema descriptions of every registered tool to send to
41    /// the LLM each turn.
42    pub fn to_schemas(&self) -> Vec<ToolSchema> {
43        self.tools
44            .values()
45            .map(|t| ToolSchema {
46                name: t.name().to_string(),
47                description: t.description().to_string(),
48                args_schema: t.args_schema(),
49            })
50            .collect()
51    }
52}
53
54/// Fluent builder for [`ToolRegistry`]. Use via [`ToolRegistry::builder`].
55#[derive(Default)]
56pub struct ToolRegistryBuilder {
57    tools: HashMap<String, Arc<dyn ToolT>>,
58}
59
60impl ToolRegistryBuilder {
61    /// Register a tool. Tools are keyed by [`ToolT::name`]; later
62    /// registrations with the same name overwrite earlier ones.
63    #[allow(clippy::should_implement_trait)]
64    #[must_use]
65    pub fn add(mut self, tool: Arc<dyn ToolT>) -> Self {
66        self.tools.insert(tool.name().to_string(), tool);
67        self
68    }
69
70    /// Freeze the builder into an immutable [`ToolRegistry`].
71    pub fn build(self) -> ToolRegistry {
72        ToolRegistry {
73            tools: Arc::new(self.tools),
74        }
75    }
76}