Hare/llm_worker_rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

restrict tool registration to mutable state

Browse Source
This commit is contained in:
Keisuke Hirata 2026-02-19 18:29:02 +09:00
parent b12785ed93
commit d63b4ea470
2 changed files with 40 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
llm-worker/src/tool_server.rs
View File

@ -50,7 +50,10 @@ pub struct ToolServerHandle {
impl ToolServerHandle {
/// Register one tool.
pub fn register_tool(&self, factory: WorkerToolDefinition) -> Result<(), ToolServerError> {
pub(crate) fn register_tool(
&self,
factory: WorkerToolDefinition,
) -> Result<(), ToolServerError> {
let (meta, instance) = factory();
let mut guard = self.tools.lock().unwrap_or_else(|e| e.into_inner());
if guard.contains_key(&meta.name) {
@ -61,7 +64,7 @@ impl ToolServerHandle {
}
/// Register many tools.
pub fn register_tools(
pub(crate) fn register_tools(
&self,
factories: impl IntoIterator<Item = WorkerToolDefinition>,
) -> Result<(), ToolServerError> {

75
llm-worker/src/worker.rs
View File

@ -300,46 +300,6 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
.push(Box::new(SubscriberTurnNotifier { subscriber }));
}
/// Register a tool
///
/// Registered tools are automatically executed when called by the LLM.
/// Registering a tool with the same name will result in an error.
///
/// # Examples
///
/// ```ignore
/// use llm_worker::tool::{ToolMeta, ToolDefinition, Tool};
/// use std::sync::Arc;
///
/// let def: ToolDefinition = Arc::new(|| {
/// (ToolMeta::new("search").description("..."), Arc::new(MyTool) as Arc<dyn Tool>)
/// });
/// worker.register_tool(def)?;
/// ```
pub fn register_tool(&mut self, factory: WorkerToolDefinition) -> Result<(), ToolRegistryError> {
match self.tool_server.register_tool(factory) {
Ok(()) => Ok(()),
Err(ToolServerError::DuplicateName(name)) => Err(ToolRegistryError::DuplicateName(name)),
Err(ToolServerError::ToolNotFound(_) | ToolServerError::ToolExecution(_)) => {
unreachable!("register_tool should only fail with DuplicateName")
}
}
}
/// Register multiple tools
pub fn register_tools(
&mut self,
factories: impl IntoIterator<Item = WorkerToolDefinition>,
) -> Result<(), ToolRegistryError> {
match self.tool_server.register_tools(factories) {
Ok(()) => Ok(()),
Err(ToolServerError::DuplicateName(name)) => Err(ToolRegistryError::DuplicateName(name)),
Err(ToolServerError::ToolNotFound(_) | ToolServerError::ToolExecution(_)) => {
unreachable!("register_tools should only fail with DuplicateName")
}
}
}
/// Get a shared tool server handle.
pub fn tool_server_handle(&self) -> ToolServerHandle {
self.tool_server.clone()
@ -1059,6 +1019,41 @@ impl<C: LlmClient> Worker<C, Mutable> {
}
}
/// Register a tool
///
/// Registered tools are automatically executed when called by the LLM.
/// Registering a tool with the same name will result in an error.
///
/// Available only in Mutable state.
pub fn register_tool(
&mut self,
factory: WorkerToolDefinition,
) -> Result<(), ToolRegistryError> {
match self.tool_server.register_tool(factory) {
Ok(()) => Ok(()),
Err(ToolServerError::DuplicateName(name)) => Err(ToolRegistryError::DuplicateName(name)),
Err(ToolServerError::ToolNotFound(_) | ToolServerError::ToolExecution(_)) => {
unreachable!("register_tool should only fail with DuplicateName")
}
}
}
/// Register multiple tools
///
/// Available only in Mutable state.
pub fn register_tools(
&mut self,
factories: impl IntoIterator<Item = WorkerToolDefinition>,
) -> Result<(), ToolRegistryError> {
match self.tool_server.register_tools(factories) {
Ok(()) => Ok(()),
Err(ToolServerError::DuplicateName(name)) => Err(ToolRegistryError::DuplicateName(name)),
Err(ToolServerError::ToolNotFound(_) | ToolServerError::ToolExecution(_)) => {
unreachable!("register_tools should only fail with DuplicateName")
}
}
}
/// Set system prompt (builder pattern)
pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.system_prompt = Some(prompt.into());