diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index ba38550..c24406c 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -278,6 +278,18 @@ public final class LanguageModelSession: @unchecked Sendable { } } +// MARK: - Transcript Manipulation + +extension LanguageModelSession { + /// Allows manual injection of Transcript entries to resume dropped sessions + /// without busting the MLX ObjectIdentifier KV-cache. + public func appendTranscriptEntry(_ entry: Transcript.Entry) { + withMutation(keyPath: \.transcript) { + state.withLock { $0.transcript.append(entry) } + } + } +} + // MARK: - String Response Convenience Methods extension LanguageModelSession { diff --git a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift index 0ef37ef..33eb62d 100644 --- a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift @@ -1010,8 +1010,12 @@ import Foundation chat.append(.tool(toolResultJSON)) } - // Continue loop to generate with tool results - continue + // Return response with appended tool outputs rather than looping to regenerate + return LanguageModelSession.Response( + content: assistantText as! Content, + rawContent: GeneratedContent(assistantText), + transcriptEntries: ArraySlice(allEntries) + ) } } }