96 lines
2.8 KiB
Rust
96 lines
2.8 KiB
Rust
//! Request-scoped tool registry.
|
|
//!
|
|
//! Tools are registered per-request (not globally) to keep the system testable
|
|
//! and allow different request contexts to have different tool sets.
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use futures::FutureExt;
|
|
|
|
use super::call::ToolError;
|
|
use super::context::ToolContext;
|
|
use super::definition::ToolDefinition;
|
|
|
|
/// Inner function pointer type for tool handlers.
|
|
type InnerHandlerFn = dyn Fn(
|
|
ToolContext,
|
|
serde_json::Value,
|
|
) -> std::pin::Pin<
|
|
Box<dyn std::future::Future<Output = Result<serde_json::Value, ToolError>> + Send>,
|
|
> + Send
|
|
+ Sync;
|
|
|
|
/// Wrapper around `Arc<dyn Fn(...)>` for `Clone` implementability.
|
|
#[derive(Clone)]
|
|
pub struct ToolHandler(std::sync::Arc<InnerHandlerFn>);
|
|
|
|
impl ToolHandler {
|
|
pub async fn execute(
|
|
&self,
|
|
ctx: ToolContext,
|
|
args: serde_json::Value,
|
|
) -> Result<serde_json::Value, ToolError> {
|
|
(self.0)(ctx, args).await
|
|
}
|
|
}
|
|
|
|
/// A request-scoped registry mapping tool names to their handlers.
|
|
#[derive(Clone, Default)]
|
|
pub struct ToolRegistry {
|
|
handlers: HashMap<String, ToolHandler>,
|
|
definitions: HashMap<String, ToolDefinition>,
|
|
}
|
|
|
|
impl ToolRegistry {
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
pub fn register_fn<F, Fut>(&mut self, name: impl Into<String>, handler: F) -> &mut Self
|
|
where
|
|
F: Fn(ToolContext, serde_json::Value) -> Fut + Send + Sync + 'static,
|
|
Fut: std::future::Future<Output = Result<serde_json::Value, ToolError>> + Send + 'static,
|
|
{
|
|
let name_str = name.into();
|
|
let def = ToolDefinition::new(&name_str);
|
|
let handler_fn: std::sync::Arc<InnerHandlerFn> =
|
|
std::sync::Arc::new(move |ctx, args| handler(ctx, args).boxed());
|
|
self.register(def, ToolHandler(handler_fn));
|
|
self
|
|
}
|
|
|
|
pub fn register(&mut self, def: ToolDefinition, handler: ToolHandler) -> &mut Self {
|
|
let name = def.name.clone();
|
|
if self.handlers.contains_key(&name) {
|
|
panic!("tool already registered: {}", name);
|
|
}
|
|
self.handlers.insert(name.clone(), handler);
|
|
self.definitions.insert(name, def);
|
|
self
|
|
}
|
|
|
|
/// Looks up a handler by tool name.
|
|
pub fn get(&self, name: &str) -> Option<&ToolHandler> {
|
|
self.handlers.get(name)
|
|
}
|
|
|
|
pub fn definitions(&self) -> std::collections::hash_map::Values<'_, String, ToolDefinition> {
|
|
self.definitions.values()
|
|
}
|
|
|
|
pub fn to_openai_tools(&self) -> Vec<async_openai::types::chat::ChatCompletionTool> {
|
|
self.definitions
|
|
.values()
|
|
.map(|d| d.to_openai_tool())
|
|
.collect()
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.handlers.len()
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.handlers.is_empty()
|
|
}
|
|
}
|