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}