Skip to content

Commit

Permalink
🛠️ Fix bug when request body can be composed improperly with streamin…
Browse files Browse the repository at this point in the history
…g response ON
  • Loading branch information
Renset committed May 22, 2024
1 parent 66ac51a commit ae6b638
Showing 1 changed file with 25 additions and 12 deletions.
37 changes: 25 additions & 12 deletions macai/UI/Chat/ChatView.swift
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ struct ChatView: View {
@State private var editSystemMessage: Bool = false
@State private var isStreaming: Bool = false
@State private var isHovered = false
@State private var currentStreamingMessage: String = ""
@StateObject private var store = ChatStore(persistenceController: PersistenceController.shared)
@AppStorage("useChatGptForNames") var useChatGptForNames: Bool = false
@AppStorage("useStream") var useStream: Bool = true
Expand Down Expand Up @@ -185,12 +186,11 @@ extension ChatView {
return
}

let defaultRole = "assistant"
let dataString = String(data: data, encoding: .utf8)
if (dataString == "[DONE]") {
handleChatCompletion()
handleChatCompletion(role: defaultRole)
return
} else {
print(dataString)
}

do {
Expand All @@ -203,11 +203,12 @@ extension ChatView {
let delta = firstChoice["delta"] as? [String: String],
let contentPart = delta["content"] {
self.isStreaming = true
self.currentStreamingMessage += contentPart
DispatchQueue.main.async {
self.updateUIWithResponse(content: contentPart, role: "assistant")
self.updateUIWithResponse(content: self.currentStreamingMessage, role: defaultRole)
}
if let finishReason = firstChoice["finish_reason"] as? String, finishReason == "stop" {
handleChatCompletion()
handleChatCompletion(role: defaultRole)
}
}

Expand All @@ -217,12 +218,13 @@ extension ChatView {
let done = dict["done"] as? Bool,
let messageContent = message["content"] as? String {
self.isStreaming = true
self.currentStreamingMessage += messageContent
DispatchQueue.main.async {
self.updateUIWithResponse(content: messageContent, role: messageRole)
self.updateUIWithResponse(content: self.currentStreamingMessage, role: messageRole)
}

if done {
handleChatCompletion()
handleChatCompletion(role: messageRole)
}
}
}
Expand All @@ -233,16 +235,22 @@ extension ChatView {
}
}

private func handleChatCompletion() {
private func handleChatCompletion(role: String) {
print("Chat interaction completed.")
DispatchQueue.main.async {
self.resetCurrentMessage()
// TODO: force the child view for update to reflect the new state
self.isStreaming = false
self.addNewMessageToRequestMessages(content: self.currentStreamingMessage, role: role)
resetCurrentStreamingMessage()
generateChatNameIfNeeded()
}
}

private func resetCurrentStreamingMessage() {
self.currentStreamingMessage = ""
}

func resetCurrentMessage() {
self.viewContext.rollback()
}
Expand All @@ -254,7 +262,7 @@ extension ChatView {
saveNewMessageInStore(with: messageBody)
}

let request = prepareRequest(with: messageBody)
resetCurrentStreamingMessage()
self.waitingForResponse = true

if useStream {
Expand All @@ -271,6 +279,7 @@ extension ChatView {
}
}
} else {
let request = prepareRequest(with: messageBody)
send(using: request) { data, response, error in
processResponse(with: data, response: response, error: error)
generateChatNameIfNeeded()
Expand Down Expand Up @@ -478,7 +487,7 @@ extension ChatView {
if lastMessage.own {
addNewMessageToChat(content: content, role: role)
} else {
lastMessage.body += content
lastMessage.body = content
lastMessage.own = false
lastMessage.timestamp = Date()
lastMessage.waitingForResponse = false
Expand All @@ -487,9 +496,9 @@ extension ChatView {
self.viewContext.saveWithRetry(attempts: 3)
}
}

} else {
addNewMessageToChat(content: content, role: role)
addNewMessageToRequestMessages(content: content, role: role)
}

}
Expand All @@ -508,13 +517,17 @@ extension ChatView {
self.chat.updatedDate = Date()
self.chat.addToMessages(receivedMessage)
self.viewContext.saveWithRetry(attempts: 3)
self.chat.requestMessages.append(["role": role, "content": content])


self.chat.objectWillChange.send()
#if DEBUG
print("Value of choices[\(self.chat.requestMessages.count - 1)].message.content: \(content)")
#endif
}

private func addNewMessageToRequestMessages(content: String, role: String) {
self.chat.requestMessages.append(["role": role, "content": content])
}

}

Expand Down

0 comments on commit ae6b638

Please sign in to comment.