Count functions tokens

Author: intitni

Core/Sources/ChatPlugin/AskChatGPT.swift

@@ -9,7 +9,11 @@ public func askChatGPT(
 ) async throws -> String? {
     let configuration = UserPreferenceChatGPTConfiguration()
         .overriding(.init(temperature: temperature))
-    let memory = AutoManagedChatGPTMemory(systemPrompt: systemPrompt, configuration: configuration)
+    let memory = AutoManagedChatGPTMemory(
+        systemPrompt: systemPrompt,
+        configuration: configuration,
+        functionProvider: NoChatGPTFunctionProvider()
+    )
     let service = ChatGPTService(
         memory: memory,
         configuration: configuration

Core/Sources/ChatPlugin/CallAIFunction.swift

@@ -21,7 +21,8 @@ func callAIFunction(
     let service = ChatGPTService(
         memory: AutoManagedChatGPTMemory(
             systemPrompt: "You are now the following python function: ```# \(description)\n\(function)```\n\nOnly respond with your `return` value.",
-            configuration: configuration
+            configuration: configuration,
+            functionProvider: NoChatGPTFunctionProvider()
         ),
         configuration: configuration
     )

Core/Sources/PromptToCodeService/OpenAIPromptToCodeAPI.swift

@@ -151,7 +151,8 @@ final class OpenAIPromptToCodeAPI: PromptToCodeAPI {
             .overriding(.init(temperature: 0))
         let memory = AutoManagedChatGPTMemory(
             systemPrompt: systemPrompt,
-            configuration: configuration
+            configuration: configuration,
+            functionProvider: NoChatGPTFunctionProvider()
         )
         let chatGPTService = ChatGPTService(
             memory: memory,

Tool/Sources/LangChain/ChatModel/OpenAIChat.swift

@@ -22,7 +22,11 @@ public struct OpenAIChat: ChatModel {
             temperature: temperature,
             stop: stops
         ))
-        let memory = AutoManagedChatGPTMemory(systemPrompt: "", configuration: configuration)
+        let memory = AutoManagedChatGPTMemory(
+            systemPrompt: "",
+            configuration: configuration,
+            functionProvider: NoChatGPTFunctionProvider()
+        )
         let service = ChatGPTService(memory: memory, configuration: configuration)
         for message in prompt {
             let role: OpenAIService.ChatMessage.Role = {

Tool/Sources/OpenAIService/ChatGPTService.swift

@@ -65,7 +65,8 @@ public class ChatGPTService: ChatGPTServiceType {
     public init(
         memory: ChatGPTMemory = AutoManagedChatGPTMemory(
             systemPrompt: "",
-            configuration: UserPreferenceChatGPTConfiguration()
+            configuration: UserPreferenceChatGPTConfiguration(),
+            functionProvider: NoChatGPTFunctionProvider()
         ),
         configuration: ChatGPTConfiguration = UserPreferenceChatGPTConfiguration(),
         functionProvider: ChatGPTFunctionProvider = NoChatGPTFunctionProvider()

Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemory.swift

@@ -13,14 +13,20 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
     }
 
     public var configuration: ChatGPTConfiguration
+    public var functionProvider: ChatGPTFunctionProvider
 
     static let encoder: TokenEncoder = GPTEncoder()
 
     var onHistoryChange: () -> Void = {}
 
-    public init(systemPrompt: String, configuration: ChatGPTConfiguration) {
+    public init(
+        systemPrompt: String,
+        configuration: ChatGPTConfiguration,
+        functionProvider: ChatGPTFunctionProvider
+    ) {
         self.systemPrompt = .init(role: .system, content: systemPrompt)
         self.configuration = configuration
+        self.functionProvider = functionProvider
     }
 
     public func mutateHistory(_ update: (inout [ChatMessage]) -> Void) {
@@ -48,11 +54,22 @@ public actor AutoManagedChatGPTMemory: ChatGPTMemory {
             message.tokensCount = count
             return count
         }
-        
+
         var all: [ChatMessage] = []
         let systemMessageTokenCount = countToken(&systemPrompt)
-        var allTokensCount = systemPrompt.isEmpty ? 0 : systemMessageTokenCount
-        
+        let functionTokenCount = functionProvider.functions.reduce(into: 0) { partial, function in
+            var count = encoder.countToken(text: function.name)
+                + encoder.countToken(text: function.description)
+            if let data = try? JSONEncoder().encode(function.argumentSchema),
+               let string = String(data: data, encoding: .utf8)
+            {
+                count += encoder.countToken(text: string)
+            }
+            partial += count
+        }
+        var allTokensCount = functionTokenCount
+        allTokensCount += systemPrompt.isEmpty ? 0 : systemMessageTokenCount
+
         for (index, message) in history.enumerated().reversed() {
             if maxNumberOfMessages > 0, all.count >= maxNumberOfMessages { break }
             if message.isEmpty { continue }

Tool/Sources/OpenAIService/Tokenizer.swift

@@ -23,8 +23,8 @@ extension TokenEncoder {
         return total
     }
 
-    func countToken(functionJSONSchema: String) -> Int {
-        encode(text: functionJSONSchema).count
+    func countToken(text: String) -> Int {
+        encode(text: text).count
     }
 }