From b67195ea13cffa1a3a6408e3d13350f7ca9ba213 Mon Sep 17 00:00:00 2001 From: Kevin Hermawan <84965338+kevinhermawan@users.noreply.github.com> Date: Tue, 21 Nov 2023 23:15:29 +0700 Subject: [PATCH] feat: allows to cancel generation (#22) --- Ollamac.xcodeproj/project.pbxproj | 2 +- .../xcshareddata/swiftpm/Package.resolved | 4 +- Ollamac/ViewModels/MessageViewModel.swift | 65 +++++++++++-------- Ollamac/Views/MessageViews/MessageView.swift | 14 +++- 4 files changed, 51 insertions(+), 34 deletions(-) diff --git a/Ollamac.xcodeproj/project.pbxproj b/Ollamac.xcodeproj/project.pbxproj index 64086a2..426424e 100644 --- a/Ollamac.xcodeproj/project.pbxproj +++ b/Ollamac.xcodeproj/project.pbxproj @@ -596,7 +596,7 @@ repositoryURL = "https://github.com/kevinhermawan/OllamaKit"; requirement = { kind = upToNextMajorVersion; - minimumVersion = 1.0.0; + minimumVersion = 2.0.0; }; }; /* End XCRemoteSwiftPackageReference section */ diff --git a/Ollamac.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/Ollamac.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved index c9467d6..1a16db6 100644 --- a/Ollamac.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved +++ b/Ollamac.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved @@ -23,8 +23,8 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/kevinhermawan/OllamaKit", "state" : { - "revision" : "10d467cfa4f0484754bebf3726e59ab8312f545f", - "version" : "1.0.1" + "revision" : "6f3258eb91145bbf26a263123ff88dbbb95208c2", + "version" : "2.0.0" } }, { diff --git a/Ollamac/ViewModels/MessageViewModel.swift b/Ollamac/ViewModels/MessageViewModel.swift index daebbc1..cf0a571 100644 --- a/Ollamac/ViewModels/MessageViewModel.swift +++ b/Ollamac/ViewModels/MessageViewModel.swift @@ -5,6 +5,7 @@ // Created by Kevin Hermawan on 04/11/23. // +import Combine import Foundation import OllamaKit import SwiftData @@ -12,6 +13,8 @@ import ViewState @Observable final class MessageViewModel { + private var generation: AnyCancellable? + private var modelContext: ModelContext private var ollamaKit: OllamaKit @@ -23,6 +26,10 @@ final class MessageViewModel { self.ollamaKit = ollamaKit } + deinit { + self.stopGenerate() + } + func fetch(for chat: Chat) throws { let chatId = chat.id let predicate = #Predicate{ $0.chat?.id == chatId } @@ -41,21 +48,19 @@ final class MessageViewModel { try? modelContext.saveChanges() if await ollamaKit.reachable() { - ollamaKit.generate(data: message.convertToOKGenerateRequestData()) { [weak self] stream in - guard let self = self else { return } - - switch stream.event { - case let .stream(result): - switch result { - case .success(let response): - self.handleSuccess(response) + let data = message.convertToOKGenerateRequestData() + + generation = ollamaKit.generate(data: data) + .sink(receiveCompletion: { [weak self] completion in + switch completion { + case .finished: + self?.handleComplete() case .failure(let error): - self.handleError(error.localizedDescription) + self?.handleError(error.localizedDescription) } - case .complete: - self.handleComplete() - } - } + }, receiveValue: { [weak self] response in + self?.handleReceive(response) + }) } else { self.handleError(AppMessages.ollamaServerUnreachable) } @@ -69,27 +74,31 @@ final class MessageViewModel { try? modelContext.saveChanges() if await ollamaKit.reachable() { - ollamaKit.generate(data: message.convertToOKGenerateRequestData()) { [weak self] stream in - guard let self = self else { return } - - switch stream.event { - case let .stream(result): - switch result { - case .success(let response): - self.handleSuccess(response) + let data = message.convertToOKGenerateRequestData() + + generation = ollamaKit.generate(data: data) + .sink(receiveCompletion: { [weak self] completion in + switch completion { + case .finished: + self?.handleComplete() case .failure(let error): - self.handleError(error.localizedDescription) + self?.handleError(error.localizedDescription) } - case .complete: - self.handleComplete() - } - } + }, receiveValue: { [weak self] response in + self?.handleReceive(response) + }) } else { self.handleError(AppMessages.ollamaServerUnreachable) } } - private func handleSuccess(_ response: OKGenerateResponse) { + func stopGenerate() { + self.sendViewState = nil + self.generation?.cancel() + try? self.modelContext.saveChanges() + } + + private func handleReceive(_ response: OKGenerateResponse) { if self.messages.isEmpty { return } let lastIndex = self.messages.count - 1 @@ -117,7 +126,7 @@ final class MessageViewModel { let lastIndex = self.messages.count - 1 self.messages[lastIndex].error = false self.messages[lastIndex].done = true - + try? self.modelContext.saveChanges() self.sendViewState = nil } diff --git a/Ollamac/Views/MessageViews/MessageView.swift b/Ollamac/Views/MessageViews/MessageView.swift index f9572a9..1fcac9f 100644 --- a/Ollamac/Views/MessageViews/MessageView.swift +++ b/Ollamac/Views/MessageViews/MessageView.swift @@ -79,13 +79,21 @@ struct MessageView: View { .focused($isEditorFocused) .onSubmit(sendAction) + Button(action: messageViewModel.stopGenerate) { + Image(systemName: "stop.circle.fill") + .padding(8) + .help("Stop generate") + } + .buttonStyle(.borderedProminent) + .visible(if: isGenerating, removeCompletely: true) + Button(action: sendAction) { - Label("Send", systemImage: "paperplane.fill") + Image(systemName: "paperplane.fill") .padding(8) - .foregroundStyle(.white) - .help("Send message (Return)") + .help("Send message") } .buttonStyle(.borderedProminent) + .hide(if: isGenerating, removeCompletely: true) } .padding(.horizontal) }