From befc05d2834fd52f7e061964cb0ab773e5324071 Mon Sep 17 00:00:00 2001 From: Claude Opus 5 Date: Mon, 3 Aug 2026 11:51:26 -0400 Subject: [PATCH] Add streaming tool calling to the MLX provider MLXLanguageModel.streamResponse passed `tools: nil` to makeUserInput and its generation loop ignored `.toolCall` items, so tool calls were never requested from the model and never handled. Streaming with tools silently behaved as if the session had no tools, while the non-streaming `respond` path supported them fully. streamResponse now passes the session's tool specs via the existing mlxToolSpecs(for:) helper and wraps generation in a turn loop that mirrors the semantics of AnthropicLanguageModel.streamResponse. Each turn accumulates `.chunk` text and collects `.toolCall` items; when a turn ends with tool calls pending they are executed and the assistant turn plus tool results are appended to the chat before generating again. The loop exits when a turn completes without tool calls. Tool execution reuses the same makeTranscriptToolCalls and resolveToolCalls functions the non-streaming path uses, including its max-iteration ceiling and repeated-tool-call-signature loop detection, so both paths abort identically on a runaway model. A `.stop` resolution appends the tool calls and finishes the stream without executing them. The transcript is updated live: growStreamingTranscript as text arrives, and appendTranscriptEntry for tool activity with `.toolCalls` always appended before the `.toolOutput` entries of the same turn. Accumulated text resets per turn so the trailing transcript response entry tracks only the turn currently being generated, matching the Anthropic provider. Resource management is unchanged in shape: the GPU memory scope and generation slot are still released through the idempotent finishScope/finishGenerationSlot pair guarded by didEndScope/didReleaseGenerationSlot, now covering the loop's normal exit, the `.stop` early return, the thrown-error path, and cancellation. The KV cache is resolved and stored per turn, as in respond. --- .../Models/MLXLanguageModel.swift | 152 +++++++++++++----- .../MLXLanguageModelTests.swift | 50 ++++++ 2 files changed, 163 insertions(+), 39 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift index cc2ece85..6a9fbd37 100644 --- a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift @@ -1103,54 +1103,128 @@ import Foundation let userInputProcessing = options[custom: MLXLanguageModel.self]?.processingForUserInput ?? .init(resize: nil) - let chat = convertTranscriptToMLXChat( + let toolSpecs = mlxToolSpecs(for: session) + var chat = convertTranscriptToMLXChat( session: session, fallbackPrompt: prompt.description ) - let userInput = makeUserInput( - chat: chat, - tools: nil, - processing: userInputProcessing, - additionalContext: additionalContext - ) - let lmInput = try await context.processor.prepare(input: userInput) - let resolved = resolveCache( - session: session, - lmInput: lmInput, - generateParameters: generateParameters, - context: context - ) + let maxToolIterations = 8 + var toolIteration = 0 + var previousToolCallSignature: String? + + // Loop until a turn completes without requesting tool calls. + generationLoop: while true { + let userInput = makeUserInput( + chat: chat, + tools: toolSpecs, + processing: userInputProcessing, + additionalContext: additionalContext + ) + let lmInput = try await context.processor.prepare(input: userInput) + let resolved = resolveCache( + session: session, + lmInput: lmInput, + generateParameters: generateParameters, + context: context + ) + + let mlxStream = try MLXLMCommon.generate( + input: resolved.input, + cache: resolved.cache, + parameters: generateParameters, + context: context + ) + + // Text restarts per turn so the trailing transcript response entry + // tracks only the text of the turn currently being generated. + var accumulatedText = "" + var collectedToolCalls: [MLXLMCommon.ToolCall] = [] + + for await item in mlxStream { + if Task.isCancelled { break } + + switch item { + case .chunk(let text): + accumulatedText += text + + // Grow the observable transcript so a Transcript-driven UI updates live. + session.growStreamingTranscript(text: accumulatedText) + + let raw = GeneratedContent(accumulatedText) + let content: Content.PartiallyGenerated = (accumulatedText as! Content) + .asPartiallyGenerated() + continuation.yield(.init(content: content, rawContent: raw)) + case .info: + break + case .toolCall(let call): + collectedToolCalls.append(call) + } + } - let mlxStream = try MLXLMCommon.generate( - input: resolved.input, - cache: resolved.cache, - parameters: generateParameters, - context: context - ) + storeSessionCache( + cache: resolved.cache, + fullTokens: resolved.fullTokens, + generateParameters: generateParameters, + session: session + ) + + if Task.isCancelled { break generationLoop } - var accumulatedText = "" - for await item in mlxStream { - if Task.isCancelled { break } - - switch item { - case .chunk(let text): - accumulatedText += text - let raw = GeneratedContent(accumulatedText) - let content: Content.PartiallyGenerated = (accumulatedText as! Content) - .asPartiallyGenerated() - continuation.yield(.init(content: content, rawContent: raw)) - case .info, .toolCall: - break + // Add the assistant turn to the chat history so the next + // generation sees what the model just said. + if !accumulatedText.isEmpty { + chat.append(.assistant(accumulatedText)) + } + + // No tool calls means the turn is complete. + if collectedToolCalls.isEmpty { break generationLoop } + + toolIteration += 1 + if toolIteration > maxToolIterations { + let unresolvedCalls = try makeTranscriptToolCalls(from: collectedToolCalls) + session.appendTranscriptEntry(.toolCalls(Transcript.ToolCalls(unresolvedCalls))) + throw Self.maxToolIterationsExceededError(limit: maxToolIterations) + } + + let signature = + collectedToolCalls + .map { "\($0.function.name):\($0.function.arguments)" } + .joined(separator: "|") + if signature == previousToolCallSignature { + let unresolvedCalls = try makeTranscriptToolCalls(from: collectedToolCalls) + session.appendTranscriptEntry(.toolCalls(Transcript.ToolCalls(unresolvedCalls))) + throw Self.repeatedToolCallLoopError() + } + previousToolCallSignature = signature + + let resolution = try await resolveMLXToolCalls(collectedToolCalls, session: session) + switch resolution { + case .stop(let calls): + if !calls.isEmpty { + session.appendTranscriptEntry(.toolCalls(Transcript.ToolCalls(calls))) + } + finishScope() + finishGenerationSlot() + continuation.finish() + return + case .invocations(let invocations): + if invocations.isEmpty { break generationLoop } + + // Tool calls must be appended before the outputs of the same turn. + session.appendTranscriptEntry( + .toolCalls(Transcript.ToolCalls(invocations.map(\.call))) + ) + + for invocation in invocations { + session.appendTranscriptEntry(.toolOutput(invocation.output)) + + // Convert tool output to a string payload for MLX + chat.append(.tool(toolOutputToJSON(invocation.output))) + } } } - storeSessionCache( - cache: resolved.cache, - fullTokens: resolved.fullTokens, - generateParameters: generateParameters, - session: session - ) finishScope() finishGenerationSlot() continuation.finish() diff --git a/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift index bb048d3e..2299830d 100644 --- a/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift @@ -144,6 +144,56 @@ import Testing } } + @Test func streamWithTools() async throws { + let weatherTool = spy(on: WeatherTool()) + let session = LanguageModelSession( + model: model, + tools: [weatherTool], + instructions: "You are a helpful assistant. Use available tools when needed." + ) + + let stream = session.streamResponse(to: "How's the weather in San Francisco?") + + var snapshots: [LanguageModelSession.ResponseStream.Snapshot] = [] + + var toolAppearedInTranscript = false + var toolResponseAppearedInTranscript = 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." + ) + + var foundToolOutput = false + for case let .toolOutput(toolOutput) in session.transcript { + #expect(!toolOutput.id.isEmpty) + #expect(toolOutput.toolName == weatherTool.name) + foundToolOutput = true + } + #expect(foundToolOutput, "Expected the 'getWeather' tool to exist in the final transcript.") + + let calls = await weatherTool.calls + #expect(calls.count >= 1) + if let first = calls.first { + #expect(first.arguments.city.contains("San Francisco")) + } + } + @Test func multimodalWithImageURL() async throws { let transcript = Transcript(entries: [ .prompt(