Skip to main content

rig_core/tool/
contextual.rs

1//! Contextual tool authoring and JSON-argument dispatch adapters.
2//! [`ToolContext`] carries typed inbound values and host-only result metadata;
3//! model-visible outputs retain their text, JSON, or multimodal representation.
4//!
5//! ```
6//! use rig_core::{message::ToolName, tool::{DynamicTool, ToolOutput}};
7//!
8//! let tool = DynamicTool::new_with_context(ToolName::new("echo")?, "Echo JSON", serde_json::json!({}),
9//!     |_context, args| Box::pin(async move { Ok(ToolOutput::json(args)) }));
10//! assert_eq!(tool.name(), "echo");
11//! # Ok::<(), rig_core::message::EmptyToolName>(())
12//! ```
13
14use std::{future::Future, sync::Arc};
15
16use serde::{Deserialize, Serialize};
17
18use crate::{
19    completion::{ToolDefinition, message::ToolName},
20    effect::{EffectKind, Outcome},
21    serve::{ErasedHandler, adapters::ToolCallback, adapters::ToolFn},
22    wasm_compat::{WasmBoxedFuture, WasmCompatSend, WasmCompatSync},
23};
24
25use super::{
26    IntoToolOutput, PublishedContext, ToolContext, ToolExecutionError, ToolOutput, ToolResult,
27};
28
29/// A typed LLM tool.
30///
31/// Tool authors provide metadata and exactly one execution method. Runtime
32/// context and host-only result metadata share the [`ToolContext`] path. Rig's
33/// object-safe dispatch boundary is private; use [`DynamicTool`] when the tool
34/// name or callback is only known at runtime.
35pub trait Tool: Sized + WasmCompatSend + WasmCompatSync {
36    /// Unique registration and provider-facing name.
37    const NAME: &'static str;
38    /// Typed JSON arguments.
39    type Args: for<'de> Deserialize<'de> + WasmCompatSend + WasmCompatSync;
40    /// Output convertible into Rig's canonical model presentation.
41    ///
42    /// Every owned serializable value implements [`IntoToolOutput`]
43    /// automatically. [`ToolResultContent`](crate::message::ToolResultContent)
44    /// and `Vec<ToolResultContent>` preserve rich content when returned
45    /// directly; use [`ToolOutput`] when constructing the presentation
46    /// explicitly.
47    type Output: IntoToolOutput;
48    /// Typed error returned by direct calls to this tool.
49    ///
50    /// Rig normalizes this error into [`ToolExecutionError`] only at the erased
51    /// dispatch boundary. This keeps ordinary `?` propagation and typed unit
52    /// tests available to tool authors without creating a second runtime error
53    /// representation.
54    type Error: std::error::Error + WasmCompatSend + WasmCompatSync + 'static;
55
56    /// Model-facing description.
57    fn description(&self) -> String;
58
59    /// JSON Schema for arguments.
60    fn parameters(&self) -> serde_json::Value;
61
62    /// Normalize a typed author-facing error for runtime policy and telemetry.
63    ///
64    /// The default preserves concrete sources for operators and exposes safe
65    /// kind-level model feedback. An existing [`ToolExecutionError`] retains its
66    /// classification and model output. Override to supply deliberate domain feedback.
67    fn map_error(&self, error: Self::Error) -> ToolExecutionError {
68        ToolExecutionError::from_error(error)
69    }
70
71    /// Execute the tool.
72    fn call(
73        &self,
74        context: &mut ToolContext,
75        args: Self::Args,
76    ) -> impl Future<Output = Result<Self::Output, Self::Error>> + WasmCompatSend;
77}
78
79impl<T> Tool for T
80where
81    T: super::PortableTool,
82{
83    const NAME: &'static str = <T as super::PortableTool>::NAME;
84    type Args = <T as super::PortableTool>::Args;
85    type Output = <T as super::PortableTool>::Output;
86    type Error = <T as super::PortableTool>::Error;
87
88    fn description(&self) -> String {
89        super::PortableTool::description(self)
90    }
91
92    fn parameters(&self) -> serde_json::Value {
93        super::PortableTool::parameters(self)
94    }
95
96    fn map_error(&self, error: Self::Error) -> ToolExecutionError {
97        super::PortableTool::map_error(self, error)
98    }
99
100    async fn call(
101        &self,
102        _context: &mut ToolContext,
103        args: Self::Args,
104    ) -> Result<Self::Output, Self::Error> {
105        super::PortableTool::call(self, args).await
106    }
107}
108
109/// A tool that can be stored in a vector store and reconstructed for RAG.
110pub trait ToolEmbedding: Tool {
111    /// Error returned while reconstructing the tool.
112    type InitError: std::error::Error + WasmCompatSend + WasmCompatSync + 'static;
113    /// Serializable static context.
114    type Context: for<'de> Deserialize<'de> + Serialize;
115    /// Runtime initialization state.
116    type State: WasmCompatSend;
117
118    /// Documents used to retrieve the tool.
119    fn embedding_docs(&self) -> Vec<String>;
120    /// Serializable tool context.
121    fn context(&self) -> Self::Context;
122    /// Reconstruct the tool.
123    fn init(state: Self::State, context: Self::Context) -> Result<Self, Self::InitError>;
124}
125
126impl<T> ToolEmbedding for T
127where
128    T: super::PortableToolEmbedding,
129{
130    type InitError = <T as super::PortableToolEmbedding>::InitError;
131    type Context = <T as super::PortableToolEmbedding>::Context;
132    type State = <T as super::PortableToolEmbedding>::State;
133
134    fn embedding_docs(&self) -> Vec<String> {
135        super::PortableToolEmbedding::embedding_docs(self)
136    }
137
138    fn context(&self) -> Self::Context {
139        super::PortableToolEmbedding::context(self)
140    }
141
142    fn init(state: Self::State, context: Self::Context) -> Result<Self, Self::InitError> {
143        super::PortableToolEmbedding::init(state, context)
144    }
145}
146
147fn parse_tool_args<A>(args: &str) -> Result<A, ToolExecutionError>
148where
149    A: serde::de::DeserializeOwned,
150{
151    match serde_json::from_str(args) {
152        Ok(parsed) => Ok(parsed),
153        Err(original) if args.trim() == "null" => serde_json::from_str("{}").map_err(|_| {
154            ToolExecutionError::invalid_args(format!("failed to parse tool arguments: {original}"))
155                .with_source(original)
156        }),
157        Err(error) => Err(ToolExecutionError::invalid_args(format!(
158            "failed to parse tool arguments: {error}"
159        ))
160        .with_source(error)),
161    }
162}
163
164/// Parses JSON arguments and runs a contextual callback, returning parse,
165/// execution, or output-conversion failures as failed tool results.
166pub(crate) async fn execute_callback<F>(
167    callback: &F,
168    args: String,
169    context: &mut ToolContext,
170) -> ToolResult
171where
172    F: for<'a> Fn(
173        &'a mut ToolContext,
174        serde_json::Value,
175    ) -> WasmBoxedFuture<'a, Result<ToolOutput, ToolExecutionError>>,
176{
177    let args = match parse_tool_args::<serde_json::Value>(&args) {
178        Ok(args) => args,
179        Err(error) => return ToolResult::failed(error),
180    };
181    tool_result_from(callback(context, args).await)
182}
183
184fn tool_result_from<O>(outcome: Result<O, ToolExecutionError>) -> ToolResult
185where
186    O: IntoToolOutput,
187{
188    match outcome.and_then(IntoToolOutput::into_tool_output) {
189        Ok(output) => ToolResult::success(output),
190        Err(error) => ToolResult::failed(error),
191    }
192}
193
194/// The object-safe form of [`Tool`]: raw JSON arguments in, a
195/// [`ToolResult`] out. This is the impl-side contract the bus's
196/// `ToolAdapter` calls; nothing stores it behind a vtable.
197pub trait ErasedTool: WasmCompatSend + WasmCompatSync {
198    /// The tool's name.
199    fn name(&self) -> String;
200    /// The tool's description.
201    fn description(&self) -> String;
202    /// The JSON schema of the tool's arguments.
203    fn parameters(&self) -> serde_json::Value;
204    /// Run the tool on raw arguments, shaping the answer into a result.
205    fn execute<'a>(
206        &'a self,
207        args: String,
208        context: &'a mut ToolContext,
209    ) -> WasmBoxedFuture<'a, ToolResult>;
210}
211
212impl<T> ErasedTool for T
213where
214    T: Tool,
215{
216    fn name(&self) -> String {
217        T::NAME.to_string()
218    }
219
220    fn description(&self) -> String {
221        Tool::description(self)
222    }
223
224    fn parameters(&self) -> serde_json::Value {
225        Tool::parameters(self)
226    }
227
228    fn execute<'a>(
229        &'a self,
230        args: String,
231        context: &'a mut ToolContext,
232    ) -> WasmBoxedFuture<'a, ToolResult> {
233        Box::pin(async move {
234            let args = match parse_tool_args::<T::Args>(&args) {
235                Ok(args) => args,
236                Err(error) => return ToolResult::failed(error),
237            };
238            tool_result_from(
239                Tool::call(self, context, args)
240                    .await
241                    .map_err(|error| Tool::map_error(self, error)),
242            )
243        })
244    }
245}
246
247/// Reports whether a tool's owner still serves it. Registries check lazily on
248/// reads or reconciliation; calls before retirement may fail with a transport
249/// error rather than `HandlerUnavailable`.
250#[cfg(not(target_family = "wasm"))]
251pub type LivenessFn = Arc<dyn Fn() -> bool + Send + Sync>;
252/// A liveness probe (browser wasm: no `Send + Sync`, no threads).
253#[cfg(target_family = "wasm")]
254pub type LivenessFn = Arc<dyn Fn() -> bool>;
255
256/// A tool defined at runtime: a name, a schema and a callback. The callback
257/// is the handler ([`ToolFn`]); this struct is its definition plus the
258/// erased handler a registry stages until a bus takes it. The optional
259/// liveness probe supports registry retirement; inline execution does not
260/// consult it.
261#[derive(Clone)]
262pub struct DynamicTool {
263    definition: ToolDefinition,
264    handler: ErasedHandler,
265    liveness: Option<LivenessFn>,
266}
267
268impl DynamicTool {
269    /// Define a tool from a context-free callback over owned arguments.
270    pub fn new<F>(
271        name: ToolName,
272        description: impl Into<String>,
273        parameters: serde_json::Value,
274        callback: F,
275    ) -> Self
276    where
277        F: Fn(
278                serde_json::Value,
279            ) -> WasmBoxedFuture<'static, Result<ToolOutput, ToolExecutionError>>
280            + WasmCompatSend
281            + WasmCompatSync
282            + 'static,
283    {
284        Self::new_with_context(
285            name,
286            description,
287            parameters,
288            move |_context: &mut ToolContext, arguments| callback(arguments),
289        )
290    }
291
292    /// Define a tool from a callback over the dispatch-scoped context.
293    pub fn new_with_context<F>(
294        name: ToolName,
295        description: impl Into<String>,
296        parameters: serde_json::Value,
297        callback: F,
298    ) -> Self
299    where
300        F: ToolCallback + 'static,
301    {
302        let description = description.into();
303        let handler = ErasedHandler::new(ToolFn::new(
304            name.to_string(),
305            description.clone(),
306            parameters.clone(),
307            callback,
308        ));
309        Self {
310            definition: ToolDefinition {
311                name,
312                description,
313                parameters,
314            },
315            handler,
316            liveness: None,
317        }
318    }
319
320    /// Attach a liveness probe.
321    pub fn with_liveness<F>(mut self, is_live: F) -> Self
322    where
323        F: Fn() -> bool + WasmCompatSend + WasmCompatSync + 'static,
324    {
325        self.liveness = Some(Arc::new(is_live));
326        self
327    }
328
329    /// The tool's name.
330    pub fn name(&self) -> &ToolName {
331        &self.definition.name
332    }
333
334    /// The tool's definition.
335    pub fn definition(&self) -> ToolDefinition {
336        self.definition.clone()
337    }
338
339    /// The erased handler behind this definition.
340    pub fn handler(&self) -> &ErasedHandler {
341        &self.handler
342    }
343
344    /// Consumes the tool into its definition, handler, and optional liveness probe.
345    pub fn into_parts(self) -> (ToolDefinition, ErasedHandler, Option<LivenessFn>) {
346        (self.definition, self.handler, self.liveness)
347    }
348
349    /// Whether the tool's owner still serves it (`true` without a probe).
350    pub fn is_live(&self) -> bool {
351        self.liveness.as_ref().is_none_or(|probe| probe())
352    }
353
354    /// Run the tool inline with an empty context.
355    pub async fn execute(
356        &self,
357        arguments: serde_json::Value,
358    ) -> Result<ToolOutput, ToolExecutionError> {
359        let mut context = ToolContext::new();
360        self.execute_with(&mut context, arguments).await
361    }
362
363    /// Run the tool inline with isolated inbound values. A completed call
364    /// replaces only `context`'s result metadata, including when the tool
365    /// returns an error. Dropping the execution future leaves the caller's
366    /// context unchanged; it does not publish partial mutations.
367    pub async fn execute_with(
368        &self,
369        context: &mut ToolContext,
370        arguments: serde_json::Value,
371    ) -> Result<ToolOutput, ToolExecutionError> {
372        let published = PublishedContext::new();
373        let outcome = crate::serve::serve_inline_with(
374            &self.handler,
375            EffectKind::ToolCall {
376                name: self.definition.name.to_string(),
377                args: arguments.to_string(),
378            },
379            vec![
380                Arc::new(context.for_dispatch()),
381                published.clone() as Arc<dyn std::any::Any + Send + Sync>,
382            ],
383        )
384        .await;
385        match outcome {
386            Ok(Outcome::ToolResult { result }) => {
387                context.accept_dispatch_result(published.take().unwrap_or_default());
388                result.into_result()
389            }
390            Ok(other) => Err(ToolExecutionError::other(format!(
391                "tool handler answered with a {} outcome",
392                other.family()
393            ))),
394            Err(report) => Err(report.into()),
395        }
396    }
397}
398
399impl std::fmt::Debug for DynamicTool {
400    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
401        f.debug_struct("DynamicTool")
402            .field("name", &self.definition.name)
403            .finish_non_exhaustive()
404    }
405}
406
407/// A typed tool's [`Tool::NAME`] as a [`ToolName`].
408///
409/// An empty `NAME` fails the build where this is instantiated.
410///
411/// ```
412/// use rig_core::{message::ToolChoice, tool::{Tool, ToolContext, tool_name}};
413///
414/// struct Ping;
415///
416/// impl Tool for Ping {
417///     const NAME: &'static str = "ping";
418///     type Args = serde_json::Value;
419///     type Output = String;
420///     type Error = std::convert::Infallible;
421///
422///     fn description(&self) -> String {
423///         "Replies with pong.".to_string()
424///     }
425///
426///     fn parameters(&self) -> serde_json::Value {
427///         serde_json::json!({ "type": "object" })
428///     }
429///
430///     async fn call(
431///         &self,
432///         _context: &mut ToolContext,
433///         _args: Self::Args,
434///     ) -> Result<String, Self::Error> {
435///         Ok("pong".to_string())
436///     }
437/// }
438///
439/// let choice = ToolChoice::Specific { function_names: vec![tool_name::<Ping>()] };
440/// assert_eq!(tool_name::<Ping>(), "ping");
441/// # let _ = choice;
442/// ```
443pub fn tool_name<T: Tool>() -> ToolName {
444    const { assert!(!T::NAME.is_empty(), "Tool::NAME cannot be empty") };
445    ToolName::new_unchecked(T::NAME)
446}
447
448/// A tool's [`ToolDefinition`] from a typed tool.
449pub fn tool_definition<T: Tool>(tool: &T) -> ToolDefinition {
450    ToolDefinition {
451        name: tool_name::<T>(),
452        description: tool.description(),
453        parameters: tool.parameters(),
454    }
455}