diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index ba38550e..e0b6445b 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -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 { @@ -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 + /// 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 = [] + ) { self.content = content self.rawContent = rawContent + self.transcriptEntries = transcriptEntries } } } @@ -887,7 +900,7 @@ extension LanguageModelSession.ResponseStream: AsyncSequence { return LanguageModelSession.Response( content: finalContent, rawContent: last.rawContent, - transcriptEntries: [] + transcriptEntries: last.transcriptEntries ) } } @@ -902,7 +915,7 @@ extension LanguageModelSession.ResponseStream: AsyncSequence { return LanguageModelSession.Response( content: finalContent, rawContent: fallbackSnapshot.rawContent, - transcriptEntries: [] + transcriptEntries: fallbackSnapshot.transcriptEntries ) } diff --git a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift index 0ef37efe..138ba0f0 100644 --- a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift @@ -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() diff --git a/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift index bb048d3e..2ac68bc5 100644 --- a/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift @@ -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.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(