Skip to main content

llm_tool/types/
definition.rs

1//! Tool execution effect classification and [`ToolDefinition`] metadata.
2
3use alloc::{string::String, sync::Arc};
4
5use serde::{Deserialize, Serialize};
6
7use super::OutputSanitization;
8use crate::compat::OnceLock;
9
10/// Effect classification for a tool.
11#[derive(
12    Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, schemars::JsonSchema,
13)]
14#[serde(rename_all = "snake_case")]
15pub enum ToolEffect {
16    /// Read-only, side-effect free invocation. Safe for concurrent wave execution.
17    ReadOnly,
18    /// Mutates system state (default). Safe for normal execution.
19    #[default]
20    Mutating,
21    /// Destructive operation (e.g. deleting files, wiping data, dropping resources).
22    /// Requires explicit confirmation before parallel batch dispatch.
23    Destructive,
24}
25
26impl ToolEffect {
27    /// Returns `true` if the tool effect is read-only / idempotent.
28    #[must_use]
29    pub const fn is_read_only(self) -> bool {
30        matches!(self, Self::ReadOnly)
31    }
32
33    /// Returns `true` if the tool effect is destructive.
34    #[must_use]
35    pub const fn is_destructive(self) -> bool {
36        matches!(self, Self::Destructive)
37    }
38}
39
40/// Describes a custom tool that can be registered with an agent.
41#[non_exhaustive]
42#[derive(Debug, Serialize)]
43pub struct ToolDefinition {
44    /// Unique tool name (e.g. `"flash_device"`).
45    pub name: String,
46    /// Human-readable description shown to the model.
47    pub description: String,
48    /// JSON Schema describing the tool's parameters.
49    pub parameter_schema: serde_json::Value,
50    /// Execution effect classification of the tool.
51    pub effect: ToolEffect,
52    /// Automatic prompt-injection sanitization policy applied to the tool's output.
53    #[serde(default, skip_serializing_if = "OutputSanitization::is_none")]
54    pub output_sanitization: OutputSanitization,
55    /// Pre-serialized compact JSON Schema string cached lazily on first access.
56    #[serde(skip)]
57    parameter_schema_json: OnceLock<Arc<str>>,
58}
59
60impl Clone for ToolDefinition {
61    fn clone(&self) -> Self {
62        Self {
63            name: self.name.clone(),
64            description: self.description.clone(),
65            parameter_schema: self.parameter_schema.clone(),
66            effect: self.effect,
67            output_sanitization: self.output_sanitization,
68            parameter_schema_json: OnceLock::new(),
69        }
70    }
71}
72
73impl PartialEq for ToolDefinition {
74    fn eq(&self, other: &Self) -> bool {
75        self.name == other.name
76            && self.description == other.description
77            && self.parameter_schema == other.parameter_schema
78            && self.effect == other.effect
79            && self.output_sanitization == other.output_sanitization
80    }
81}
82
83impl Eq for ToolDefinition {}
84
85impl<'de> Deserialize<'de> for ToolDefinition {
86    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
87    where
88        D: serde::Deserializer<'de>,
89    {
90        #[derive(Deserialize)]
91        struct RawDef {
92            name: String,
93            description: String,
94            parameter_schema: serde_json::Value,
95            #[serde(default)]
96            effect: Option<ToolEffect>,
97            #[serde(default)]
98            idempotent: Option<bool>,
99            #[serde(default)]
100            output_sanitization: OutputSanitization,
101        }
102
103        let raw = RawDef::deserialize(deserializer)?;
104        let effect = match (raw.effect, raw.idempotent) {
105            (Some(eff), _) => eff,
106            (None, Some(true)) => ToolEffect::ReadOnly,
107            (None, Some(false) | None) => ToolEffect::Mutating,
108        };
109
110        Ok(Self {
111            name: raw.name,
112            description: raw.description,
113            parameter_schema: raw.parameter_schema,
114            effect,
115            output_sanitization: raw.output_sanitization,
116            parameter_schema_json: OnceLock::new(),
117        })
118    }
119}
120
121impl ToolDefinition {
122    /// Construct a [`ToolDefinition`] with default [`ToolEffect::Mutating`]
123    /// and [`OutputSanitization::None`].
124    #[must_use]
125    pub fn new(
126        name: impl Into<String>,
127        description: impl Into<String>,
128        parameter_schema: serde_json::Value,
129    ) -> Self {
130        Self::with_effect(name, description, parameter_schema, ToolEffect::Mutating)
131    }
132
133    /// Construct a [`ToolDefinition`] with an explicit [`ToolEffect`].
134    #[must_use]
135    pub fn with_effect(
136        name: impl Into<String>,
137        description: impl Into<String>,
138        parameter_schema: serde_json::Value,
139        effect: ToolEffect,
140    ) -> Self {
141        Self {
142            name: name.into(),
143            description: description.into(),
144            parameter_schema,
145            effect,
146            output_sanitization: OutputSanitization::None,
147            parameter_schema_json: OnceLock::new(),
148        }
149    }
150
151    /// Attach an [`OutputSanitization`] policy to this [`ToolDefinition`]. Chainable.
152    #[must_use]
153    pub const fn with_output_sanitization(mut self, sanitization: OutputSanitization) -> Self {
154        self.output_sanitization = sanitization;
155        self
156    }
157
158    /// Update `parameter_schema` and invalidate any cached JSON serialization.
159    pub fn set_parameter_schema(&mut self, parameter_schema: serde_json::Value) {
160        self.parameter_schema = parameter_schema;
161        self.parameter_schema_json = OnceLock::new();
162    }
163
164    /// Borrow `parameter_schema` mutably and invalidate any cached JSON serialization.
165    pub fn parameter_schema_mut(&mut self) -> &mut serde_json::Value {
166        self.parameter_schema_json = OnceLock::new();
167        &mut self.parameter_schema
168    }
169
170    /// Return the pre-serialized JSON schema string (`&str`), lazily derived
171    /// from `parameter_schema`.
172    #[must_use]
173    pub fn parameter_schema_json(&self) -> &str {
174        self.parameter_schema_json_arc()
175    }
176
177    /// Return the pre-serialized JSON schema string (`&Arc<str>`), lazily derived
178    /// from `parameter_schema`.
179    #[must_use]
180    pub fn parameter_schema_json_arc(&self) -> &Arc<str> {
181        self.parameter_schema_json.get_or_init(|| {
182            match serde_json::to_string(&self.parameter_schema) {
183                Ok(s) => Arc::from(s),
184                Err(err) => {
185                    tracing::warn!(error = %err, tool = %self.name, "Failed to serialize parameter_schema to JSON");
186                    Arc::from("{}")
187                }
188            }
189        })
190    }
191
192    /// Whether repeated invocations with identical arguments are side-effect free / idempotent.
193    ///
194    /// Preserved for backward compatibility.
195    #[must_use]
196    pub const fn is_idempotent(&self) -> bool {
197        matches!(self.effect, ToolEffect::ReadOnly)
198    }
199}