Skip to content

Commit 335642f

Browse files
committed
Fix the prompt truncation is not ended after cancellation
1 parent 9a508d7 commit 335642f

1 file changed

Lines changed: 3 additions & 0 deletions

File tree

Tool/Sources/OpenAIService/Memory/TemplateChatGPTMemory.swift

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ public actor TemplateChatGPTMemory: ChatGPTMemory {
4646

4747
while !(await checkTokenCount()) {
4848
do {
49+
try Task.checkCancellation()
4950
try await memoryTemplate.truncate()
5051
} catch {
5152
Logger.service.error("Failed to truncate prompt template: \(error)")
@@ -180,6 +181,8 @@ public struct MemoryTemplate {
180181
}
181182

182183
mutating func truncate() async throws {
184+
if Task.isCancelled { return }
185+
183186
if let truncateRule = truncateRule {
184187
try await truncateRule(&messages, &followUpMessages)
185188
return

0 commit comments

Comments
 (0)