Skip to content

Support tool calling in MLXLanguageModel.streamResponse() #181

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

Open
noorbhatia wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Support tool calling in MLXLanguageModel.streamResponse() #181

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
21 changes: 17 additions & 4 deletions Sources/AnyLanguageModel/LanguageModelSession.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,10 @@ public final class LanguageModelSession: @unchecked Sendable {
)
)
session.withMutation(keyPath: \.transcript) {
session.state.withLock { $0.transcript.append(responseEntry) }
session.state.withLock {
$0.transcript.append(contentsOf: lastSnapshot.transcriptEntries)
$0.transcript.append(responseEntry)
}
}
}
} catch {
Expand Down Expand Up @@ -818,13 +821,23 @@ extension LanguageModelSession {
public var content: Content.PartiallyGenerated
public var rawContent: GeneratedContent

/// Transcript entries (tool calls and outputs) produced so far while streaming.
/// Cumulative across tool rounds; empty for providers that don't stream tool activity.
public var transcriptEntries: ArraySlice<Transcript.Entry>

/// Creates a snapshot from partially generated content and raw content.
/// - Parameters:
/// - content: The partially generated content.
/// - rawContent: The raw content produced by the model.
public init(content: Content.PartiallyGenerated, rawContent: GeneratedContent) {
/// - transcriptEntries: Transcript entries accumulated so far (tool calls/outputs).
public init(
content: Content.PartiallyGenerated,
rawContent: GeneratedContent,
transcriptEntries: ArraySlice<Transcript.Entry> = []
) {
self.content = content
self.rawContent = rawContent
self.transcriptEntries = transcriptEntries
}
}
}
Expand Down Expand Up @@ -887,7 +900,7 @@ extension LanguageModelSession.ResponseStream: AsyncSequence {
return LanguageModelSession.Response(
content: finalContent,
rawContent: last.rawContent,
transcriptEntries: []
transcriptEntries: last.transcriptEntries
)
}
}
Expand All @@ -902,7 +915,7 @@ extension LanguageModelSession.ResponseStream: AsyncSequence {
return LanguageModelSession.Response(
content: finalContent,
rawContent: fallbackSnapshot.rawContent,
transcriptEntries: []
transcriptEntries: fallbackSnapshot.transcriptEntries
)
}

Expand Down
149 changes: 110 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 @@ -1087,54 +1087,125 @@ 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
)
// Accumulators live outside the tool loop so streamed snapshots stay
// monotonic across rounds: text never shrinks, entries only grow.
var accumulatedText = ""
var accumulatedEntries: [Transcript.Entry] = []
let maxToolIterations = 8
var toolIteration = 0
var previousToolCallSignature: String?

// Yields a snapshot carrying the cumulative text and tool entries so far.
func yieldSnapshot() {
let raw = GeneratedContent(accumulatedText)
let content: Content.PartiallyGenerated = (accumulatedText as! Content)
.asPartiallyGenerated()
continuation.yield(
.init(
content: content,
rawContent: raw,
transcriptEntries: ArraySlice(accumulatedEntries)
)
)
}

let mlxStream = try MLXLMCommon.generate(
input: resolved.input,
cache: resolved.cache,
parameters: generateParameters,
context: context
)
// Loop until the model stops without pending tool calls (mirrors `respond()`).
toolLoop: 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
)

let roundStartTextCount = accumulatedText.count
var collectedToolCalls: [MLXLMCommon.ToolCall] = []

for await item in mlxStream {
if Task.isCancelled { break toolLoop }

switch item {
case .chunk(let text):
accumulatedText += text
yieldSnapshot()
case .toolCall(let call):
collectedToolCalls.append(call)
case .info:
break
}
}

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
storeSessionCache(
cache: resolved.cache,
fullTokens: resolved.fullTokens,
generateParameters: generateParameters,
session: session
)

// Feed this round's assistant text back into the chat history.
let roundText = String(accumulatedText.dropFirst(roundStartTextCount))
if !roundText.isEmpty {
chat.append(.assistant(roundText))
}

guard !collectedToolCalls.isEmpty else { break }

toolIteration += 1
if toolIteration > maxToolIterations {
throw Self.maxToolIterationsExceededError(limit: maxToolIterations)
}

let signature =
collectedToolCalls
.map { "\($0.function.name):\($0.function.arguments)" }
.joined(separator: "|")
if signature == previousToolCallSignature {
throw Self.repeatedToolCallLoopError()
}
previousToolCallSignature = signature

let resolution = try await resolveToolCalls(collectedToolCalls, session: session)
switch resolution {
case .stop(let calls):
if !calls.isEmpty {
accumulatedEntries.append(.toolCalls(Transcript.ToolCalls(calls)))
yieldSnapshot()
}
break toolLoop
case .invocations(let invocations):
if invocations.isEmpty { break toolLoop }

accumulatedEntries.append(
.toolCalls(Transcript.ToolCalls(invocations.map(\.call)))
)
for invocation in invocations {
accumulatedEntries.append(.toolOutput(invocation.output))
chat.append(.tool(toolOutputToJSON(invocation.output)))
}
yieldSnapshot()
}
}

storeSessionCache(
cache: resolved.cache,
fullTokens: resolved.fullTokens,
generateParameters: generateParameters,
session: session
)
finishScope()
finishGenerationSlot()
continuation.finish()
Expand Down
39 changes: 39 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,45 @@ import Testing
}
}

@Test func streamingWithTools() 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?")

// Iterate the stream, keeping the last snapshot as the final state.
var snapshotCount = 0
var lastSnapshot: LanguageModelSession.ResponseStream<String>.Snapshot?
for try await snapshot in stream {
snapshotCount += 1
lastSnapshot = snapshot
}

// The stream yielded incremental snapshots and produced text.
#expect(snapshotCount >= 1)
#expect(!(lastSnapshot?.content.isEmpty ?? true))

// The tool actually executed.
let calls = await weatherTool.calls
#expect(calls.count >= 1)
if let first = calls.first {
#expect(first.arguments.city.contains("San Francisco"))
}

// Tool activity surfaces through the stream's transcript entries.
var foundToolOutput = false
for case let .toolOutput(toolOutput) in lastSnapshot?.transcriptEntries ?? [] {
#expect(!toolOutput.id.isEmpty)
#expect(toolOutput.toolName == weatherTool.name)
foundToolOutput = true
}
#expect(foundToolOutput)
}

@Test func multimodalWithImageURL() async throws {
let transcript = Transcript(entries: [
.prompt(
Expand Down