Pass configuration to DynamicContextController

Author: intitni

Core/Sources/ChatContextCollector/ChatContextCollector.swift

@@ -14,7 +14,8 @@ public protocol ChatContextCollector {
     func generateContext(
         history: [ChatMessage],
         scopes: Set<String>,
-        content: String
+        content: String,
+        configuration: ChatGPTConfiguration
     ) -> ChatContext?
 }
 

Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/ActiveDocumentChatContextCollector.swift

@@ -15,7 +15,8 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
     public func generateContext(
         history: [ChatMessage],
         scopes: Set<String>,
-        content: String
+        content: String,
+        configuration: ChatGPTConfiguration
     ) -> ChatContext? {
         guard let info = getEditorInformation() else { return nil }
         let context = getActiveDocumentContext(info)

Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/LegacyActiveDocumentChatContextCollector.swift

@@ -11,7 +11,8 @@ public struct LegacyActiveDocumentChatContextCollector: ChatContextCollector {
     public func generateContext(
         history: [ChatMessage],
         scopes: Set<String>,
-        content: String
+        content: String,
+        configuration: ChatGPTConfiguration
     ) -> ChatContext? {
         guard let content = getEditorInformation() else { return nil }
         let relativePath = content.relativePath

Core/Sources/ChatContextCollectors/SystemInfoChatContextCollector/SystemInfoChatContextCollector.swift

@@ -14,7 +14,8 @@ public final class SystemInfoChatContextCollector: ChatContextCollector {
     public func generateContext(
         history: [ChatMessage],
         scopes: Set<String>,
-        content: String
+        content: String,
+        configuration: ChatGPTConfiguration
     ) -> ChatContext? {
         return .init(
             systemPrompt: """

Core/Sources/ChatContextCollectors/WebChatContextCollector/SearchFunction.swift

@@ -28,6 +28,8 @@ struct SearchFunction: ChatGPTFunction {
             }.joined(separator: "\n")
         }
     }
+    
+    let maxTokens: Int
 
     var reportProgress: (String) async -> Void = { _ in }
 
@@ -73,11 +75,9 @@ struct SearchFunction: ChatGPTFunction {
                 searchURL: UserDefaults.shared.value(for: \.bingSearchEndpoint)
             )
 
-            #warning("request chat service to pass in the token length")
-            
             let result = try await bingSearch.search(
                 query: arguments.query,
-                numberOfResult: UserDefaults.shared.value(for: \.chatGPTMaxToken) > 5000 ? 5 : 3,
+                numberOfResult: maxTokens > 5000 ? 5 : 3,
                 freshness: arguments.freshness
             )
 

Core/Sources/ChatContextCollectors/WebChatContextCollector/WebChatContextCollector.swift

@@ -10,12 +10,13 @@ public final class WebChatContextCollector: ChatContextCollector {
     public func generateContext(
         history: [ChatMessage],
         scopes: Set<String>,
-        content: String
+        content: String,
+        configuration: ChatGPTConfiguration
     ) -> ChatContext? {
         guard scopes.contains("web") || scopes.contains("w") else { return nil }
         let links = Self.detectLinks(from: history) + Self.detectLinks(from: content)
         let functions: [(any ChatGPTFunction)?] = [
-            SearchFunction(),
+            SearchFunction(maxTokens: configuration.maxTokens),
             // allow this function only when there is a link in the memory.
             links.isEmpty ? nil : QueryWebsiteFunction(),
         ]

Core/Sources/ChatService/ContextAwareAutoManagedChatGPTMemory.swift

@@ -35,6 +35,7 @@ public final class ContextAwareAutoManagedChatGPTMemory: ChatGPTMemory {
         contextController = DynamicContextController(
             memory: memory,
             functionProvider: functionProvider,
+            configuration: configuration,
             contextCollectors: allContextCollectors
         )
         self.functionProvider = functionProvider

Core/Sources/ChatService/DynamicContextController.swift

@@ -9,40 +9,50 @@ final class DynamicContextController {
     let contextCollectors: [ChatContextCollector]
     let memory: AutoManagedChatGPTMemory
     let functionProvider: ChatFunctionProvider
+    let configuration: ChatGPTConfiguration
     var defaultScopes = [] as Set<String>
 
     convenience init(
         memory: AutoManagedChatGPTMemory,
         functionProvider: ChatFunctionProvider,
+        configuration: ChatGPTConfiguration,
         contextCollectors: ChatContextCollector...
     ) {
         self.init(
             memory: memory,
             functionProvider: functionProvider,
+            configuration: configuration,
             contextCollectors: contextCollectors
         )
     }
 
     init(
         memory: AutoManagedChatGPTMemory,
         functionProvider: ChatFunctionProvider,
+        configuration: ChatGPTConfiguration,
         contextCollectors: [ChatContextCollector]
     ) {
         self.memory = memory
         self.functionProvider = functionProvider
+        self.configuration = configuration
         self.contextCollectors = contextCollectors
     }
 
     func updatePromptToMatchContent(systemPrompt: String, content: String) async throws {
         var content = content
         var scopes = Self.parseScopes(&content)
         scopes.formUnion(defaultScopes)
-        
+
         functionProvider.removeAll()
         let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
         let oldMessages = await memory.history
         let contexts = contextCollectors.compactMap {
-            $0.generateContext(history: oldMessages, scopes: scopes, content: content)
+            $0.generateContext(
+                history: oldMessages,
+                scopes: scopes,
+                content: content,
+                configuration: configuration
+            )
         }
         let contextualSystemPrompt = """
         \(language.isEmpty ? "" : "You must always reply in \(language)")