Support auto token count management

Author: intitni

Core/Sources/OpenAIService/ChatGPTService.swift

@@ -1,5 +1,6 @@
 import AsyncAlgorithms
 import Foundation
+import GPTEncoder
 import Preferences
 
 public protocol ChatGPTServiceType: ObservableObject {
@@ -112,12 +113,14 @@ public actor ChatGPTService: ChatGPTServiceType {
         )
         history.append(newMessage)
 
+        let (messages, remainingTokens) = combineHistoryWithSystemPrompt()
+        
         let requestBody = CompletionRequestBody(
             model: model,
-            messages: combineHistoryWithSystemPrompt(),
+            messages: messages,
             temperature: temperature,
             stream: true,
-            max_tokens: maxToken
+            max_tokens: remainingTokens
         )
 
         isReceivingMessage = true
@@ -190,12 +193,14 @@ public actor ChatGPTService: ChatGPTServiceType {
         )
         history.append(newMessage)
 
+        let (messages, remainingTokens) = combineHistoryWithSystemPrompt()
+        
         let requestBody = CompletionRequestBody(
             model: model,
-            messages: combineHistoryWithSystemPrompt(),
+            messages: messages,
             temperature: temperature,
             stream: true,
-            max_tokens: maxToken
+            max_tokens: remainingTokens
         )
 
         isReceivingMessage = true
@@ -210,10 +215,10 @@ public actor ChatGPTService: ChatGPTServiceType {
                 role: choice.message.role,
                 content: choice.message.content
             ))
-            
+
             return choice.message.content
         }
-        
+
         return nil
     }
 
@@ -250,17 +255,34 @@ extension ChatGPTService {
         uuidGenerator = generator
     }
 
-    func combineHistoryWithSystemPrompt() -> [CompletionRequestBody.Message] {
+    func combineHistoryWithSystemPrompt(
+        minimumReplyTokens: Int = 200,
+        maxNumberOfMessages: Int = 5,
+        maxTokens: Int =  UserDefaults.shared.value(for: \.chatGPTMaxToken),
+        encoder: TokenEncoder = GPTEncoder()
+    )
+        -> (messages: [CompletionRequestBody.Message], remainingTokens: Int)
+    {
         var all: [CompletionRequestBody.Message] = []
-        var count = 0
+        var allTokensCount = encoder.encode(text: systemPrompt).count
         for message in history.reversed() {
-            if count >= 5 { break }
+            if all.count >= maxNumberOfMessages { break }
             if message.content.isEmpty { continue }
+            let tokensCount = encoder.encode(text: message.content).count
+            if tokensCount + allTokensCount > maxTokens - minimumReplyTokens {
+                break
+            }
+            allTokensCount += tokensCount
             all.append(.init(role: message.role, content: message.content))
-            count += 1
         }
 
         all.append(.init(role: .system, content: systemPrompt))
-        return all.reversed()
+        return (all.reversed(), max(minimumReplyTokens, maxTokens - allTokensCount))
     }
 }
+
+protocol TokenEncoder {
+    func encode(text: String) -> [Int]
+}
+
+extension GPTEncoder: TokenEncoder {}

Core/Tests/OpenAIServiceTests/LimitMessagesTests.swift

@@ -0,0 +1,124 @@
+import Foundation
+import XCTest
+
+@testable import OpenAIService
+
+final class LimitMessagesTests: XCTestCase {
+    func test_send_all_messages_if_not_reached_token_limit() async {
+        let service = await createService(systemPrompt: "system", messages: [
+            "hi",
+            "hello",
+            "world",
+        ])
+
+        let (messages, remainingTokens) = await runService(
+            service,
+            minimumReplyTokens: 200,
+            maxNumberOfMessages: 100,
+            maxTokens: 10000
+        )
+        XCTAssertEqual(messages, [
+            "system",
+            "hi",
+            "hello",
+            "world",
+        ])
+        
+        XCTAssertEqual(remainingTokens, 10000 - 12 - 6)
+    }
+    
+    func test_send_max_message_if_not_reached_token_limit() async {
+        let service = await createService(systemPrompt: "system", messages: [
+            "hi",
+            "hello",
+            "world",
+        ])
+
+        let (messages, remainingTokens) = await runService(
+            service,
+            minimumReplyTokens: 200,
+            maxNumberOfMessages: 2,
+            maxTokens: 10000
+        )
+        XCTAssertEqual(messages, [
+            "system",
+            "hello",
+            "world",
+        ], "Count from end to start.")
+        
+        XCTAssertEqual(remainingTokens, 10000 - 10 - 6)
+    }
+    
+    func test_reached_token_limit() async {
+        let service = await createService(systemPrompt: "system", messages: [
+            "hi",
+            "hello",
+            "world",
+        ])
+
+        let (messages, remainingTokens) = await runService(
+            service,
+            minimumReplyTokens: 200,
+            maxNumberOfMessages: 100,
+            maxTokens: 212
+        )
+        XCTAssertEqual(messages, [
+            "system",
+            "world",
+        ])
+        
+        XCTAssertEqual(remainingTokens, 201)
+    }
+    
+    func test_minimum_reply_tokens_count() async {
+        let service = await createService(systemPrompt: "system", messages: [
+            "hi",
+            "hello",
+            "world",
+        ])
+
+        let (messages, remainingTokens) = await runService(
+            service,
+            minimumReplyTokens: 200,
+            maxNumberOfMessages: 100,
+            maxTokens: 200
+        )
+        XCTAssertEqual(messages, [
+            "system",
+        ])
+        
+        XCTAssertEqual(remainingTokens, 200)
+    }
+}
+
+class MockEncoder: TokenEncoder {
+    func encode(text: String) -> [Int] {
+        return .init(repeating: 0, count: text.count)
+    }
+}
+
+private func createService(systemPrompt: String, messages: [String]) async -> ChatGPTService {
+    let service = ChatGPTService(systemPrompt: systemPrompt)
+    await service.mutateHistory { history in
+        messages.forEach { message in
+            history.append(.init(role: .user, content: message))
+        }
+    }
+    return service
+}
+
+private func runService(
+    _ service: ChatGPTService,
+    minimumReplyTokens: Int,
+    maxNumberOfMessages: Int,
+    maxTokens: Int
+) async -> (messages: [String], remainingTokens: Int) {
+    let (messages, remainingTokens) = await service.combineHistoryWithSystemPrompt(
+        minimumReplyTokens: minimumReplyTokens,
+        maxNumberOfMessages: maxNumberOfMessages,
+        maxTokens: maxTokens,
+        encoder: MockEncoder()
+    )
+
+    return (messages.map(\.content), remainingTokens)
+}