Skip to content

Commit c009a90

Browse files
intitniintitni
committed
Move memory of ChatGPTService to ChatGPTMemory
1 parent a78d9d7 commit c009a90

23 files changed

+416
-264
lines changed

‎Core/Sources/ChatPlugin/AITerminalChatPlugin.swift‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public actor AITerminalChatPlugin: ChatPlugin {
2727

2828
do {
2929
if let command {
30-
await chatGPTService.mutateHistory { history in
30+
await chatGPTService.memory.mutateHistory { history in
3131
history.append(.init(role: .user, content: content))
3232
}
3333
delegate?.pluginDidStartResponding(self)
@@ -43,14 +43,14 @@ public actor AITerminalChatPlugin: ChatPlugin {
4343
case .cancellation:
4444
delegate?.pluginDidEndResponding(self)
4545
delegate?.pluginDidEnd(self)
46-
await chatGPTService.mutateHistory { history in
46+
await chatGPTService.memory.mutateHistory { history in
4747
history.append(.init(role: .assistant, content: "Cancelled"))
4848
}
4949
case .modification:
5050
let result = try await modifyCommand(command: command, requirement: content)
5151
self.command = result
5252
delegate?.pluginDidEndResponding(self)
53-
await chatGPTService.mutateHistory { history in
53+
await chatGPTService.memory.mutateHistory { history in
5454
history.append(.init(role: .assistant, content: """
5555
Should I run this command? You can instruct me to modify it again.
5656
```
@@ -60,15 +60,15 @@ public actor AITerminalChatPlugin: ChatPlugin {
6060
}
6161
case .other:
6262
delegate?.pluginDidEndResponding(self)
63-
await chatGPTService.mutateHistory { history in
63+
await chatGPTService.memory.mutateHistory { history in
6464
history.append(.init(
6565
role: .assistant,
6666
content: "Sorry, I don't understand. Do you want me to run it?"
6767
))
6868
}
6969
}
7070
} else {
71-
await chatGPTService.mutateHistory { history in
71+
await chatGPTService.memory.mutateHistory { history in
7272
history.append(.init(
7373
role: .user,
7474
content: originalMessage,
@@ -79,7 +79,7 @@ public actor AITerminalChatPlugin: ChatPlugin {
7979
let result = try await generateCommand(task: content)
8080
command = result
8181
if isCancelled { return }
82-
await chatGPTService.mutateHistory { history in
82+
await chatGPTService.memory.mutateHistory { history in
8383
history.append(.init(role: .assistant, content: """
8484
Should I run this command? You can instruct me to modify it.
8585
```
@@ -90,7 +90,7 @@ public actor AITerminalChatPlugin: ChatPlugin {
9090
delegate?.pluginDidEndResponding(self)
9191
}
9292
} catch {
93-
await chatGPTService.mutateHistory { history in
93+
await chatGPTService.memory.mutateHistory { history in
9494
history.append(.init(role: .assistant, content: error.localizedDescription))
9595
}
9696
delegate?.pluginDidEndResponding(self)

‎Core/Sources/ChatPlugin/AskChatGPT.swift‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,12 @@ public func askChatGPT(
77
question: String,
88
temperature: Double? = nil
99
) async throws -> String? {
10+
let configuration = OverridingUserPreferenceChatGPTConfiguration(
11+
overriding: .init(temperature: temperature)
12+
)
13+
let memory = AutoManagedChatGPTMemory(systemPrompt: systemPrompt, configuration: configuration)
1014
let service = ChatGPTService(
11-
systemPrompt: systemPrompt,
15+
memory: memory,
1216
configuration: OverridingUserPreferenceChatGPTConfiguration(
1317
overriding: .init(temperature: temperature)
1418
)

‎Core/Sources/ChatPlugin/CallAIFunction.swift‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,16 @@ func callAIFunction(
1616
}
1717
}
1818
let argsString = args.joined(separator: ", ")
19+
let configuration = OverridingUserPreferenceChatGPTConfiguration(
20+
overriding: .init(temperature: 0)
21+
)
1922
let service = ChatGPTService(
20-
systemPrompt: "You are now the following python function: ```# \(description)\n\(function)```\n\nOnly respond with your `return` value."
23+
memory: AutoManagedChatGPTMemory(
24+
systemPrompt: "You are now the following python function: ```# \(description)\n\(function)```\n\nOnly respond with your `return` value.",
25+
configuration: configuration
26+
),
27+
configuration: configuration
2128
)
2229
return try await service.sendAndWait(content: argsString)
2330
}
31+

‎Core/Sources/ChatPlugin/TerminalChatPlugin.swift‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ public actor TerminalChatPlugin: ChatPlugin {
4242
return try await Environment.guessProjectRootURLForFile(fileURL)
4343
}()
4444

45-
await chatGPTService.mutateHistory { history in
45+
await chatGPTService.memory.mutateHistory { history in
4646
history.append(
4747
.init(
4848
role: .user,
@@ -69,7 +69,7 @@ public actor TerminalChatPlugin: ChatPlugin {
6969

7070
for try await content in output {
7171
if isCancelled { throw CancellationError() }
72-
await chatGPTService.mutateHistory { history in
72+
await chatGPTService.memory.mutateHistory { history in
7373
if history.last?.id == id {
7474
history.removeLast()
7575
}
@@ -78,23 +78,23 @@ public actor TerminalChatPlugin: ChatPlugin {
7878
}
7979
}
8080
outputContent += "\n[finished]"
81-
await chatGPTService.mutateHistory { history in
81+
await chatGPTService.memory.mutateHistory { history in
8282
if history.last?.id == id {
8383
history.removeLast()
8484
}
8585
history.append(message)
8686
}
8787
} catch let error as Terminal.TerminationError {
8888
outputContent += "\n[error: \(error.status)]"
89-
await chatGPTService.mutateHistory { history in
89+
await chatGPTService.memory.mutateHistory { history in
9090
if history.last?.id == id {
9191
history.removeLast()
9292
}
9393
history.append(message)
9494
}
9595
} catch {
9696
outputContent += "\n[error: \(error.localizedDescription)]"
97-
await chatGPTService.mutateHistory { history in
97+
await chatGPTService.memory.mutateHistory { history in
9898
if history.last?.id == id {
9999
history.removeLast()
100100
}

‎Core/Sources/ChatPlugins/MathChatPlugin/MathChatPlugin.swift‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,15 @@ public actor MathChatPlugin: ChatPlugin {
2525
async let translatedAnswer = translate(text: "Answer:")
2626
var reply = ChatMessage(id: id, role: .assistant, content: "")
2727

28-
await chatGPTService.mutateHistory { history in
28+
await chatGPTService.memory.mutateHistory { history in
2929
history.append(.init(role: .user, content: originalMessage, summary: content))
3030
}
3131

3232
do {
3333
let result = try await solveMathProblem(content)
3434
let formattedResult = "\(await translatedAnswer) \(result)"
3535
if !isCancelled {
36-
await chatGPTService.mutateHistory { history in
36+
await chatGPTService.memory.mutateHistory { history in
3737
if history.last?.id == id {
3838
history.removeLast()
3939
}
@@ -43,7 +43,7 @@ public actor MathChatPlugin: ChatPlugin {
4343
}
4444
} catch {
4545
if !isCancelled {
46-
await chatGPTService.mutateHistory { history in
46+
await chatGPTService.memory.mutateHistory { history in
4747
if history.last?.id == id {
4848
history.removeLast()
4949
}

‎Core/Sources/ChatPlugins/SearchChatPlugin/SearchChatPlugin.swift‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,7 @@ public actor SearchChatPlugin: ChatPlugin {
2323
let id = "\(Self.command)-\(UUID().uuidString)"
2424
var reply = ChatMessage(id: id, role: .assistant, content: "")
2525

26-
await chatGPTService.mutateHistory { history in
27-
history.append(.init(role: .user, content: originalMessage, summary: content))
28-
}
26+
await chatGPTService.memory.appendMessage(.init(role: .user, content: originalMessage, summary: content))
2927

3028
do {
3129
let (eventStream, cancelAgent) = try await search(content)
@@ -54,7 +52,7 @@ public actor SearchChatPlugin: ChatPlugin {
5452
"""
5553
}
5654

57-
await chatGPTService.mutateHistory { history in
55+
await chatGPTService.memory.mutateHistory { history in
5856
if history.last?.id == id {
5957
history.removeLast()
6058
}
@@ -77,7 +75,7 @@ public actor SearchChatPlugin: ChatPlugin {
7775
}
7876

7977
} catch {
80-
await chatGPTService.mutateHistory { history in
78+
await chatGPTService.memory.mutateHistory { history in
8179
if history.last?.id == id {
8280
history.removeLast()
8381
}

‎Core/Sources/ChatPlugins/ShortcutChatPlugin/ShortcutChatPlugin.swift‎

Lines changed: 12 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -42,23 +42,17 @@ public actor ShortcutChatPlugin: ChatPlugin {
4242
guard let shortcutName, !shortcutName.isEmpty else {
4343
message.content =
4444
"Please provide the shortcut name in format: `/\(Self.command)(shortcut name)`."
45-
await chatGPTService.mutateHistory { history in
46-
history.append(message)
47-
}
45+
await chatGPTService.memory.appendMessage(message)
4846
return
4947
}
5048

5149
var input = String(content).trimmingCharacters(in: .whitespacesAndNewlines)
5250
if input.isEmpty {
5351
// if no input detected, use the previous message as input
54-
input = await chatGPTService.history.last?.content ?? ""
55-
await chatGPTService.mutateHistory { history in
56-
history.append(.init(role: .user, content: originalMessage))
57-
}
52+
input = await chatGPTService.memory.messages.last?.content ?? ""
53+
await chatGPTService.memory.appendMessage(.init(role: .user, content: originalMessage))
5854
} else {
59-
await chatGPTService.mutateHistory { history in
60-
history.append(.init(role: .user, content: originalMessage))
61-
}
55+
await chatGPTService.memory.appendMessage(.init(role: .user, content: originalMessage))
6256
}
6357

6458
do {
@@ -71,7 +65,7 @@ public actor ShortcutChatPlugin: ChatPlugin {
7165
.appendingPathComponent("\(id)-input.txt")
7266
let temporaryOutputFileURL = temporaryURL
7367
.appendingPathComponent("\(id)-output")
74-
68+
7569
try input.write(to: temporaryInputFileURL, atomically: true, encoding: .utf8)
7670

7771
let command = """
@@ -86,7 +80,7 @@ public actor ShortcutChatPlugin: ChatPlugin {
8680
currentDirectoryPath: "/",
8781
environment: [:]
8882
)
89-
83+
9084
await Task.yield()
9185

9286
if FileManager.default.fileExists(atPath: temporaryOutputFileURL.path) {
@@ -96,33 +90,25 @@ public actor ShortcutChatPlugin: ChatPlugin {
9690
if text.isEmpty {
9791
message.content = "Finished"
9892
}
99-
await chatGPTService.mutateHistory { history in
100-
history.append(message)
101-
}
93+
await chatGPTService.memory.appendMessage(message)
10294
} else {
10395
message.content = """
10496
[View File](\(temporaryOutputFileURL))
10597
"""
106-
await chatGPTService.mutateHistory { history in
107-
history.append(message)
108-
}
98+
await chatGPTService.memory.appendMessage(message)
10999
}
110-
100+
111101
return
112102
}
113-
103+
114104
message.content = "Finished"
115-
await chatGPTService.mutateHistory { history in
116-
history.append(message)
117-
}
105+
await chatGPTService.memory.appendMessage(message)
118106
} catch {
119107
message.content = error.localizedDescription
120108
if error.localizedDescription.isEmpty {
121109
message.content = "Error"
122110
}
123-
await chatGPTService.mutateHistory { history in
124-
history.append(message)
125-
}
111+
await chatGPTService.memory.appendMessage(message)
126112
}
127113
}
128114

‎Core/Sources/ChatPlugins/ShortcutChatPlugin/ShortcutInputChatPlugin.swift‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -45,16 +45,14 @@ public actor ShortcutInputChatPlugin: ChatPlugin {
4545
role: .assistant,
4646
content: "Please provide the shortcut name in format: `/\(Self.command)(shortcut name)`."
4747
)
48-
await chatGPTService.mutateHistory { history in
49-
history.append(reply)
50-
}
48+
await chatGPTService.memory.appendMessage(reply)
5149
return
5250
}
5351

5452
var input = String(content).trimmingCharacters(in: .whitespacesAndNewlines)
5553
if input.isEmpty {
5654
// if no input detected, use the previous message as input
57-
input = await chatGPTService.history.last?.content ?? ""
55+
input = await chatGPTService.memory.messages.last?.content ?? ""
5856
}
5957

6058
do {
@@ -108,9 +106,7 @@ public actor ShortcutInputChatPlugin: ChatPlugin {
108106
role: .assistant,
109107
content: error.localizedDescription
110108
)
111-
await chatGPTService.mutateHistory { history in
112-
history.append(reply)
113-
}
109+
await chatGPTService.memory.appendMessage(reply)
114110
}
115111
}
116112

‎Core/Sources/ChatService/ChatPluginController.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ final class ChatPluginController {
3535
if command == "exit" {
3636
if let plugin = runningPlugin {
3737
runningPlugin = nil
38-
_ = await chatGPTService.mutateHistory { history in
38+
_ = await chatGPTService.memory.mutateHistory { history in
3939
history.append(.init(
4040
role: .user,
4141
content: "",
@@ -48,7 +48,7 @@ final class ChatPluginController {
4848
))
4949
}
5050
} else {
51-
_ = await chatGPTService.mutateHistory { history in
51+
_ = await chatGPTService.memory.mutateHistory { history in
5252
history.append(.init(
5353
role: .system,
5454
content: "",

0 commit comments

Comments
 (0)