Exploring MLX Swift: Getting Started with Tool Use
I have been waiting a while for the PR on tool use to get merged to finally publish this post. It is about how to get started with tool use in MLX Swift.
LLMs are powerful at generating text and answering questions based on their training data but they struggle with tasks that requires real-time information or interaction with external systems, like in iOS or macOS. This is where you can take advantage of "tool use" (also usually known as "function calling").
This is an introductory post that explains what tool use is, how to augment the LLMs with extra context and how to use a weather tool to fetch the current weather in Gurgaon, India.
Also, I would want to thank DePasqualeOrg for their contribution!
What is Tool Use?
Tool use allows an LLM to interact with external functions (tools) during its response generation. Instead of directly answering a question, the LLM can, based on the user's prompt, identify the need for a specific tool, form a request to that tool, and then incorporate the tool's output into its final response. For example, you can ask for the current weather, and the model will call a function named get_current_weather.
With tools, it can look up information and interact with other services, becoming way more useful than its current capabilities.
Defining a Tool in MLX Swift
Before an LLM can use a tool, you need to define it in a structured way that the LLM can understand. This involves creating a schema that describes the tool's name, purpose, and input parameters. MLX Swift uses a format similar to OpenAI's function calling API, making it straightforward to define your tools.
Let's break down the key components of a tool definition:
Step 1: Define Input and Output Structures
First, create Codable structures for the tool's input parameters and output:
struct WeatherInput: Codable {
let location: String
let unit: String?
}
struct WeatherOutput: Codable {
let temperature: Double
let conditions: String
}Step 2: Create the Tool
Now create a Tool instance with the input/output types, parameters, and handler:
import MLXLMCommon
let currentWeatherTool = Tool<WeatherInput, WeatherOutput>(
name: "get_current_weather",
description: "Get the current weather in a given location",
parameters: [
.required(
"location",
type: .string,
description: "The city and state, e.g. San Francisco, CA"
),
.optional(
"unit",
type: .string,
description: "The temperature unit (celsius or fahrenheit)"
),
]
) { input in
// This is where you'd actually fetch weather data
// For this example, we'll return mock data
let temperature = Double.random(in: 15...30)
let conditions = ["Sunny", "Cloudy", "Rainy", "Snowy"].randomElement()!
return WeatherOutput(temperature: temperature, conditions: conditions)
}The Tool structure takes:
- name: Identifier the LLM will use to call this tool
- description: Explains what the tool does (helps the LLM decide when to use it)
- parameters: Array of
ToolParameterdescribing required and optional inputs - handler: Closure that executes when the tool is called
Step 3: Using Tools with the LLM
To use tools with your LLM, pass them in the UserInput:
let chat: [Chat.Message] = [
.system("You are a helpful assistant"),
.user("What's the weather in Gurgaon, Haryana?")
]
let userInput = UserInput(
chat: chat,
tools: [currentWeatherTool.schema]
)
// Generate with the model
let modelContainer = try await LLMModelFactory.shared.loadContainer(
configuration: LLMRegistry.qwen3_4b_4bit
)
try await modelContainer.perform { context in
let lmInput = try await context.processor.prepare(input: userInput)
let stream = try MLXLMCommon.generate(
input: lmInput,
parameters: .init(),
context: context
)
for await batch in stream {
// Handle tool calls from the stream
if let toolCall = batch.toolCall {
let result = try await toolCall.execute(with: currentWeatherTool)
print("Tool result: \(result.toolResult)")
}
// Handle regular text output
if let chunk = batch.chunk {
print(chunk, terminator: "")
}
}
}Step 4: Handling Tool Calls
When the LLM decides to use a tool, it returns a ToolCall object. You need to execute the tool and provide the result back:
private func handleToolCall(_ toolCall: ToolCall) async throws -> String {
switch toolCall.function.name {
case currentWeatherTool.name:
let result = try await toolCall.execute(with: currentWeatherTool)
return result.toolResult
default:
return "Unknown tool: \(toolCall.function.name)"
}
}The execute(with:) method:
- Parses the tool call arguments into your
Inputtype - Calls your handler closure
- Serializes the
Outputback to JSON - Returns a
ToolResultthat can be sent back to the LLM
Continuing the Conversation
After the tool executes and returns data, you can send the result back to the LLM for a natural response:
private func generate(prompt: String, toolResult: String? = nil) async {
var chat: [Chat.Message] = [
.system("You are a helpful assistant"),
.user(prompt),
]
// If we have a tool result, append it to the conversation
if let toolResult {
chat.append(.tool(toolResult))
}
let userInput = UserInput(
chat: chat,
tools: [currentWeatherTool.schema]
)
let modelContainer = try await LLMModelFactory.shared.loadContainer(
configuration: LLMRegistry.qwen3_4b_4bit
)
try await modelContainer.perform { context in
let lmInput = try await context.processor.prepare(input: userInput)
let stream = try MLXLMCommon.generate(
input: lmInput,
parameters: .init(),
context: context
)
for await batch in stream {
if let toolCall = batch.toolCall {
// Execute the tool and continue with its result
let result = try await toolCall.execute(with: currentWeatherTool)
await generate(prompt: prompt, toolResult: result.toolResult)
return
}
if let chunk = batch.chunk {
print(chunk, terminator: "")
}
}
}
}The key steps are:
- Start with user's question: "What's the weather in Gurgaon?"
- LLM decides to call
get_current_weathertool - Execute the tool handler and get weather data
- Pass tool result back to LLM with
.tool()message - LLM generates natural language response using the data
Output
You can now type in a prompt like "What's the weather in Gurgaon, Haryana?" and see the LLM call the weather tool, retrieve the data, and provide a response!
Moving Forward
This was a starter post on tool calling and there are many ways you can utilize it on-device for different use cases. You can take it forward with:
- Exploring more tools: Create tools for other tasks (e.g., calendar access, health data, web searches)
- Exploring prompts: Experiment with different prompts to guide the LLM's behavior. This is significant when playing with other models like Hermes 3 Llama 3.2 that has a different chat template for tool calling.
- Exploring interactions: This was a one shot tool calling. You can pass the tool responses to each other and multi-turn conversations as well.
Check out the LLMEval example app to see a complete implementation with multiple tools including weather, math, and time tools.
Happy MLXing!
Post Topics
Explore more in these categories: