Skip to main content

llm_stack/tool/
handler.rs

1//! Tool handler trait and implementations.
2
3use std::future::Future;
4use std::marker::PhantomData;
5use std::pin::Pin;
6
7use serde_json::Value;
8
9use super::{ToolError, ToolOutput};
10use crate::provider::ToolDefinition;
11
12/// A single tool that can be invoked by the LLM.
13///
14/// Implement this trait for tools that need complex state or lifetime
15/// management. For simple tools, use [`super::tool_fn`] to wrap a closure.
16///
17/// The trait is generic over a context type `Ctx` which is passed to
18/// `execute()`. This allows tools to access shared state like database
19/// connections or user identity without closure capture. The default
20/// context type is `()` for backwards compatibility.
21///
22/// The trait is object-safe (uses boxed futures) so handlers can be
23/// stored as `Arc<dyn ToolHandler<Ctx>>`.
24///
25/// # Example with Context
26///
27/// ```rust
28/// use llm_stack::tool::{ToolHandler, ToolOutput, ToolError};
29/// use llm_stack::{ToolDefinition, JsonSchema};
30/// use serde_json::{json, Value};
31/// use std::future::Future;
32/// use std::pin::Pin;
33///
34/// struct AppContext {
35///     user_id: String,
36/// }
37///
38/// struct UserInfoTool;
39///
40/// impl ToolHandler<AppContext> for UserInfoTool {
41///     fn definition(&self) -> ToolDefinition {
42///         ToolDefinition {
43///             name: "get_user_info".into(),
44///             description: "Get current user info".into(),
45///             parameters: JsonSchema::new(json!({"type": "object"})),
46///             retry: None,
47///         }
48///     }
49///
50///     fn execute<'a>(
51///         &'a self,
52///         _input: Value,
53///         ctx: &'a AppContext,
54///     ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + 'a>> {
55///         Box::pin(async move {
56///             Ok(ToolOutput::new(format!("User: {}", ctx.user_id)))
57///         })
58///     }
59/// }
60/// ```
61pub trait ToolHandler<Ctx = ()>: Send + Sync {
62    /// Returns the tool's definition (name, description, parameter schema).
63    fn definition(&self) -> ToolDefinition;
64
65    /// Executes the tool with the given JSON arguments and context.
66    ///
67    /// Returns a [`ToolOutput`] containing the content for the LLM and
68    /// optional metadata for application use. Providers expect tool results
69    /// as text content — callers should `serde_json::to_string()` if they
70    /// have structured data.
71    fn execute<'a>(
72        &'a self,
73        input: Value,
74        ctx: &'a Ctx,
75    ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + 'a>>;
76}
77
78/// A tool handler backed by an async closure.
79///
80/// Created via [`super::tool_fn`] or [`super::tool_fn_with_ctx`].
81pub struct FnToolHandler<Ctx, F> {
82    pub(crate) definition: ToolDefinition,
83    pub(crate) handler: F,
84    pub(crate) _ctx: PhantomData<fn(&Ctx)>,
85}
86
87impl<Ctx, F> std::fmt::Debug for FnToolHandler<Ctx, F> {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        f.debug_struct("FnToolHandler")
90            .field("name", &self.definition.name)
91            .finish_non_exhaustive()
92    }
93}
94
95impl<Ctx, F, Fut, O> ToolHandler<Ctx> for FnToolHandler<Ctx, F>
96where
97    Ctx: Send + Sync + 'static,
98    F: for<'c> Fn(Value, &'c Ctx) -> Fut + Send + Sync,
99    Fut: Future<Output = Result<O, ToolError>> + Send + 'static,
100    O: Into<ToolOutput> + Send + 'static,
101{
102    fn definition(&self) -> ToolDefinition {
103        self.definition.clone()
104    }
105
106    fn execute<'a>(
107        &'a self,
108        input: Value,
109        ctx: &'a Ctx,
110    ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + 'a>> {
111        let fut = (self.handler)(input, ctx);
112        Box::pin(async move { fut.await.map(Into::into) })
113    }
114}
115
116/// A tool handler without context, created by [`super::tool_fn`].
117pub struct NoCtxToolHandler<F> {
118    pub(crate) definition: ToolDefinition,
119    pub(crate) handler: F,
120}
121
122impl<F> std::fmt::Debug for NoCtxToolHandler<F> {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.debug_struct("NoCtxToolHandler")
125            .field("name", &self.definition.name)
126            .finish_non_exhaustive()
127    }
128}
129
130impl<F, Fut, O> ToolHandler<()> for NoCtxToolHandler<F>
131where
132    F: Fn(Value) -> Fut + Send + Sync,
133    Fut: Future<Output = Result<O, ToolError>> + Send + 'static,
134    O: Into<ToolOutput> + Send + 'static,
135{
136    fn definition(&self) -> ToolDefinition {
137        self.definition.clone()
138    }
139
140    fn execute<'a>(
141        &'a self,
142        input: Value,
143        _ctx: &'a (),
144    ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + 'a>> {
145        let fut = (self.handler)(input);
146        Box::pin(async move { fut.await.map(Into::into) })
147    }
148}