diff --git a/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift b/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift index 221119c..9d311ae 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 diff --git a/Tests/AnyLanguageModelTests/CoreMLLanguageModelTests.swift b/Tests/AnyLanguageModelTests/CoreMLLanguageModelTests.swift index 9d697d5..accd2d1 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