Add Cancel Streaming response via Task Cancellation

Author: alfianlosari

Shared/ChatGPTAPI.swift

@@ -77,6 +77,7 @@ class ChatGPTAPI: @unchecked Sendable {
         urlRequest.httpBody = try jsonBody(text: text)
 
         let (result, response) = try await urlSession.bytes(for: urlRequest)
+        try Task.checkCancellation()
 
         guard let httpResponse = response as? HTTPURLResponse else {
             throw "Invalid response"
@@ -85,6 +86,7 @@ class ChatGPTAPI: @unchecked Sendable {
         guard 200...299 ~= httpResponse.statusCode else {
             var errorText = ""
             for try await line in result.lines {
+                try Task.checkCancellation()
                 errorText += line
             }
 
@@ -95,26 +97,21 @@ class ChatGPTAPI: @unchecked Sendable {
             throw "Bad Response: \(httpResponse.statusCode), \(errorText)"
         }
 
-        return AsyncThrowingStream<String, Error> { continuation in
-            Task(priority: .userInitiated) { [weak self] in
-                guard let self else { return }
-                do {
-                    var responseText = ""
-                    for try await line in result.lines {
-                        if line.hasPrefix("data: "),
-                           let data = line.dropFirst(6).data(using: .utf8),
-                           let response = try? self.jsonDecoder.decode(StreamCompletionResponse.self, from: data),
-                           let text = response.choices.first?.delta.content {
-                            responseText += text
-                            continuation.yield(text)
-                        }
-                    }
-                    self.appendToHistoryList(userText: text, responseText: responseText)
-                    continuation.finish()
-                } catch {
-                    continuation.finish(throwing: error)
+        var responseText = ""
+        return AsyncThrowingStream { [weak self] in
+            guard let self else { return nil }
+            for try await line in result.lines {
+                try Task.checkCancellation()
+                if line.hasPrefix("data: "),
+                   let data = line.dropFirst(6).data(using: .utf8),
+                   let response = try? self.jsonDecoder.decode(StreamCompletionResponse.self, from: data),
+                   let text = response.choices.first?.delta.content {
+                    responseText += text
+                    return text
                 }
             }
+            self.appendToHistoryList(userText: text, responseText: responseText)
+            return nil
         }
     }
 
@@ -123,7 +120,7 @@ class ChatGPTAPI: @unchecked Sendable {
         urlRequest.httpBody = try jsonBody(text: text, stream: false)
 
         let (data, response) = try await urlSession.data(for: urlRequest)
-        
+        try Task.checkCancellation()
         guard let httpResponse = response as? HTTPURLResponse else {
             throw "Invalid response"
         }

Shared/ContentView.swift

@@ -73,7 +73,18 @@ struct ContentView: View {
                 .disabled(vm.isInteractingWithChatGPT)
 
             if vm.isInteractingWithChatGPT {
+                #if os(iOS)
+                Button {
+                    vm.cancelStreamingResponse()
+                } label: {
+                    Image(systemName: "stop.circle.fill")
+                        .font(.system(size: 30))
+                        .symbolRenderingMode(.multicolor)
+                        .foregroundColor(.red)
+                }
+                #else
                 DotLoadingView().frame(width: 60, height: 30)
+                #endif
             } else {
                 Button {
                     Task { @MainActor in

Shared/MessageRow.swift

@@ -33,7 +33,7 @@ struct MessageRow: Identifiable {
     var isInteractingWithChatGPT: Bool
 
     let sendImage: String
-    let send: MessageRowType
+    var send: MessageRowType
     var sendText: String {
         send.text
     }

Shared/ViewModel.swift

@@ -14,6 +14,7 @@ class ViewModel: ObservableObject {
     @Published var isInteractingWithChatGPT = false
     @Published var messages: [MessageRow] = []
     @Published var inputMessage: String = ""
+    var task: Task<Void, Never>?
 
     #if !os(watchOS)
     private var synthesizer: AVSpeechSynthesizer?
@@ -32,11 +33,15 @@ class ViewModel: ObservableObject {
 
     @MainActor
     func sendTapped() async {
-        let text = inputMessage
-        inputMessage = ""
         #if os(iOS)
-        await sendAttributed(text: text)
+        self.task = Task {
+            let text = inputMessage
+            inputMessage = ""
+            await sendAttributed(text: text)
+        }
         #else
+        let text = inputMessage
+        inputMessage = ""
         await send(text: text)
         #endif
     }
@@ -52,47 +57,61 @@ class ViewModel: ObservableObject {
 
     @MainActor
     func retry(message: MessageRow) async {
+        #if os(iOS)
+        self.task = Task {
+            guard let index = messages.firstIndex(where: { $0.id == message.id }) else {
+                return
+            }
+            self.messages.remove(at: index)
+            await sendAttributed(text: message.sendText)
+        }
+        #else
         guard let index = messages.firstIndex(where: { $0.id == message.id }) else {
             return
         }
         self.messages.remove(at: index)
-        #if os(iOS)
-        await sendAttributed(text: message.sendText)
-        #else
         await send(text: message.sendText)
         #endif
     }
 
+    func cancelStreamingResponse() {
+        self.task?.cancel()
+        self.task = nil
+    }
+    
     #if os(iOS)
     @MainActor
     private func sendAttributed(text: String) async {
         isInteractingWithChatGPT = true
-        
-        let parsingTask = ResponseParsingTask()
-        let attributedSend = await parsingTask.parse(text: text)
-        
         var streamText = ""
+        
         var messageRow = MessageRow(
             isInteractingWithChatGPT: true,
             sendImage: "profile",
-            send: .attributed(attributedSend),
+            send: .rawText(text),
             responseImage: "openai",
             responseError: nil)
-        
-        self.messages.append(messageRow)
-        
-        let parserThresholdTextCount = 64
-        var currentTextCount = 0
-        var currentOutput: AttributedOutput?
-        
+    
         do {
+            let parsingTask = ResponseParsingTask()
+            let attributedSend = await parsingTask.parse(text: text)
+            try Task.checkCancellation()
+            messageRow.send = .attributed(attributedSend)
+            
+            self.messages.append(messageRow)
+            
+            let parserThresholdTextCount = 64
+            var currentTextCount = 0
+            var currentOutput: AttributedOutput?
+            
             let stream = try await api.sendMessageStream(text: text)
             for try await text in stream {
                 streamText += text
                 currentTextCount += text.count
 
                 if currentTextCount >= parserThresholdTextCount || text.contains("```") {
                     currentOutput = await parsingTask.parse(text: streamText)
+                    try Task.checkCancellation()
                     currentTextCount = 0
                 }
 
@@ -115,17 +134,22 @@ class ViewModel: ObservableObject {
                 }
 
                 self.messages[self.messages.count - 1] = messageRow
+                if let currentString = currentOutput?.string, currentString != streamText {
+                    let output = await parsingTask.parse(text: streamText)
+                    try Task.checkCancellation()
+                    messageRow.response = .attributed(output)
+                }
             }
+        } catch is CancellationError {
+            messageRow.responseError = "The response was cancelled"
         } catch {
             messageRow.responseError = error.localizedDescription
-            messageRow.response = .rawText(streamText)
         }
 
-        if let currentString = currentOutput?.string, currentString != streamText {
-            let output = await parsingTask.parse(text: streamText)
-            messageRow.response = .attributed(output)
+        if messageRow.response == nil {
+            messageRow.response = .rawText(streamText)
         }
-        
+  
         messageRow.isInteractingWithChatGPT = false
         self.messages[self.messages.count - 1] = messageRow
         isInteractingWithChatGPT = false