Print token usage in DEBUG mode

Author: intitni

Tool/Sources/OpenAIService/EmbeddingService.swift

@@ -1,4 +1,5 @@
 import Foundation
+import Logger
 
 public struct EmbeddingResponse: Decodable {
     public struct Object: Decodable {
@@ -74,9 +75,19 @@ public struct EmbeddingService {
                 .otherError(String(data: result, encoding: .utf8) ?? "Unknown Error")
         }
 
-        return try JSONDecoder().decode(EmbeddingResponse.self, from: result)
+        let embeddingResponse = try JSONDecoder().decode(EmbeddingResponse.self, from: result)
+        #if DEBUG
+        Logger.service.info("""
+        Embedding usage
+        - number of strings: \(text.count)
+        - prompt tokens: \(embeddingResponse.usage.prompt_tokens)
+        - total tokens: \(embeddingResponse.usage.total_tokens)
+
+        """)
+        #endif
+        return embeddingResponse
     }
-    
+
     public func embed(tokens: [[Int]]) async throws -> EmbeddingResponse {
         guard let url = URL(string: configuration.endpoint) else {
             throw ChatGPTServiceError.endpointIncorrect
@@ -112,7 +123,17 @@ public struct EmbeddingService {
                 .otherError(String(data: result, encoding: .utf8) ?? "Unknown Error")
         }
 
-        return try JSONDecoder().decode(EmbeddingResponse.self, from: result)
+        let embeddingResponse = try JSONDecoder().decode(EmbeddingResponse.self, from: result)
+        #if DEBUG
+        Logger.service.info("""
+        Embedding usage
+        - number of strings: \(tokens.count)
+        - prompt tokens: \(embeddingResponse.usage.prompt_tokens)
+        - total tokens: \(embeddingResponse.usage.total_tokens)
+
+        """)
+        #endif
+        return embeddingResponse
     }
 }
 

Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemory.swift

@@ -1,11 +1,12 @@
 import Foundation
+import Logger
 import Preferences
 import TokenEncoder
 
 /// A memory that automatically manages the history according to max tokens and max message count.
 public actor AutoManagedChatGPTMemory: ChatGPTMemory {
     public private(set) var messages: [ChatMessage] = []
-    public private(set) var remainingTokens: Int? = nil
+    public private(set) var remainingTokens: Int?
 
     public var systemPrompt: ChatMessage
     public var history: [ChatMessage] = [] {
@@ -44,7 +45,7 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
             await setOnHistoryChangeBlock(onChange)
         }
     }
-    
+
     public func refresh() async {
         messages = generateSendingHistory()
         remainingTokens = generateRemainingTokens()
@@ -74,7 +75,8 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
             }
             partial += count
         }
-        var allTokensCount = functionTokenCount + 3 // every reply is primed with <|start|>assistant<|message|>
+        var allTokensCount = functionTokenCount +
+            3 // every reply is primed with <|start|>assistant<|message|>
         allTokensCount += systemPrompt.isEmpty ? 0 : systemMessageTokenCount
 
         for (index, message) in history.enumerated().reversed() {
@@ -93,6 +95,17 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
         if !systemPrompt.isEmpty {
             all.append(systemPrompt)
         }
+
+        #if DEBUG
+        Logger.service.info("""
+        Sending tokens count
+        - system prompt: \(systemMessageTokenCount)
+        - functions: \(functionTokenCount)
+        - total: \(allTokensCount)
+        
+        """)
+        #endif
+
         return all.reversed()
     }