Skip to main content

inference_gateway_adk/server/
agent_toolbox.rs

1use anyhow::Result;
2use serde_json::Value;
3
4/// Trait for handling tool calls
5#[async_trait::async_trait]
6pub trait ToolHandler: Send + Sync {
7    /// Handle a tool call with the given arguments and return the result
8    async fn handle(&self, args: Value) -> Result<String>;
9}
10
11/// A simple function-based tool handler
12pub struct FunctionToolHandler<F>
13where
14    F: Fn(Value) -> Result<String> + Send + Sync,
15{
16    handler: F,
17}
18
19impl<F> FunctionToolHandler<F>
20where
21    F: Fn(Value) -> Result<String> + Send + Sync,
22{
23    pub fn new(handler: F) -> Self {
24        Self { handler }
25    }
26}
27
28#[async_trait::async_trait]
29impl<F> ToolHandler for FunctionToolHandler<F>
30where
31    F: Fn(Value) -> Result<String> + Send + Sync,
32{
33    async fn handle(&self, args: Value) -> Result<String> {
34        (self.handler)(args)
35    }
36}
37
38/// An async function-based tool handler
39pub struct AsyncFunctionToolHandler<F, Fut>
40where
41    F: Fn(Value) -> Fut + Send + Sync,
42    Fut: std::future::Future<Output = Result<String>> + Send,
43{
44    handler: F,
45}
46
47impl<F, Fut> AsyncFunctionToolHandler<F, Fut>
48where
49    F: Fn(Value) -> Fut + Send + Sync,
50    Fut: std::future::Future<Output = Result<String>> + Send,
51{
52    pub fn new(handler: F) -> Self {
53        Self { handler }
54    }
55}
56
57#[async_trait::async_trait]
58impl<F, Fut> ToolHandler for AsyncFunctionToolHandler<F, Fut>
59where
60    F: Fn(Value) -> Fut + Send + Sync,
61    Fut: std::future::Future<Output = Result<String>> + Send,
62{
63    async fn handle(&self, args: Value) -> Result<String> {
64        (self.handler)(args).await
65    }
66}