Add support for Mistral3 tool calling

Author: atdrendel

Libraries/MLXLMCommon/Evaluate.swift

@@ -1080,6 +1080,8 @@ private func generateLoopTask<Handler: TokenLoopHandler>(
                 }
             }
 
+            handler.onGenerationEnd(emit: continuation.yield)
+
             let now = Date.timeIntervalSinceReferenceDate
             let generateTime = now - start
 
@@ -1292,6 +1294,11 @@ private protocol TokenLoopHandler: Sendable {
         emit: (sending Output) -> AsyncStream<Output>.Continuation.YieldResult
     ) -> Bool
 
+    /// Called after the token loop finishes, before the info event.
+    mutating func onGenerationEnd(
+        emit: (sending Output) -> AsyncStream<Output>.Continuation.YieldResult
+    )
+
     func infoEvent(_ info: GenerateCompletionInfo) -> Output
 }
 
@@ -1337,6 +1344,16 @@ private struct TextToolTokenLoopHandler: TokenLoopHandler, @unchecked Sendable {
         true
     }
 
+    mutating func onGenerationEnd(
+        emit: (sending Generation) -> AsyncStream<Generation>.Continuation.YieldResult
+    ) {
+        toolCallProcessor.flush()
+        while let toolCall = toolCallProcessor.toolCalls.first {
+            toolCallProcessor.toolCalls.removeFirst()
+            _ = emit(.toolCall(toolCall))
+        }
+    }
+
     func infoEvent(_ info: GenerateCompletionInfo) -> Generation {
         .info(info)
     }
@@ -1365,6 +1382,10 @@ private struct RawTokenLoopHandler: TokenLoopHandler {
         return true
     }
 
+    mutating func onGenerationEnd(
+        emit: (sending TokenGeneration) -> AsyncStream<TokenGeneration>.Continuation.YieldResult
+    ) {}
+
     func infoEvent(_ info: GenerateCompletionInfo) -> TokenGeneration {
         .info(info)
     }

Libraries/MLXLMCommon/Tool/Parsers/MistralToolCallParser.swift

@@ -0,0 +1,66 @@
+// Copyright © 2025 Apple Inc.
+
+import Foundation
+
+/// Parser for Mistral V13 tool call format: `[TOOL_CALLS]name[ARGS]{"json_args"}`
+///
+/// This format is used by Mistral3/Ministral-3 2512 models and Devstral 2.
+/// The special tokens `[TOOL_CALLS]` (token ID 9) and `[ARGS]` (ID 32) are used
+/// as delimiters. Multiple tool calls use repeated `[TOOL_CALLS]` tokens.
+///
+/// Also handles the older V11 format which includes an optional `[CALL_ID]`
+/// between the function name and `[ARGS]` (V13 does not use `[CALL_ID]`).
+///
+/// Examples:
+/// - `[TOOL_CALLS]get_weather[ARGS]{"location": "Tokyo"}`
+/// - `[TOOL_CALLS]fn1[ARGS]{...}[TOOL_CALLS]fn2[ARGS]{...}` (multiple calls)
+///
+/// The end tag is `</s>` (EOS token). Since stop tokens are intercepted at the
+/// token ID level before detokenization, the EOS text never reaches the processor
+/// — tool calls are extracted via `ToolCallProcessor.flush()` at generation end.
+public struct MistralToolCallParser: ToolCallParser, Sendable {
+    public let startTag: String? = "[TOOL_CALLS]"
+    public let endTag: String? = "</s>"
+
+    public init() {}
+
+    public func parse(content: String, tools: [[String: any Sendable]]?) -> ToolCall? {
+        var text = content
+
+        // Strip [TOOL_CALLS] prefix if present
+        if let range = text.range(of: "[TOOL_CALLS]") {
+            text = String(text[range.upperBound...])
+        }
+
+        text = text.trimmingCharacters(in: .whitespacesAndNewlines)
+
+        // Split on [ARGS] to get function name and arguments
+        guard let argsRange = text.range(of: "[ARGS]") else {
+            return nil
+        }
+
+        var namePart = String(text[..<argsRange.lowerBound])
+            .trimmingCharacters(in: .whitespacesAndNewlines)
+        let argsPart = String(text[argsRange.upperBound...])
+            .trimmingCharacters(in: .whitespacesAndNewlines)
+
+        // Handle optional [CALL_ID] between name and [ARGS]
+        if let callIdRange = namePart.range(of: "[CALL_ID]") {
+            namePart = String(namePart[..<callIdRange.lowerBound])
+                .trimmingCharacters(in: .whitespacesAndNewlines)
+        }
+
+        guard !namePart.isEmpty else { return nil }
+
+        // Parse arguments as JSON using deserialize from ParserUtilities
+        let arguments = deserialize(argsPart)
+
+        guard let argsDict = arguments as? [String: any Sendable] else {
+            return nil
+        }
+
+        return ToolCall(
+            function: ToolCall.Function(name: namePart, arguments: argsDict)
+        )
+    }
+}

Libraries/MLXLMCommon/Tool/ToolCallFormat.swift

@@ -66,6 +66,10 @@ public enum ToolCallFormat: String, Sendable, Codable, CaseIterable {
     /// Example: `<invoke name="f"><parameter name="k">v</parameter></invoke>`
     case minimaxM2 = "minimax_m2"
 
+    /// Mistral V11+ format with [TOOL_CALLS] and [ARGS] delimiters.
+    /// Example: `[TOOL_CALLS]get_weather [ARGS]{"location": "Tokyo"}`
+    case mistral
+
     // MARK: - Factory Methods
 
     /// Create the appropriate parser for this format.
@@ -87,6 +91,8 @@ public enum ToolCallFormat: String, Sendable, Codable, CaseIterable {
             return KimiK2ToolCallParser()
         case .minimaxM2:
             return MiniMaxM2ToolCallParser()
+        case .mistral:
+            return MistralToolCallParser()
         }
     }
 
@@ -115,6 +121,11 @@ public enum ToolCallFormat: String, Sendable, Codable, CaseIterable {
             return .gemma
         }
 
+        // Mistral3 family (mistral3, mistral3_text, etc.)
+        if type.hasPrefix("mistral3") {
+            return .mistral
+        }
+
         return nil
     }
 }

Libraries/MLXLMCommon/Tool/ToolCallProcessor.swift

@@ -67,6 +67,46 @@ public class ToolCallProcessor {
 
     // MARK: - Public Methods
 
+    /// Flush any buffered content, attempting to parse it as tool call(s).
+    ///
+    /// Call this when generation ends (e.g., on EOS token) to handle formats
+    /// whose end tag is never delivered as text (e.g., Mistral where `</s>`
+    /// is intercepted at the token ID level).
+    ///
+    /// For formats with end tags that appear in the text stream, the buffer
+    /// will already be empty at generation end, making this a no-op.
+    public func flush() {
+        guard state == .collectingToolCall || state == .potentialToolCall else { return }
+        guard !toolCallBuffer.isEmpty else {
+            state = .normal
+            return
+        }
+
+        if let startTag = parser.startTag {
+            // Split buffer into individual tool call segments to handle
+            // multiple tool calls (e.g., Mistral V13):
+            // "[TOOL_CALLS]fn1[ARGS]{...}[TOOL_CALLS]fn2[ARGS]{...}"
+            // → ["fn1[ARGS]{...}", "fn2[ARGS]{...}"]
+            let segments =
+                toolCallBuffer
+                .components(separatedBy: startTag)
+                .filter { !$0.isEmpty }
+
+            for segment in segments {
+                if let toolCall = parser.parse(content: segment, tools: tools) {
+                    toolCalls.append(toolCall)
+                }
+            }
+        } else {
+            if let toolCall = parser.parse(content: toolCallBuffer, tools: tools) {
+                toolCalls.append(toolCall)
+            }
+        }
+
+        toolCallBuffer = ""
+        state = .normal
+    }
+
     /// Process a generated text chunk and extract any tool call content.
     /// - Parameter chunk: The text chunk to process
     /// - Returns: Regular text that should be displayed (non-tool call content), or `nil` if buffering

Tests/MLXLMIntegrationTests/ToolCallIntegrationTests.swift

@@ -213,17 +213,15 @@ public class ToolCallIntegrationTests: XCTestCase {
 
     // MARK: - Mistral3 Tests
 
-    func testMistral3ToolCallFormatDefaultsToJSON() async throws {
+    func testMistral3ToolCallFormatAutoDetection() async throws {
         guard let container = Self.mistral3Container else {
             throw XCTSkip("Mistral3 model not available")
         }
 
         let config = await container.configuration
-        // Mistral3 uses the default JSON tool call format (infer returns nil)
-        let format = config.toolCallFormat ?? .json
         XCTAssertEqual(
-            format, .json,
-            "Mistral3 model should use default .json tool call format"
+            config.toolCallFormat, .mistral,
+            "Mistral3 model should auto-detect .mistral tool call format"
         )
     }
 
@@ -264,6 +262,69 @@ public class ToolCallIntegrationTests: XCTestCase {
         }
     }
 
+    func testMistral3MultipleToolCallGeneration() async throws {
+        guard let container = Self.mistral3Container else {
+            throw XCTSkip("Mistral3 model not available")
+        }
+
+        let multiToolSchema: [[String: any Sendable]] =
+            Self.weatherToolSchema + [
+                [
+                    "type": "function",
+                    "function": [
+                        "name": "get_time",
+                        "description": "Get the current time in a given timezone",
+                        "parameters": [
+                            "type": "object",
+                            "properties": [
+                                "timezone": [
+                                    "type": "string",
+                                    "description":
+                                        "The timezone, e.g. America/New_York, Asia/Tokyo",
+                                ] as [String: any Sendable]
+                            ] as [String: any Sendable],
+                            "required": ["timezone"],
+                        ] as [String: any Sendable],
+                    ] as [String: any Sendable],
+                ]
+            ]
+
+        let input = UserInput(
+            chat: [
+                .system(
+                    "You are a helpful assistant with access to tools. Always use the available tools to answer questions. Call multiple tools in parallel when needed."
+                ),
+                .user(
+                    "What's the weather in Tokyo and what time is it there?"
+                ),
+            ],
+            tools: multiToolSchema
+        )
+
+        let (result, toolCalls) = try await generateWithTools(
+            container: container,
+            input: input,
+            maxTokens: 150
+        )
+
+        print("Mistral3 Output: \(result)")
+        print("Mistral3 Calls: \(toolCalls)")
+
+        // Verify all returned tool calls have valid names from our schema
+        let validNames: Set<String> = ["get_weather", "get_time"]
+        for toolCall in toolCalls {
+            XCTAssertTrue(
+                validNames.contains(toolCall.function.name),
+                "Unexpected tool call: \(toolCall.function.name)"
+            )
+        }
+
+        // If the model made multiple calls, verify we got more than one
+        if toolCalls.count > 1 {
+            print("Successfully parsed \(toolCalls.count) tool calls from Mistral3")
+        }
+    }
+
     // MARK: - Helper Methods
 
     /// Generate text and collect any tool calls