feat: allows to cancel generation (kevinhermawan#22)

tom-ludwig · Nov 21, 2023 · b67195e · b67195e
1 parent 16136c9
commit b67195e
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 34 deletions.
diff --git 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
@@ -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
@@ -5,13 +5,16 @@
 //  Created by Kevin Hermawan on 04/11/23.
 //
 
+import Combine
 import Foundation
 import OllamaKit
 import SwiftData
 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<Message>{ $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
@@ -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)
             }