Skip to content

Add streaming tool calling to the MLX provider #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ActuallyTaylor merged 1 commit into from
Aug 3, 2026
Merged

Add streaming tool calling to the MLX provider #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
152 changes: 113 additions & 39 deletions Sources/AnyLanguageModel/Models/MLXLanguageModel.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
50 changes: 50 additions & 0 deletions Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>.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(
Expand Down