Skip to main content

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