From ceaed6dec902d36b2d7fc98193715663d098bb9d Mon Sep 17 00:00:00 2001 From: Claude Opus 5 Date: Mon, 3 Aug 2026 11:57:50 -0400 Subject: [PATCH 1/2] Add tool calling to the Core ML provider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Core ML provider forwarded tool specs into the chat template and then ignored the result: generated output was never inspected for tool calls, nothing was executed, and nothing reached the transcript or the tool execution delegate. Both `respond` and `streamResponse` now detect tool calls, execute them, feed the results back as another turn, and loop until the model answers without requesting tools. Core ML models have no structured tool-call channel. swift-transformers renders tool specs into the prompt through the tokenizer's Jinja chat template and stops there — it ships no tool-call parser — so the model's request comes back as ordinary generated text in whatever format its chat template taught it. `CoreMLToolCallParser` therefore supports an explicit, closed set of formats: Hermes/Qwen `` tags, Mistral `[TOOL_CALLS]`, Llama 3.x `<|python_tag|>`, fenced JSON blocks, and bare JSON. The last two are ambiguous with a model simply answering in JSON, so they only count as a tool call when the name matches a tool the session actually has. Observable behavior matches the Anthropic baseline. Tool calls go through the session's `toolExecutionDelegate`, including `.stop` and `.provideOutput`. During streaming, tool activity lands in the transcript as it happens, with `.toolCalls` always appended before the corresponding `.toolOutput`, and a `.stop` decision appends the tool calls, finishes the stream, and executes nothing. Text growth goes through `growStreamingTranscript`, and tool-call markup is held back mid-stream so it never reaches the transcript or the caller. Repeated-tool-call-loop protection and an iteration cap mirror the MLX provider. A caller-supplied `toolsHandler` still wins and is still called exactly as before. When no handler is supplied, the session's tools are now converted with the conventional OpenAI-style function schema that Hugging Face chat templates expect, instead of being silently dropped. This also fixes the free-form generation config, which tool calling depends on. `GenerationConfig.maxLength` defaults to 20 and `generate` stops as soon as the total sequence reaches it, so with the prompt unaccounted for it generated nothing at all for any realistic prompt. `eosTokenId` likewise had no default, so generation never stopped early and decoded past the end of the turn. The constrained-decoding path in this file already set both; the plain path now does the same. Co-Authored-By: Claude Opus 5 (1M context) --- .../Models/CoreMLLanguageModel.swift | 967 ++++++++++++++++-- 1 file changed, 891 insertions(+), 76 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift b/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift index 221119ce..9d311ae7 100644 --- a/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift @@ -95,45 +95,101 @@ ) } - // Convert AnyLanguageModel GenerationOptions to swift-transformers GenerationConfig - let generationConfig = toGenerationConfig(options) + let toolSpecs = resolvedToolSpecs(for: session) + let toolNames = Set(session.tools.map(\.name)) + var promptSource = makePromptSource(session: session, prompt: prompt) + + var visibleChunks: [String] = [] + var allEntries: [Transcript.Entry] = [] + var toolIteration = 0 + var previousToolCallSignature: String? + + // Generate, then keep looping for as long as the model answers with tool calls. + while true { + let tokens = try encodePrompt(promptSource, toolSpecs: toolSpecs) + let generationConfig = toGenerationConfig(options, promptTokenCount: tokens.count) + + // Reset model state for new generation + await model.resetState() + + let outputTokens = await model.generate( + config: generationConfig, + tokens: tokens, + model: model.callAsFunction + ) + let assistantText = decodeAssistantText(from: outputTokens, promptTokenCount: tokens.count) - let tokens: [Int] - if let chatTemplateHandler = chatTemplateHandler { - // Use chat template handler with optional tools - let messages = chatTemplateHandler(session.instructions, prompt) - let toolSpecs: [ToolSpec]? = toolsHandler?(session.tools) - tokens = try tokenizer.applyChatTemplate(messages: messages, tools: toolSpecs) - } else { - // Fall back to direct tokenizer encoding - tokens = tokenizer.encode(text: prompt.description) - } + // Tool calling requires the chat-template path: there is no other way to + // render tool specs into the prompt or feed tool results back to the model. + guard case .chat(var messages) = promptSource, !session.tools.isEmpty else { + visibleChunks.append(assistantText) + break + } - // Reset model state for new generation - await model.resetState() + let parsed = CoreMLToolCallParser.parse(assistantText, knownToolNames: toolNames) + visibleChunks.append(parsed.visibleText) - let outputTokens = await model.generate( - config: generationConfig, - tokens: tokens, - model: model.callAsFunction - ) + if !parsed.calls.isEmpty { + toolIteration += 1 + if toolIteration > Self.maximumToolIterations { + allEntries.append( + .toolCalls(Transcript.ToolCalls(makeTranscriptToolCalls(from: parsed.calls))) + ) + throw Self.maxToolIterationsExceededError(limit: Self.maximumToolIterations) + } - // Strip the prompt at the token level to avoid issues with - // normalization or whitespace differences in decoded strings - let assistantTokenSlice: ArraySlice - if outputTokens.count >= tokens.count { - assistantTokenSlice = outputTokens.dropFirst(tokens.count) - } else { - // Fallback: if the model did not echo the full prompt, - // treat the entire output as assistant tokens - assistantTokenSlice = outputTokens[outputTokens.indices] + // Guard against a model that keeps asking for the exact same tool call. + let signature = CoreMLToolCallParser.signature(for: parsed.calls) + if signature == previousToolCallSignature { + allEntries.append( + .toolCalls(Transcript.ToolCalls(makeTranscriptToolCalls(from: parsed.calls))) + ) + throw Self.repeatedToolCallLoopError() + } + previousToolCallSignature = signature + + let transcriptCalls = makeTranscriptToolCalls(from: parsed.calls) + let resolution = try await resolveToolCalls(transcriptCalls, session: session) + switch resolution { + case .stop(let calls): + if !calls.isEmpty { + allEntries.append(.toolCalls(Transcript.ToolCalls(calls))) + } + return LanguageModelSession.Response( + content: "" as! Content, + rawContent: GeneratedContent(""), + transcriptEntries: ArraySlice(allEntries) + ) + case .invocations(let invocations): + if !invocations.isEmpty { + allEntries.append(.toolCalls(Transcript.ToolCalls(invocations.map(\.call)))) + + messages.append( + assistantToolCallMessage( + text: parsed.visibleText, + transcriptCalls: transcriptCalls, + parsedCalls: parsed.calls + ) + ) + for invocation in invocations { + allEntries.append(.toolOutput(invocation.output)) + messages.append(toolResultMessage(for: invocation.output)) + } + + promptSource = .chat(messages) + continue + } + } + } + + break } - let assistantText = tokenizer.decode(tokens: Array(assistantTokenSlice)) + let assistantText = visibleChunks.joined() return LanguageModelSession.Response( content: assistantText as! Content, rawContent: GeneratedContent(assistantText), - transcriptEntries: ArraySlice([]) + transcriptEntries: ArraySlice(allEntries) ) } @@ -166,63 +222,152 @@ ) } - // Convert AnyLanguageModel GenerationOptions to swift-transformers GenerationConfig - let generationConfig = toGenerationConfig(options) - // Transform the generation into ResponseStream snapshots let stream: AsyncThrowingStream.Snapshot, any Error> = .init { @Sendable continuation in let task = Task { do { - let tokens: [Int] - if let chatTemplateHandler = chatTemplateHandler { - // Use chat template handler with optional tools - let messages = chatTemplateHandler(session.instructions, prompt) - let toolSpecs: [ToolSpec]? = toolsHandler?(session.tools) - tokens = try tokenizer.applyChatTemplate(messages: messages, tools: toolSpecs) - } else { - // Fall back to direct tokenizer encoding - tokens = tokenizer.encode(text: prompt.description) - } + let toolSpecs = resolvedToolSpecs(for: session) + let toolNames = Set(session.tools.map(\.name)) + let toolsEnabled = !session.tools.isEmpty + var promptSource = makePromptSource(session: session, prompt: prompt) - await model.resetState() - - let promptTokenCount = tokens.count - var accumulatedText = "" - - _ = await model.generate( - config: generationConfig, - tokens: tokens, - model: model.callAsFunction - ) { tokenIds in - let assistantTokenSlice: ArraySlice - if tokenIds.count >= promptTokenCount { - assistantTokenSlice = tokenIds.dropFirst(promptTokenCount) - } else { - assistantTokenSlice = tokenIds[tokenIds.indices] - } - let assistantText = tokenizer.decode(tokens: Array(assistantTokenSlice)) - - // Compute delta vs accumulated text and yield - if assistantText.count >= accumulatedText.count, - assistantText.hasPrefix(accumulatedText) - { - let startIdx = assistantText.index( - assistantText.startIndex, - offsetBy: accumulatedText.count + var toolIteration = 0 + var previousToolCallSignature: String? + + // Generate, then keep looping for as long as the model answers with tool calls. + while true { + if Task.isCancelled { break } + + let tokens = try encodePrompt(promptSource, toolSpecs: toolSpecs) + let generationConfig = toGenerationConfig(options, promptTokenCount: tokens.count) + + await model.resetState() + + let promptTokenCount = tokens.count + // Text already published to the transcript and to the stream for this turn. + var publishedText = "" + + let outputTokens = await model.generate( + config: generationConfig, + tokens: tokens, + model: model.callAsFunction + ) { tokenIds in + if Task.isCancelled { return } + + let assistantText = decodeAssistantText( + from: tokenIds, + promptTokenCount: promptTokenCount + ) + // Hold back text that looks like the start of a tool call so that + // raw tool-call markup never reaches the transcript or the caller. + let visibleText = + toolsEnabled + ? CoreMLToolCallParser.visibleTextForStreaming(assistantText) + : assistantText + guard visibleText != publishedText else { return } + publishedText = visibleText + + // Grow the observable transcript so a Transcript-driven UI updates live. + session.growStreamingTranscript(text: visibleText) + continuation.yield( + .init( + content: (visibleText as! Content).asPartiallyGenerated(), + rawContent: GeneratedContent(visibleText) + ) ) - let delta = String(assistantText[startIdx...]) - accumulatedText += delta - } else { - accumulatedText = assistantText } - continuation.yield( - .init( - content: (accumulatedText as! Content).asPartiallyGenerated(), - rawContent: GeneratedContent(accumulatedText) + if Task.isCancelled { break } + + let assistantText = decodeAssistantText( + from: outputTokens, + promptTokenCount: promptTokenCount + ) + + // Tool calling requires the chat-template path: there is no other way to + // render tool specs into the prompt or feed tool results back to the model. + guard case .chat(var messages) = promptSource, toolsEnabled else { + publish( + assistantText, + ifDifferentFrom: publishedText, + in: session, + to: continuation ) + break + } + + let parsed = CoreMLToolCallParser.parse(assistantText, knownToolNames: toolNames) + + // Reconcile: the streaming heuristic is deliberately conservative, so the + // final parse is the authoritative split between prose and tool calls. + publish( + parsed.visibleText, + ifDifferentFrom: publishedText, + in: session, + to: continuation ) + + if !parsed.calls.isEmpty { + toolIteration += 1 + if toolIteration > Self.maximumToolIterations { + session.appendTranscriptEntry( + .toolCalls( + Transcript.ToolCalls(makeTranscriptToolCalls(from: parsed.calls)) + ) + ) + throw Self.maxToolIterationsExceededError( + limit: Self.maximumToolIterations + ) + } + + // Guard against a model that keeps asking for the exact same tool call. + let signature = CoreMLToolCallParser.signature(for: parsed.calls) + if signature == previousToolCallSignature { + session.appendTranscriptEntry( + .toolCalls( + Transcript.ToolCalls(makeTranscriptToolCalls(from: parsed.calls)) + ) + ) + throw Self.repeatedToolCallLoopError() + } + previousToolCallSignature = signature + + let transcriptCalls = makeTranscriptToolCalls(from: parsed.calls) + let resolution = try await resolveToolCalls(transcriptCalls, session: session) + switch resolution { + case .stop(let calls): + if !calls.isEmpty { + session.appendTranscriptEntry(.toolCalls(Transcript.ToolCalls(calls))) + } + continuation.finish() + return + case .invocations(let invocations): + if !invocations.isEmpty { + // Tool calls must land in the transcript before their outputs. + session.appendTranscriptEntry( + .toolCalls(Transcript.ToolCalls(invocations.map(\.call))) + ) + + messages.append( + assistantToolCallMessage( + text: parsed.visibleText, + transcriptCalls: transcriptCalls, + parsedCalls: parsed.calls + ) + ) + for invocation in invocations { + session.appendTranscriptEntry(.toolOutput(invocation.output)) + messages.append(toolResultMessage(for: invocation.output)) + } + + promptSource = .chat(messages) + continue + } + } + } + + break } continuation.finish() @@ -239,6 +384,26 @@ return LanguageModelSession.ResponseStream(stream: stream) } + /// Publishes `text` to the session transcript and the response stream when it differs + /// from what was already published for the current turn. + private func publish( + _ text: String, + ifDifferentFrom publishedText: String, + in session: LanguageModelSession, + to continuation: AsyncThrowingStream< + LanguageModelSession.ResponseStream.Snapshot, any Error + >.Continuation + ) where Content: Generable { + guard text != publishedText else { return } + session.growStreamingTranscript(text: text) + continuation.yield( + .init( + content: (text as! Content).asPartiallyGenerated(), + rawContent: GeneratedContent(text) + ) + ) + } + // MARK: - Image Validation private func validateNoImageSegments(in session: LanguageModelSession) throws { @@ -295,6 +460,25 @@ @available(macOS 15.0, iOS 18.0, tvOS 18.0, visionOS 2.0, watchOS 11.0, *) extension CoreMLLanguageModel { + /// Builds a generation config for a free-form (non-constrained) generation. + /// + /// - Important: `GenerationConfig.maxLength` defaults to 20 and `generate` stops as soon as + /// the total sequence reaches it, so it has to be sized against the prompt or nothing is + /// generated at all. `eosTokenId` likewise has no default, and without it generation never + /// stops early — it always runs to the token budget and decodes past the end of the turn, + /// which makes tool-call text impossible to parse reliably. The constrained-decoding path + /// in this file already sets both; this does the same for the plain path. + private func toGenerationConfig( + _ options: GenerationOptions, + promptTokenCount: Int + ) -> GenerationConfig { + var config = toGenerationConfig(options) + config.maxLength = config.maxNewTokens + promptTokenCount + config.eosTokenId = tokenizer.eosTokenId + config.bosTokenId = tokenizer.bosTokenId + return config + } + private func toGenerationConfig(_ options: GenerationOptions) -> GenerationConfig { var config = GenerationConfig(maxNewTokens: options.maximumResponseTokens ?? 2048) @@ -608,4 +792,635 @@ } } + + // MARK: - Tool Calling + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, visionOS 2.0, watchOS 11.0, *) + extension CoreMLLanguageModel { + /// The maximum number of tool round-trips allowed for a single response. + fileprivate static var maximumToolIterations: Int { 8 } + + fileprivate static func maxToolIterationsExceededError( + limit: Int + ) -> LanguageModelSession.GenerationError { + .decodingFailure( + .init( + debugDescription: + "Exceeded maximum tool iterations (\(limit)) while processing Core ML tool calls." + ) + ) + } + + fileprivate static func repeatedToolCallLoopError() -> LanguageModelSession.GenerationError { + .decodingFailure( + .init( + debugDescription: + "Detected repeated Core ML tool-call signature and aborted to avoid an infinite tool loop." + ) + ) + } + + // MARK: Prompt construction + + /// How the prompt for a generation turn is expressed. + fileprivate enum PromptSource { + /// Chat messages rendered through the tokenizer's chat template. Tool calling is only + /// available on this path. + case chat([Message]) + /// Raw text encoded directly by the tokenizer, with no chat template involved. + case rawText(String) + } + + /// Resolves the tool specifications handed to the chat template. + /// + /// A caller-supplied `toolsHandler` always wins, and is still called exactly as before — + /// including when the session has no tools. When no handler is supplied, the session's tools + /// are converted with the conventional OpenAI-style function schema that Hugging Face chat + /// templates expect, instead of silently dropping them. + fileprivate func resolvedToolSpecs(for session: LanguageModelSession) -> [ToolSpec]? { + if let toolsHandler { + return toolsHandler(session.tools) + } + guard !session.tools.isEmpty else { return nil } + return session.tools.map { convertToolToToolSpec($0) } + } + + fileprivate func makePromptSource( + session: LanguageModelSession, + prompt: Prompt + ) -> PromptSource { + if let chatTemplateHandler { + return .chat(chatTemplateHandler(session.instructions, prompt)) + } + + // Without a chat template handler the model is normally prompted with raw text, but tool + // specs and tool results can only be expressed through the chat template. When the + // session actually has tools, build the minimal message list the handler would have + // produced rather than dropping the tools. + if !session.tools.isEmpty, tokenizer.hasChatTemplate { + var messages: [Message] = [] + if let instructions = session.instructions { + let content = instructions.description + if !content.isEmpty { + messages.append(["role": "system", "content": content]) + } + } + messages.append(["role": "user", "content": prompt.description]) + return .chat(messages) + } + + return .rawText(prompt.description) + } + + fileprivate func encodePrompt( + _ source: PromptSource, + toolSpecs: [ToolSpec]? + ) throws -> [Int] { + switch source { + case .chat(let messages): + return try tokenizer.applyChatTemplate(messages: messages, tools: toolSpecs) + case .rawText(let text): + return tokenizer.encode(text: text) + } + } + + /// Strips the prompt at the token level to avoid issues with normalization or whitespace + /// differences in decoded strings. + fileprivate func decodeAssistantText(from outputTokens: [Int], promptTokenCount: Int) -> String { + let assistantTokenSlice: ArraySlice + if outputTokens.count >= promptTokenCount { + assistantTokenSlice = outputTokens.dropFirst(promptTokenCount) + } else { + // Fallback: if the model did not echo the full prompt, + // treat the entire output as assistant tokens + assistantTokenSlice = outputTokens[outputTokens.indices] + } + return tokenizer.decode(tokens: Array(assistantTokenSlice)) + } + + // MARK: Tool specs + + /// Converts a tool to the OpenAI-style function schema used by Hugging Face chat templates. + private func convertToolToToolSpec(_ tool: any Tool) -> ToolSpec { + let parametersDict: [String: any Sendable] + do { + let resolvedSchema = tool.parameters.withResolvedRoot() ?? tool.parameters + let data = try JSONEncoder().encode(resolvedSchema) + if let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] { + parametersDict = CoreMLToolCallParser.sendableJSONObject(from: json) + } else { + parametersDict = CoreMLToolCallParser.emptyJSONSchemaObject + } + } catch { + parametersDict = CoreMLToolCallParser.emptyJSONSchemaObject + } + + let functionSpec: [String: any Sendable] = [ + "name": tool.name, + "description": tool.description, + "parameters": parametersDict, + ] + + return [ + "type": "function", + "function": functionSpec, + ] + } + + // MARK: Feeding results back to the model + + /// The assistant turn that requested the tools. + /// + /// The tool calls are expressed as structured `tool_calls`, which is what the Hermes/Qwen, + /// Llama 3.x and Mistral chat templates read. `content` holds only the model's prose so that + /// templates rendering both `content` and `tool_calls` do not emit the call twice. + fileprivate func assistantToolCallMessage( + text: String, + transcriptCalls: [Transcript.ToolCall], + parsedCalls: [CoreMLToolCallParser.ParsedToolCall] + ) -> Message { + var toolCalls: [[String: any Sendable]] = [] + toolCalls.reserveCapacity(parsedCalls.count) + for (transcriptCall, parsedCall) in zip(transcriptCalls, parsedCalls) { + let function: [String: any Sendable] = [ + "name": parsedCall.name, + "arguments": parsedCall.arguments, + ] + toolCalls.append([ + "id": transcriptCall.id, + "type": "function", + "function": function, + ]) + } + + return [ + "role": "assistant", + "content": text, + "tool_calls": toolCalls, + ] + } + + /// The tool result turn. + /// + /// `role: "tool"` is what the Hermes/Qwen and Mistral templates expect; the Llama 3.x + /// templates accept either `"tool"` or `"ipython"`. `tool_call_id` and `name` are extra keys + /// that templates which do not use them simply ignore. + fileprivate func toolResultMessage(for output: Transcript.ToolOutput) -> Message { + [ + "role": "tool", + "tool_call_id": output.id, + "name": output.toolName, + "content": toolOutputText(output), + ] + } + + private func toolOutputText(_ output: Transcript.ToolOutput) -> String { + var textParts: [String] = [] + for segment in output.segments { + switch segment { + case .text(let textSegment): + textParts.append(textSegment.content) + case .structure(let structuredSegment): + textParts.append(structuredSegment.content.jsonString) + case .image: + // Image segments are not supported in Core ML tool output. + break + } + } + return textParts.joined(separator: "\n") + } + + // MARK: Tool invocation + + fileprivate struct ToolInvocationResult { + let call: Transcript.ToolCall + let output: Transcript.ToolOutput + } + + fileprivate enum ToolResolutionOutcome { + case stop(calls: [Transcript.ToolCall]) + case invocations([ToolInvocationResult]) + } + + fileprivate func makeTranscriptToolCalls( + from parsedCalls: [CoreMLToolCallParser.ParsedToolCall] + ) -> [Transcript.ToolCall] { + parsedCalls.map { parsedCall in + let arguments = + (try? GeneratedContent(json: parsedCall.argumentsJSON)) + ?? GeneratedContent(kind: .structure(properties: [:], orderedKeys: [])) + return Transcript.ToolCall( + id: CoreMLToolCallParser.makeToolCallID(), + toolName: parsedCall.name, + arguments: arguments + ) + } + } + + // NOTE: Every provider keeps its own file-private `resolveToolCalls`. This is Core ML's copy + // of that shared shape, written here because the existing ones are private to their files. + fileprivate func resolveToolCalls( + _ transcriptCalls: [Transcript.ToolCall], + session: LanguageModelSession + ) async throws -> ToolResolutionOutcome { + guard !transcriptCalls.isEmpty else { return .invocations([]) } + + var toolsByName: [String: any Tool] = [:] + for tool in session.tools where toolsByName[tool.name] == nil { + toolsByName[tool.name] = tool + } + + if let delegate = session.toolExecutionDelegate { + await delegate.didGenerateToolCalls(transcriptCalls, in: session) + } + + var decisions: [ToolExecutionDecision] = [] + decisions.reserveCapacity(transcriptCalls.count) + + if let delegate = session.toolExecutionDelegate { + for call in transcriptCalls { + let decision = await delegate.toolCallDecision(for: call, in: session) + if case .stop = decision { + return .stop(calls: transcriptCalls) + } + decisions.append(decision) + } + } else { + decisions = Array(repeating: .execute, count: transcriptCalls.count) + } + + var results: [ToolInvocationResult] = [] + results.reserveCapacity(transcriptCalls.count) + + for (index, call) in transcriptCalls.enumerated() { + switch decisions[index] { + case .stop: + // This branch should be unreachable because `.stop` returns during decision + // collection. Keep it as a defensive guard in case that logic changes. + return .stop(calls: transcriptCalls) + case .provideOutput(let segments): + let output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: segments + ) + if let delegate = session.toolExecutionDelegate { + await delegate.didExecuteToolCall(call, output: output, in: session) + } + results.append(ToolInvocationResult(call: call, output: output)) + case .execute: + guard let tool = toolsByName[call.toolName] else { + let message = Transcript.Segment.text( + .init(content: "Tool not found: \(call.toolName)") + ) + let output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: [message] + ) + if let delegate = session.toolExecutionDelegate { + await delegate.didExecuteToolCall(call, output: output, in: session) + } + results.append(ToolInvocationResult(call: call, output: output)) + continue + } + + do { + let segments = try await tool.makeOutputSegments(from: call.arguments) + let output = Transcript.ToolOutput( + id: call.id, + toolName: tool.name, + segments: segments + ) + if let delegate = session.toolExecutionDelegate { + await delegate.didExecuteToolCall(call, output: output, in: session) + } + results.append(ToolInvocationResult(call: call, output: output)) + } catch { + if let delegate = session.toolExecutionDelegate { + await delegate.didFailToolCall(call, error: error, in: session) + } + throw LanguageModelSession.ToolCallError(tool: tool, underlyingError: error) + } + } + } + + return .invocations(results) + } + } + + // MARK: - Tool Call Parsing + + /// Parses tool calls out of the plain text a Core ML hosted model generates. + /// + /// Core ML models have no structured tool-call channel. `swift-transformers` renders tool specs + /// into the prompt through the tokenizer's Jinja chat template and stops there — it has no + /// tool-call parser at all — so the model's request comes back as ordinary generated text in + /// whatever format its chat template taught it. Parsing is therefore necessarily + /// format-specific, and this parser supports an explicit, closed set of formats: + /// + /// 1. Hermes / Qwen / NousResearch tags: `{"name": …, "arguments": {…}}`, + /// repeated for multiple calls. `…` is accepted as a synonym. + /// 2. Mistral: `[TOOL_CALLS] [{"name": …, "arguments": {…}}]`. + /// 3. Llama 3.x: `<|python_tag|>{"name": …, "parameters": {…}}`, with `;` separating calls. + /// 4. A fenced ```` ```json ```` block whose content has tool-call shape. + /// 5. A bare JSON object or array that has tool-call shape and nothing else in the turn. + /// + /// Formats 4 and 5 are ambiguous with a model simply answering in JSON, so they only count as a + /// tool call when the name matches a tool the session actually has. Formats 1 to 3 are keyed on + /// unambiguous markers, so an unknown name there is reported as a call to a missing tool. + /// + /// Not supported: DeepSeek's `<｜tool▁calls▁begin｜>` markers, Llama's `<|python_tag|>` when it + /// carries actual Python rather than JSON, and any XML-argument format such as Claude's + /// `` blocks. + /// - Note: Internal rather than file-private only so that the parsing rules can be unit tested + /// without a downloaded Core ML model. + enum CoreMLToolCallParser { + struct ParsedToolCall { + let name: String + /// Canonical (sorted-key) JSON text of the arguments object. + let argumentsJSON: String + /// The same arguments, ready to be handed back to the chat template. + let arguments: [String: any Sendable] + } + + struct ParseResult { + /// The generated text with tool-call markup removed. + let visibleText: String + let calls: [ParsedToolCall] + } + + static let emptyJSONSchemaObject: [String: any Sendable] = [ + "type": "object", + "properties": [String: any Sendable](), + "required": [String](), + ] + + /// Paired tags that unambiguously wrap a tool call. + private static let tagPairs = [ + ("", ""), + ("", ""), + ] + + /// Markers that introduce a tool call and run to the end of the turn. + private static let leadingMarkers = ["[TOOL_CALLS]", "<|python_tag|>"] + + /// Everything that signals a tool call may be starting, used to hold text back mid-stream. + private static var streamingMarkers: [String] { + tagPairs.map(\.0) + leadingMarkers + } + + // MARK: Streaming + + /// The portion of a partially generated turn that is safe to show the caller. + /// + /// This runs on every token, so it stays cheap and conservative: text is cut at the first + /// tool-call marker, and a turn that starts with `{` or `[` is withheld entirely because it + /// may still turn out to be a bare-JSON tool call. A fenced ```` ```json ```` tool call is + /// not withheld and is therefore briefly visible before ``parse(_:knownToolNames:)`` removes + /// it at the end of the turn. + static func visibleTextForStreaming(_ text: String) -> String { + var cutoff = text.endIndex + for marker in streamingMarkers { + if let range = text.range(of: marker), range.lowerBound < cutoff { + cutoff = range.lowerBound + } + } + + let visibleText = String(text[text.startIndex ..< cutoff]) + let trimmed = visibleText.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.hasPrefix("{") || trimmed.hasPrefix("[") { + return "" + } + return visibleText + } + + // MARK: Parsing + + static func parse(_ text: String, knownToolNames: Set) -> ParseResult { + if let result = parseTaggedBlocks(text) { return result } + if let result = parseLeadingMarker(text) { return result } + if let result = parseFencedJSON(text, knownToolNames: knownToolNames) { return result } + if let result = parseBareJSON(text, knownToolNames: knownToolNames) { return result } + return ParseResult(visibleText: text, calls: []) + } + + /// A stable identity for a set of calls, used to detect a model looping on the same request. + static func signature(for calls: [ParsedToolCall]) -> String { + calls.map { "\($0.name):\($0.argumentsJSON)" }.joined(separator: "|") + } + + /// Mistral's chat template rejects tool call ids that are not exactly nine alphanumeric + /// characters, and no other template cares about the shape, so use that everywhere. + static func makeToolCallID() -> String { + let alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + return String((0 ..< 9).compactMap { _ in alphabet.randomElement() }) + } + + // MARK: Format-specific parsing + + private static func parseTaggedBlocks(_ text: String) -> ParseResult? { + var calls: [ParsedToolCall] = [] + var visibleText = "" + var index = text.startIndex + + while index < text.endIndex { + var earliest: (open: Range, close: String)? + for (open, close) in tagPairs { + guard let range = text.range(of: open, range: index ..< text.endIndex) else { continue } + if earliest == nil || range.lowerBound < earliest!.open.lowerBound { + earliest = (range, close) + } + } + guard let (openRange, closeTag) = earliest else { break } + + visibleText += text[index ..< openRange.lowerBound] + + let bodyStart = openRange.upperBound + let bodyEnd: String.Index + if let closeRange = text.range(of: closeTag, range: bodyStart ..< text.endIndex) { + bodyEnd = closeRange.lowerBound + index = closeRange.upperBound + } else { + // Truncated generation: take the rest of the turn as the call body. + bodyEnd = text.endIndex + index = text.endIndex + } + + calls.append(contentsOf: toolCalls(fromJSONText: String(text[bodyStart ..< bodyEnd]))) + } + + // A tag with an unparseable body is treated as ordinary text rather than a lost call. + guard !calls.isEmpty else { return nil } + + visibleText += text[index ..< text.endIndex] + return ParseResult( + visibleText: visibleText.trimmingCharacters(in: .whitespacesAndNewlines), + calls: calls + ) + } + + private static func parseLeadingMarker(_ text: String) -> ParseResult? { + var earliest: Range? + for marker in leadingMarkers { + guard let range = text.range(of: marker) else { continue } + if earliest == nil || range.lowerBound < earliest!.lowerBound { + earliest = range + } + } + guard let markerRange = earliest else { return nil } + + let body = String(text[markerRange.upperBound...]) + var calls = toolCalls(fromJSONText: body) + if calls.isEmpty { + // Llama 3.1 separates multiple calls with `;`. + calls = body.split(separator: ";").flatMap { toolCalls(fromJSONText: String($0)) } + } + guard !calls.isEmpty else { return nil } + + let visibleText = String(text[text.startIndex ..< markerRange.lowerBound]) + return ParseResult( + visibleText: visibleText.trimmingCharacters(in: .whitespacesAndNewlines), + calls: calls + ) + } + + private static func parseFencedJSON( + _ text: String, + knownToolNames: Set + ) -> ParseResult? { + let fence = "```" + var calls: [ParsedToolCall] = [] + var visibleText = "" + var index = text.startIndex + + while index < text.endIndex { + guard let openRange = text.range(of: fence, range: index ..< text.endIndex), + let closeRange = text.range(of: fence, range: openRange.upperBound ..< text.endIndex) + else { break } + + var body = String(text[openRange.upperBound ..< closeRange.lowerBound]) + // Drop an optional language tag on the opening fence. + if let newline = body.firstIndex(of: "\n") { + let firstLine = body[body.startIndex ..< newline].trimmingCharacters( + in: .whitespaces + ) + if firstLine.isEmpty || firstLine.allSatisfy({ $0.isLetter }) { + body = String(body[body.index(after: newline)...]) + } + } + + let blockCalls = toolCalls(fromJSONText: body) + if blockCalls.isEmpty || !blockCalls.allSatisfy({ knownToolNames.contains($0.name) }) { + // Not a tool call: keep the fenced block as visible text. + visibleText += text[index ..< closeRange.upperBound] + } else { + visibleText += text[index ..< openRange.lowerBound] + calls.append(contentsOf: blockCalls) + } + index = closeRange.upperBound + } + + guard !calls.isEmpty else { return nil } + + visibleText += text[index ..< text.endIndex] + return ParseResult( + visibleText: visibleText.trimmingCharacters(in: .whitespacesAndNewlines), + calls: calls + ) + } + + private static func parseBareJSON( + _ text: String, + knownToolNames: Set + ) -> ParseResult? { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard trimmed.hasPrefix("{") || trimmed.hasPrefix("[") else { return nil } + + let calls = toolCalls(fromJSONText: trimmed) + guard !calls.isEmpty, calls.allSatisfy({ knownToolNames.contains($0.name) }) else { + return nil + } + return ParseResult(visibleText: "", calls: calls) + } + + // MARK: JSON shaping + + private static func toolCalls(fromJSONText text: String) -> [ParsedToolCall] { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty, let data = trimmed.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: data) + else { return [] } + + if let array = json as? [Any] { + return array.compactMap { element in + (element as? [String: Any]).flatMap { toolCall(from: $0) } + } + } + if let object = json as? [String: Any], let call = toolCall(from: object) { + return [call] + } + return [] + } + + private static func toolCall(from object: [String: Any]) -> ParsedToolCall? { + // Unwrap the OpenAI-style `{"type": "function", "function": {…}}` envelope. + let source = (object["function"] as? [String: Any]) ?? object + + guard let name = source["name"] as? String, !name.isEmpty else { return nil } + + var argumentsValue = source["arguments"] ?? source["parameters"] ?? source["args"] + // Some templates emit the arguments as a JSON-encoded string. + if let argumentsString = argumentsValue as? String { + argumentsValue = + argumentsString.data(using: .utf8) + .flatMap { try? JSONSerialization.jsonObject(with: $0) } + } + let argumentsObject = (argumentsValue as? [String: Any]) ?? [:] + + guard + let data = try? JSONSerialization.data( + withJSONObject: argumentsObject, + options: [.sortedKeys] + ), + let argumentsJSON = String(data: data, encoding: .utf8) + else { return nil } + + return ParsedToolCall( + name: name, + argumentsJSON: argumentsJSON, + arguments: sendableJSONObject(from: argumentsObject) + ) + } + + /// Re-types a `JSONSerialization` object graph as `Sendable` so it can be handed to the + /// tokenizer's chat template. + static func sendableJSONObject(from object: [String: Any]) -> [String: any Sendable] { + var converted: [String: any Sendable] = [:] + converted.reserveCapacity(object.count) + for (key, value) in object { + converted[key] = sendableJSONValue(from: value) + } + return converted + } + + private static func sendableJSONValue(from value: Any) -> any Sendable { + if let string = value as? String { return string } + if let array = value as? [Any] { return array.map { sendableJSONValue(from: $0) } } + if let dictionary = value as? [String: Any] { return sendableJSONObject(from: dictionary) } + if value is NSNull { + // Erased `nil`, which the template engine degrades to null. + let null: String? = nil + return null as any Sendable + } + if let number = value as? NSNumber { + // `as? Bool` is not reliable here: JSON booleans and the integers 0 and 1 both + // bridge to NSNumber, so distinguish them by the underlying CoreFoundation type. + if CFGetTypeID(number) == CFBooleanGetTypeID() { return number.boolValue } + if let integer = Int(exactly: number) { return integer } + return number.doubleValue + } + return String(describing: value) + } + } #endif // CoreML From 59b415d89ed56886c424d7845e1d25fd372f08fb Mon Sep 17 00:00:00 2001 From: Claude Opus 5 Date: Mon, 3 Aug 2026 11:57:59 -0400 Subject: [PATCH 2/2] Test Core ML tool calling Adds `withTools` and `streamWithTools` to the existing Core ML suite, modeled on the Anthropic equivalents. `streamWithTools` additionally asserts that tool-call markup never appears in the streamed assistant text. Both follow the suite's existing gating and so need the downloaded Core ML model. Because that model is not available in most environments, this also adds an ungated `CoreMLToolCallParsing` suite covering the tool-call text formats the provider claims to support: Hermes/Qwen tags, Mistral `[TOOL_CALLS]`, Llama `<|python_tag|>`, arguments encoded as a JSON string, the OpenAI function envelope, bare JSON gated on known tool names, the mid-stream hold-back heuristic, loop-detection signatures, and the nine-character alphanumeric tool call ids that Mistral's chat template requires. These need no model and run wherever the CoreML trait is enabled. Co-Authored-By: Claude Opus 5 (1M context) --- .../CoreMLLanguageModelTests.swift | 205 ++++++++++++++++++ 1 file changed, 205 insertions(+) diff --git a/Tests/AnyLanguageModelTests/CoreMLLanguageModelTests.swift b/Tests/AnyLanguageModelTests/CoreMLLanguageModelTests.swift index 9d697d56..accd2d1d 100644 --- a/Tests/AnyLanguageModelTests/CoreMLLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/CoreMLLanguageModelTests.swift @@ -300,6 +300,68 @@ import Testing #expect(!response.content.colors.isEmpty) } + /// Requires a model whose chat template supports tools. The template renders the tool specs, + /// the model answers with tool-call text, and the parser turns that back into transcript + /// entries. + @Test @available(macOS 15.0, iOS 18.0, tvOS 18.0, visionOS 2.0, watchOS 11.0, *) + func withTools() async throws { + let model = try await getModel() + let weatherTool = WeatherTool() + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + let response = try await session.respond(to: "How's the weather in San Francisco?") + + var foundToolOutput = false + for case let .toolOutput(toolOutput) in response.transcriptEntries { + #expect(!toolOutput.id.isEmpty) + #expect(toolOutput.toolName == "getWeather") + foundToolOutput = true + } + #expect(foundToolOutput) + } + + @Test @available(macOS 15.0, iOS 18.0, tvOS 18.0, visionOS 2.0, watchOS 11.0, *) + func streamWithTools() async throws { + let model = try await getModel() + let weatherTool = WeatherTool() + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + let stream = session.streamResponse(to: "How's the weather in San Francisco?") + + var snapshots: [LanguageModelSession.ResponseStream.Snapshot] = [] + + var toolAppearedInTranscript: Bool = false + var toolResponseAppearedInTranscript: Bool = false + + for try await snapshot in stream { + snapshots.append(snapshot) + + for entry in session.transcript { + switch entry { + case .toolCalls: + toolAppearedInTranscript = true + case .toolOutput: + toolResponseAppearedInTranscript = true + default: break + } + } + } + + #expect(toolAppearedInTranscript, "Expected a tool call to appear in the transcript during streaming.") + #expect( + toolResponseAppearedInTranscript, + "Expected a tool output to appear in the transcript during streaming." + ) + + // Tool-call markup must never be published as assistant text. + if #available(macOS 26.0, iOS 26.0, tvOS 26.0, visionOS 26.0, watchOS 26.0, *) { + for snapshot in snapshots { + #expect(!snapshot.content.contains("")) + #expect(!snapshot.content.contains("[TOOL_CALLS]")) + } + } + } + @Test @available(macOS 15.0, iOS 18.0, tvOS 18.0, visionOS 2.0, watchOS 11.0, *) func structuredGenerationNestedStruct() async throws { let model = try await getModel() @@ -317,4 +379,147 @@ import Testing #expect(!response.content.address.city.isEmpty) } } + + /// Exercises the tool-call text formats the Core ML provider claims to support. These need no + /// downloaded model, so unlike the suite above they run everywhere the CoreML trait is enabled. + @Suite("CoreMLToolCallParsing") + struct CoreMLToolCallParsingTests { + private let knownToolNames: Set = ["getWeather"] + + @Test func parsesHermesStyleTaggedCall() { + let result = CoreMLToolCallParser.parse( + "Let me look that up.\n\n{\"name\": \"getWeather\", \"arguments\": {\"city\": \"Paris\"}}\n", + knownToolNames: knownToolNames + ) + + #expect(result.visibleText == "Let me look that up.") + #expect(result.calls.count == 1) + #expect(result.calls.first?.name == "getWeather") + #expect(result.calls.first?.argumentsJSON == "{\"city\":\"Paris\"}") + } + + @Test func parsesMultipleTaggedCalls() { + let result = CoreMLToolCallParser.parse( + "{\"name\": \"getWeather\", \"arguments\": {\"city\": \"Paris\"}}" + + "{\"name\": \"getWeather\", \"arguments\": {\"city\": \"Oslo\"}}", + knownToolNames: knownToolNames + ) + + #expect(result.calls.count == 2) + #expect(result.visibleText.isEmpty) + } + + @Test func parsesMistralToolCallsMarker() { + let result = CoreMLToolCallParser.parse( + "[TOOL_CALLS] [{\"name\": \"getWeather\", \"arguments\": {\"city\": \"Paris\"}}]", + knownToolNames: knownToolNames + ) + + #expect(result.calls.count == 1) + #expect(result.calls.first?.name == "getWeather") + #expect(result.visibleText.isEmpty) + } + + @Test func parsesLlamaPythonTagWithParametersKey() { + let result = CoreMLToolCallParser.parse( + "<|python_tag|>{\"name\": \"getWeather\", \"parameters\": {\"city\": \"Paris\"}}", + knownToolNames: knownToolNames + ) + + #expect(result.calls.count == 1) + #expect(result.calls.first?.argumentsJSON == "{\"city\":\"Paris\"}") + } + + @Test func parsesArgumentsEncodedAsJSONString() { + let result = CoreMLToolCallParser.parse( + "{\"name\": \"getWeather\", \"arguments\": \"{\\\"city\\\": \\\"Paris\\\"}\"}", + knownToolNames: knownToolNames + ) + + #expect(result.calls.first?.argumentsJSON == "{\"city\":\"Paris\"}") + } + + @Test func parsesOpenAIFunctionEnvelope() { + let result = CoreMLToolCallParser.parse( + "{\"type\": \"function\", \"function\": {\"name\": \"getWeather\", \"arguments\": {\"city\": \"Paris\"}}}", + knownToolNames: knownToolNames + ) + + #expect(result.calls.first?.name == "getWeather") + } + + @Test func parsesBareJSONOnlyForKnownTools() { + let text = "{\"name\": \"getWeather\", \"arguments\": {\"city\": \"Paris\"}}" + + let known = CoreMLToolCallParser.parse(text, knownToolNames: knownToolNames) + #expect(known.calls.count == 1) + + let unknown = CoreMLToolCallParser.parse(text, knownToolNames: []) + #expect(unknown.calls.isEmpty) + #expect(unknown.visibleText == text) + } + + @Test func treatsPlainProseAsText() { + let result = CoreMLToolCallParser.parse( + "The weather in Paris is sunny.", + knownToolNames: knownToolNames + ) + + #expect(result.calls.isEmpty) + #expect(result.visibleText == "The weather in Paris is sunny.") + } + + @Test func treatsUnparseableTagBodyAsText() { + let text = "not json" + let result = CoreMLToolCallParser.parse(text, knownToolNames: knownToolNames) + + #expect(result.calls.isEmpty) + #expect(result.visibleText == text) + } + + @Test func withholdsToolCallMarkupWhileStreaming() { + #expect( + CoreMLToolCallParser.visibleTextForStreaming("Checking. {\"na") + == "Checking. " + ) + #expect(CoreMLToolCallParser.visibleTextForStreaming("[TOOL_CALLS] [{\"na").isEmpty) + #expect(CoreMLToolCallParser.visibleTextForStreaming("{\"name\": \"get").isEmpty) + #expect( + CoreMLToolCallParser.visibleTextForStreaming("The weather is") == "The weather is" + ) + } + + @Test func signatureDistinguishesArgumentsAndIgnoresKeyOrder() { + let first = CoreMLToolCallParser.parse( + "{\"name\": \"getWeather\", \"arguments\": {\"city\": \"Paris\", \"unit\": \"C\"}}", + knownToolNames: knownToolNames + ) + let reordered = CoreMLToolCallParser.parse( + "{\"name\": \"getWeather\", \"arguments\": {\"unit\": \"C\", \"city\": \"Paris\"}}", + knownToolNames: knownToolNames + ) + let different = CoreMLToolCallParser.parse( + "{\"name\": \"getWeather\", \"arguments\": {\"city\": \"Oslo\"}}", + knownToolNames: knownToolNames + ) + + #expect( + CoreMLToolCallParser.signature(for: first.calls) + == CoreMLToolCallParser.signature(for: reordered.calls) + ) + #expect( + CoreMLToolCallParser.signature(for: first.calls) + != CoreMLToolCallParser.signature(for: different.calls) + ) + } + + @Test func toolCallIDsAreNineAlphanumericCharacters() { + // Mistral's chat template rejects anything else. + for _ in 0 ..< 32 { + let id = CoreMLToolCallParser.makeToolCallID() + #expect(id.count == 9) + #expect(id.allSatisfy { $0.isLetter || $0.isNumber }) + } + } + } #endif // CoreML