Skip to content

Commit 2b73ce6

Browse files
intitniintitni
committed
Move isReceivingMessage to ChatService
1 parent 84723cc commit 2b73ce6

4 files changed

Lines changed: 15 additions & 35 deletions

File tree

‎Core/Sources/ChatService/ChatPluginController.swift‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ final class ChatPluginController {
77
let chatGPTService: any ChatGPTServiceType
88
let plugins: [String: ChatPlugin.Type]
99
var runningPlugin: ChatPlugin?
10+
weak var chatService: ChatService?
1011

1112
init(chatGPTService: any ChatGPTServiceType, plugins: [ChatPlugin.Type]) {
1213
self.chatGPTService = chatGPTService
@@ -107,15 +108,11 @@ final class ChatPluginController {
107108

108109
extension ChatPluginController: ChatPluginDelegate {
109110
public func pluginDidStartResponding(_: ChatPlugin) {
110-
Task {
111-
await chatGPTService.markReceivingMessage(true)
112-
}
111+
chatService?.isReceivingMessage = true
113112
}
114113

115114
public func pluginDidEndResponding(_: ChatPlugin) {
116-
Task {
117-
await chatGPTService.markReceivingMessage(false)
118-
}
115+
chatService?.isReceivingMessage = false
119116
}
120117

121118
public func pluginDidStart(_ plugin: ChatPlugin) {

‎Core/Sources/ChatService/ChatService.swift‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,11 @@ import OpenAIService
66

77
public final class ChatService: ObservableObject {
88
public let chatGPTService: any ChatGPTServiceType
9+
public var allPluginCommands: [String] { allPlugins.map { $0.command } }
910
let pluginController: ChatPluginController
1011
let contextController: DynamicContextController
1112
var cancellable = Set<AnyCancellable>()
13+
@Published public internal(set) var isReceivingMessage = false
1214
@Published public internal(set) var systemPrompt = UserDefaults.shared
1315
.value(for: \.defaultChatSystemPrompt)
1416
@Published public internal(set) var extraSystemPrompt = ""
@@ -21,30 +23,37 @@ public final class ChatService: ObservableObject {
2123
contextCollectors: ActiveDocumentChatContextCollector()
2224
)
2325

26+
pluginController.chatService = self
2427
chatGPTService.objectWillChange.sink { [weak self] _ in
2528
self?.objectWillChange.send()
2629
}.store(in: &cancellable)
2730
}
2831

2932
public func send(content: String) async throws {
33+
guard !isReceivingMessage else { throw CancellationError() }
3034
let handledInPlugin = try await pluginController.handleContent(content)
3135
if handledInPlugin { return }
3236
try await contextController.updatePromptToMatchContent(systemPrompt: """
3337
\(systemPrompt)
3438
\(extraSystemPrompt)
3539
""", content: content)
3640

37-
_ = try await chatGPTService.send(content: content, summary: nil)
41+
let stream = try await chatGPTService.send(content: content, summary: nil)
42+
isReceivingMessage = true
43+
for try await _ in stream {}
44+
isReceivingMessage = false
3845
}
3946

4047
public func stopReceivingMessage() async {
4148
await pluginController.stopResponding()
4249
await chatGPTService.stopReceivingMessage()
50+
isReceivingMessage = false
4351
}
4452

4553
public func clearHistory() async {
4654
await pluginController.cancel()
4755
await chatGPTService.clearHistory()
56+
isReceivingMessage = false
4857
}
4958

5059
public func resetPrompt() async {

‎Core/Sources/Service/GUI/ChatProvider+Service.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ extension ChatProvider {
1111
onCloseChat: @escaping () -> Void,
1212
onSwitchContext: @escaping () -> Void
1313
) {
14-
self.init()
14+
self.init(pluginIdentifiers: service.allPluginCommands)
1515

1616
let cancellable = service.objectWillChange.sink { [weak self] in
1717
guard let self else { return }
@@ -23,7 +23,7 @@ extension ChatProvider {
2323
text: message.summary ?? message.content
2424
)
2525
}
26-
self.isReceivingMessage = await service.chatGPTService.isReceivingMessage
26+
self.isReceivingMessage = service.isReceivingMessage
2727
self.systemPrompt = service.systemPrompt
2828
self.extraSystemPrompt = service.extraSystemPrompt
2929
}

‎Tool/Sources/OpenAIService/ChatGPTService.swift‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,12 @@ import GPTEncoder
44
import Preferences
55

66
public protocol ChatGPTServiceType: ObservableObject {
7-
var isReceivingMessage: Bool { get async }
87
var history: [ChatMessage] { get async }
98
func send(content: String, summary: String?) async throws -> AsyncThrowingStream<String, Error>
109
func stopReceivingMessage() async
1110
func clearHistory() async
1211
func mutateSystemPrompt(_ newPrompt: String) async
1312
func mutateHistory(_ mutate: (inout [ChatMessage]) -> Void) async
14-
func markReceivingMessage(_ receiving: Bool) async
1513
}
1614

1715
public enum ChatGPTServiceError: Error, LocalizedError {
@@ -105,10 +103,6 @@ public actor ChatGPTService: ChatGPTServiceType {
105103
didSet { objectWillChange.send() }
106104
}
107105

108-
public internal(set) var isReceivingMessage = false {
109-
didSet { objectWillChange.send() }
110-
}
111-
112106
var stop: [String]
113107
var uuidGenerator: () -> String = { UUID().uuidString }
114108
var cancelTask: Cancellable?
@@ -131,7 +125,6 @@ public actor ChatGPTService: ChatGPTServiceType {
131125
content: String,
132126
summary: String? = nil
133127
) async throws -> AsyncThrowingStream<String, Error> {
134-
guard !isReceivingMessage else { throw CancellationError() }
135128
guard let url = URL(string: endpoint) else { throw ChatGPTServiceError.endpointIncorrect }
136129

137130
if !content.isEmpty || summary != nil {
@@ -155,8 +148,6 @@ public actor ChatGPTService: ChatGPTServiceType {
155148
max_tokens: maxTokenForReply(model: model, remainingTokens: remainingTokens)
156149
)
157150

158-
isReceivingMessage = true
159-
160151
let api = buildCompletionStreamAPI(
161152
apiKey,
162153
designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider),
@@ -168,10 +159,6 @@ public actor ChatGPTService: ChatGPTServiceType {
168159
Task {
169160
do {
170161
let (trunks, cancel) = try await api()
171-
guard isReceivingMessage else {
172-
continuation.finish()
173-
return
174-
}
175162
cancelTask = cancel
176163
for try await trunk in trunks {
177164
guard let delta = trunk.choices.first?.delta else { continue }
@@ -199,19 +186,15 @@ public actor ChatGPTService: ChatGPTServiceType {
199186
}
200187

201188
continuation.finish()
202-
isReceivingMessage = false
203189
} catch let error as CancellationError {
204-
isReceivingMessage = false
205190
continuation.finish(throwing: error)
206191
} catch let error as NSError where error.code == NSURLErrorCancelled {
207-
isReceivingMessage = false
208192
continuation.finish(throwing: error)
209193
} catch {
210194
history.append(.init(
211195
role: .assistant,
212196
content: error.localizedDescription
213197
))
214-
isReceivingMessage = false
215198
continuation.finish(throwing: error)
216199
}
217200
}
@@ -222,7 +205,6 @@ public actor ChatGPTService: ChatGPTServiceType {
222205
content: String,
223206
summary: String? = nil
224207
) async throws -> String? {
225-
guard !isReceivingMessage else { throw CancellationError() }
226208
guard let url = URL(string: endpoint) else { throw ChatGPTServiceError.endpointIncorrect }
227209

228210
if !content.isEmpty || summary != nil {
@@ -246,9 +228,6 @@ public actor ChatGPTService: ChatGPTServiceType {
246228
max_tokens: maxTokenForReply(model: model, remainingTokens: remainingTokens)
247229
)
248230

249-
isReceivingMessage = true
250-
defer { isReceivingMessage = false }
251-
252231
let api = buildCompletionAPI(
253232
apiKey,
254233
designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider),
@@ -273,7 +252,6 @@ public actor ChatGPTService: ChatGPTServiceType {
273252
public func stopReceivingMessage() {
274253
cancelTask?()
275254
cancelTask = nil
276-
isReceivingMessage = false
277255
}
278256

279257
public func clearHistory() {
@@ -288,10 +266,6 @@ public actor ChatGPTService: ChatGPTServiceType {
288266
public func mutateHistory(_ mutate: (inout [ChatMessage]) -> Void) async {
289267
mutate(&history)
290268
}
291-
292-
public func markReceivingMessage(_ receiving: Bool) {
293-
isReceivingMessage = receiving
294-
}
295269
}
296270

297271
extension ChatGPTService {

0 commit comments

Comments
 (0)