lha_core/tools/
registry.rs

1use crate::tools::context::ToolInvocation;
2use crate::tools::context::ToolOutput;
3use crate::tools::context::ToolPayload;
4use async_trait::async_trait;
5use lha_llm::ToolDescriptor;
6use lha_llm::ToolResultItem;
7use std::collections::HashMap;
8use std::sync::Arc;
9use thiserror::Error;
10use tokio_util::sync::CancellationToken;
11
12#[derive(Debug, Error, Clone, PartialEq, Eq)]
13pub enum ToolError {
14    #[error("{0}")]
15    Fatal(String),
16    #[error("{0}")]
17    RespondToModel(String),
18}
19
20#[derive(Debug, Clone)]
21pub struct ConfiguredTool {
22    pub spec: ToolDescriptor,
23    pub supports_parallel_tool_calls: bool,
24}
25
26#[async_trait]
27pub trait ToolHandler: Send + Sync {
28    fn spec(&self) -> ToolDescriptor;
29
30    fn supports_parallel_tool_calls(&self) -> bool {
31        false
32    }
33
34    async fn handle(
35        &self,
36        invocation: ToolInvocation,
37        cancellation_token: CancellationToken,
38    ) -> Result<ToolOutput, ToolError>;
39}
40
41pub struct ToolRegistry {
42    handlers: HashMap<String, Arc<dyn ToolHandler>>,
43    specs: Vec<ConfiguredTool>,
44}
45
46impl ToolRegistry {
47    pub fn new(
48        handlers: HashMap<String, Arc<dyn ToolHandler>>,
49        specs: Vec<ConfiguredTool>,
50    ) -> Self {
51        Self { handlers, specs }
52    }
53
54    pub fn specs(&self) -> Vec<ToolDescriptor> {
55        self.specs.iter().map(|tool| tool.spec.clone()).collect()
56    }
57
58    pub fn supports_parallel_tool_calls(&self, tool_name: &str) -> bool {
59        self.specs
60            .iter()
61            .find(|tool| tool.spec.name() == tool_name)
62            .is_some_and(|tool| tool.supports_parallel_tool_calls)
63    }
64
65    pub fn any_parallel_tool_calls(&self) -> bool {
66        self.specs
67            .iter()
68            .any(|tool| tool.supports_parallel_tool_calls)
69    }
70
71    pub async fn dispatch(
72        &self,
73        invocation: ToolInvocation,
74        cancellation_token: CancellationToken,
75    ) -> Result<ToolResultItem, ToolError> {
76        let tool_name = invocation.tool_name.clone();
77        let handler = self.handlers.get(tool_name.as_str()).ok_or_else(|| {
78            ToolError::RespondToModel(unsupported_tool_call_message(
79                &invocation.payload,
80                tool_name.as_str(),
81            ))
82        })?;
83
84        let output = handler
85            .handle(invocation.clone(), cancellation_token)
86            .await?;
87        Ok(output.into_response(
88            invocation.call_id.as_str(),
89            invocation.tool_name.as_str(),
90            &invocation.payload,
91        ))
92    }
93}
94
95#[derive(Default)]
96pub struct ToolRegistryBuilder {
97    handlers: HashMap<String, Arc<dyn ToolHandler>>,
98    specs: Vec<ConfiguredTool>,
99}
100
101impl ToolRegistryBuilder {
102    pub fn new() -> Self {
103        Self::default()
104    }
105
106    pub fn register_handler(&mut self, handler: Arc<dyn ToolHandler>) {
107        let spec = handler.spec();
108        let name = spec.name().to_string();
109        self.specs.push(ConfiguredTool {
110            spec,
111            supports_parallel_tool_calls: handler.supports_parallel_tool_calls(),
112        });
113        self.handlers.insert(name, handler);
114    }
115
116    pub fn push_spec(&mut self, spec: ToolDescriptor) {
117        self.specs.push(ConfiguredTool {
118            spec,
119            supports_parallel_tool_calls: false,
120        });
121    }
122
123    pub fn build(self) -> ToolRegistry {
124        ToolRegistry::new(self.handlers, self.specs)
125    }
126}
127
128fn unsupported_tool_call_message(payload: &ToolPayload, tool_name: &str) -> String {
129    match payload {
130        ToolPayload::TextInput { .. } => format!("unsupported custom tool call: {tool_name}"),
131        ToolPayload::JsonArguments { .. } => {
132            format!("unsupported call: {tool_name}")
133        }
134    }
135}
lha_core/tools/registry.rs

lha_core/tools/
registry.rs
