aimux_core/
tool.rs

1//! Tool / function-calling types.
2
3use std::collections::HashMap;
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use ts_rs::TS;
8
9use crate::error::AiMuxError;
10use crate::types::ProviderMetadata;
11
12/// A tool definition passed to the model in `CallOptions.tools`.
13#[derive(Debug, Clone, Serialize, Deserialize, TS)]
14#[ts(export)]
15pub struct FunctionTool {
16    /// Unique tool name.
17    pub name: String,
18    /// Human-readable description.
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub description: Option<String>,
21    /// JSON Schema describing the tool's parameters.
22    pub input_schema: Value,
23    /// Whether the tool supports strict schema enforcement (OpenAI strict mode).
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub strict: Option<bool>,
26    /// Provider-specific options, keyed by provider name
27    /// (e.g. `{"anthropic": {"eagerInputStreaming": true}}`). Aligned with the
28    /// V4 `providerOptions` field on function tools.
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub provider_options: Option<HashMap<String, Value>>,
31    /// Example inputs for the tool (V4 `inputExamples`), used by some providers
32    /// to emit `input_examples` in the request body.
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub input_examples: Option<Vec<Value>>,
35}
36
37impl FunctionTool {
38    /// Create a `FunctionTool` with the required fields and all optional fields
39    /// set to `None`. Combined with struct-update syntax this avoids spelling
40    /// out every `None` field in tests.
41    pub fn new(name: impl Into<String>, input_schema: Value) -> Self {
42        FunctionTool {
43            name: name.into(),
44            description: None,
45            input_schema,
46            strict: None,
47            provider_options: None,
48            input_examples: None,
49        }
50    }
51
52    /// Builder-style setter for `description`.
53    #[must_use]
54    pub fn with_description(mut self, description: impl Into<String>) -> Self {
55        self.description = Some(description.into());
56        self
57    }
58
59    /// Builder-style setter for `strict`.
60    #[must_use]
61    pub fn with_strict(mut self, strict: bool) -> Self {
62        self.strict = Some(strict);
63        self
64    }
65}
66
67/// A provider-defined tool (e.g. `anthropic.web_search_20250305`,
68/// `google.googleSearch`). Mirrors the V4 `LanguageModelV4ProviderTool`.
69///
70/// Provider tools are specific to a certain provider. The input and output
71/// schemas are defined by the provider, and some are executed on the provider's
72/// servers (e.g. web search, code execution).
73#[derive(Debug, Clone, Serialize, Deserialize, TS)]
74#[ts(export)]
75pub struct ProviderTool {
76    /// The ID of the tool, following the format `<provider-id>.<unique-tool-name>`.
77    pub id: String,
78    /// The name of the tool, unique within this model call.
79    pub name: String,
80    /// The arguments for configuring the tool. Must match the expected
81    /// arguments defined by the provider for this tool.
82    pub args: Value,
83}
84
85/// A tool that can be either a user-defined function tool or a
86/// provider-defined tool. Mirrors the V4
87/// `Array<LanguageModelV4FunctionTool | LanguageModelV4ProviderTool>`.
88#[derive(Debug, Clone, Serialize, Deserialize, TS)]
89#[serde(tag = "type", rename_all = "snake_case")]
90#[ts(export)]
91pub enum Tool {
92    /// A user-defined function tool.
93    Function(FunctionTool),
94    /// A provider-defined tool.
95    Provider(ProviderTool),
96}
97
98impl From<FunctionTool> for Tool {
99    fn from(ft: FunctionTool) -> Self {
100        Tool::Function(ft)
101    }
102}
103
104/// A tool call requested by the model.
105#[derive(Debug, Clone, Serialize, Deserialize, TS)]
106#[ts(export)]
107pub struct ToolCall {
108    /// Provider-assigned call id.
109    pub tool_call_id: String,
110    /// Tool name.
111    pub tool_name: String,
112    /// Parsed arguments, or the original string when `invalid` is true and the
113    /// provider input was not valid JSON.
114    pub input: Value,
115    /// Whether the tool call will be executed by the provider.
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub provider_executed: Option<bool>,
118    /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools).
119    #[serde(default, skip_serializing_if = "Option::is_none")]
120    pub dynamic: Option<bool>,
121    /// Provider-assigned thought signature (e.g. Google Gemini
122    /// `thoughtSignature`). Must be echoed back verbatim on the follow-up turn
123    /// when the tool result is sent; thinking models reject the request
124    /// otherwise.
125    #[serde(default, skip_serializing_if = "Option::is_none")]
126    pub thought_signature: Option<String>,
127    /// Additional provider-specific metadata associated with this call.
128    #[serde(default, skip_serializing_if = "Option::is_none")]
129    pub provider_metadata: Option<ProviderMetadata>,
130    /// Set when lookup, JSON parsing, or schema validation still failed after
131    /// the optional repair attempt.
132    #[serde(default, skip_serializing_if = "Option::is_none")]
133    pub invalid: Option<bool>,
134    /// Typed failure associated with an invalid tool call.
135    #[serde(default, skip_serializing_if = "Option::is_none")]
136    pub error: Option<AiMuxError>,
137}
138
139/// The result of executing a tool call, to be sent back to the model.
140#[derive(Debug, Clone, Serialize, Deserialize, TS)]
141#[ts(export)]
142pub struct ToolResult {
143    /// Must match the corresponding `ToolCall::tool_call_id`.
144    pub tool_call_id: String,
145    /// The tool's output (usually a JSON-serializable value or plain text).
146    pub result: Value,
147    /// Whether the result is an error or error message.
148    #[serde(default, skip_serializing_if = "Option::is_none")]
149    pub is_error: Option<bool>,
150    /// Whether the result is preliminary (replaces prior, e.g. image previews).
151    #[serde(default, skip_serializing_if = "Option::is_none")]
152    pub preliminary: Option<bool>,
153}
154
155/// How the model should choose tools.
156///
157/// Wire format aligns with Vercel AI SDK `toolChoice`:
158/// `"auto" | "none" | "required" | { "type": "tool", "toolName": "..." }`.
159/// This is a mixed tagged/untagged shape (unit variants as bare strings,
160/// `Tool` variant as a tagged object), so serde is implemented by hand.
161#[derive(Debug, Clone, Default, PartialEq, TS)]
162#[ts(
163    export,
164    type = "\"auto\" | \"none\" | \"required\" | { type: \"tool\"; toolName: string }"
165)]
166pub enum ToolChoice {
167    /// Model decides whether to call a tool (default).
168    #[default]
169    Auto,
170    /// Model must not call any tools.
171    None,
172    /// Model must call at least one tool.
173    Required,
174    /// Model must call the specified tool.
175    Tool { tool_name: String },
176}
177
178impl Serialize for ToolChoice {
179    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
180    where
181        S: serde::Serializer,
182    {
183        use serde::ser::SerializeStruct;
184
185        match self {
186            ToolChoice::Auto => serializer.serialize_str("auto"),
187            ToolChoice::None => serializer.serialize_str("none"),
188            ToolChoice::Required => serializer.serialize_str("required"),
189            ToolChoice::Tool { tool_name } => {
190                let mut s = serializer.serialize_struct("ToolChoice", 2)?;
191                s.serialize_field("type", "tool")?;
192                s.serialize_field("toolName", tool_name)?;
193                s.end()
194            }
195        }
196    }
197}
198
199impl<'de> Deserialize<'de> for ToolChoice {
200    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
201    where
202        D: serde::Deserializer<'de>,
203    {
204        use serde::de::Error;
205
206        #[derive(Deserialize)]
207        #[serde(untagged)]
208        enum Repr {
209            Str(String),
210            Tool {
211                #[serde(rename = "type")]
212                typ: String,
213                #[serde(rename = "toolName")]
214                tool_name: Option<String>,
215            },
216        }
217
218        match Repr::deserialize(deserializer)? {
219            Repr::Str(s) => match s.as_str() {
220                "auto" => Ok(ToolChoice::Auto),
221                "none" => Ok(ToolChoice::None),
222                "required" => Ok(ToolChoice::Required),
223                other => Err(D::Error::custom(format!(
224                    "unknown toolChoice '{other}' (expected auto/none/required or {{type:\"tool\",toolName}})"
225                ))),
226            },
227            Repr::Tool { typ, tool_name } => {
228                if typ != "tool" {
229                    return Err(D::Error::custom(format!(
230                        "unknown toolChoice type '{typ}' (expected \"tool\")"
231                    )));
232                }
233                let tool_name = tool_name.ok_or_else(|| {
234                    D::Error::custom("toolChoice {type:\"tool\"} requires a toolName field")
235                })?;
236                Ok(ToolChoice::Tool { tool_name })
237            }
238        }
239    }
240}
aimux_core/tool.rs

aimux_core/
tool.rs
