Cache token count

Author: intitni

Core/Sources/OpenAIService/ChatGPTService.swift

@@ -270,10 +270,11 @@ extension ChatGPTService {
     {
         var all: [CompletionRequestBody.Message] = []
         var allTokensCount = encoder.encode(text: systemPrompt).count
-        for message in history.reversed() {
+        for (index, message) in history.enumerated().reversed() {
             if maxNumberOfMessages > 0, all.count >= maxNumberOfMessages { break }
             if message.content.isEmpty { continue }
-            let tokensCount = encoder.encode(text: message.content).count
+            let tokensCount = message.tokensCount ?? encoder.encode(text: message.content).count
+            history[index].tokensCount = tokensCount
             if tokensCount + allTokensCount > maxTokens - minimumReplyTokens {
                 break
             }

Core/Sources/OpenAIService/Models.swift

@@ -15,9 +15,14 @@ public struct ChatMessage: Equatable, Codable {
     }
 
     public var role: Role
-    public var content: String
+    public var content: String {
+        didSet {
+            tokensCount = nil
+        }
+    }
     public var summary: String?
     public var id: String
+    public var tokensCount: Int?
 
     public init(
         id: String = UUID().uuidString,

Core/Tests/OpenAIServiceTests/ChatGPTServiceTests.swift

@@ -62,7 +62,10 @@ final class ChatGPTServiceTests: XCTestCase {
             .init(role: .user, content: "Hello"),
         ], "System prompt is included")
         XCTAssertEqual(all, ["hello", "my", "friends"], "Text stream is correct")
-        let history = await service.history
+        var history = await service.history
+        for (i, _) in history.enumerated() {
+            history[i].tokensCount = nil
+        }
         XCTAssertEqual(history, [
             .init(id: "0", role: .user, content: "Hello"),
             .init(id: "1", role: .assistant, content: "hellomyfriends"),

Core/Tests/OpenAIServiceTests/LimitMessagesTests.swift

@@ -23,10 +23,16 @@ final class LimitMessagesTests: XCTestCase {
             "hello",
             "world",
         ])
-        
+
         XCTAssertEqual(remainingTokens, 10000 - 12 - 6)
+        let history = await service.history
+        XCTAssertEqual(history.map(\.tokensCount), [
+            2,
+            5,
+            5,
+        ])
     }
-    
+
     func test_send_max_message_if_not_reached_token_limit() async {
         let service = await createService(systemPrompt: "system", messages: [
             "hi",
@@ -45,10 +51,10 @@ final class LimitMessagesTests: XCTestCase {
             "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",
@@ -66,10 +72,10 @@ final class LimitMessagesTests: XCTestCase {
             "system",
             "world",
         ])
-        
+
         XCTAssertEqual(remainingTokens, 201)
     }
-    
+
     func test_minimum_reply_tokens_count() async {
         let service = await createService(systemPrompt: "system", messages: [
             "hi",
@@ -86,7 +92,7 @@ final class LimitMessagesTests: XCTestCase {
         XCTAssertEqual(messages, [
             "system",
         ])
-        
+
         XCTAssertEqual(remainingTokens, 200)
     }
 }