Update chat context collector to handle multiple scopes and control functions

Author: intitni

Core/Sources/ChatContextCollector/ActiveDocumentChatContextCollector.swift

@@ -7,13 +7,17 @@ import XcodeInspector
 public struct ActiveDocumentChatContextCollector: ChatContextCollector {
     public init() {}
 
-    public func generateSystemPrompt(history: [ChatMessage], content prompt: String) -> String {
+    public func generateContext(
+        history: [ChatMessage],
+        scopes: Set<String>,
+        content: String
+    ) -> ChatContext? {
         let content = getEditorInformation()
         let relativePath = content.documentURL.path
             .replacingOccurrences(of: content.projectURL.path, with: "")
         let selectionRange = content.editorContent?.selections.first ?? .outOfScope
         let editorContent = {
-            if prompt.hasPrefix("@file") {
+            if scopes.contains("file") {
                 return """
                 File Content:```\(content.language.rawValue)
                 \(content.editorContent?.content ?? "")
@@ -53,7 +57,7 @@ public struct ActiveDocumentChatContextCollector: ChatContextCollector {
                 """
             }
 
-            if prompt.hasPrefix("@selection") {
+            if scopes.contains("selection") {
                 return """
                 Selected Code \
                 (start from line \(selectionRange.start.line)):```\(content.language.rawValue)
@@ -72,27 +76,30 @@ public struct ActiveDocumentChatContextCollector: ChatContextCollector {
             """
         }()
 
-        return """
-        Active Document Context:###
-        Document Relative Path: \(relativePath)
-        Selection Range Start: \
-        Line \(selectionRange.start.line) \
-        Character \(selectionRange.start.character)
-        Selection Range End: \
-        Line \(selectionRange.end.line) \
-        Character \(selectionRange.end.character)
-        Cursor Position: \
-        Line \(selectionRange.end.line) \
-        Character \(selectionRange.end.character)
-        \(editorContent)
-        Line Annotations:
-        \(
-            content.editorContent?.lineAnnotations
-                .map { "  - \($0)" }
-                .joined(separator: "\n") ?? "N/A"
+        return .init(
+            systemPrompt: """
+            Active Document Context:###
+            Document Relative Path: \(relativePath)
+            Selection Range Start: \
+            Line \(selectionRange.start.line) \
+            Character \(selectionRange.start.character)
+            Selection Range End: \
+            Line \(selectionRange.end.line) \
+            Character \(selectionRange.end.character)
+            Cursor Position: \
+            Line \(selectionRange.end.line) \
+            Character \(selectionRange.end.character)
+            \(editorContent)
+            Line Annotations:
+            \(
+                content.editorContent?.lineAnnotations
+                    .map { "  - \($0)" }
+                    .joined(separator: "\n") ?? "N/A"
+            )
+            ###
+            """,
+            functions: []
         )
-        ###
-        """
     }
 }
 

Core/Sources/ChatContextCollector/ChatContextCollector.swift

@@ -1,7 +1,20 @@
 import Foundation
 import OpenAIService
 
+public struct ChatContext {
+    public var systemPrompt: String
+    public var functions: [any ChatGPTFunction]
+    public init(systemPrompt: String, functions: [any ChatGPTFunction]) {
+        self.systemPrompt = systemPrompt
+        self.functions = functions
+    }
+}
+
 public protocol ChatContextCollector {
-    func generateSystemPrompt(history: [ChatMessage], content: String) -> String
+    func generateContext(
+        history: [ChatMessage],
+        scopes: Set<String>,
+        content: String
+    ) -> ChatContext?
 }
 

Core/Sources/ChatService/DynamicContextController.swift

@@ -1,32 +1,83 @@
 import ChatContextCollector
 import Foundation
 import OpenAIService
+import Parsing
 import Preferences
 import XcodeInspector
 
 final class DynamicContextController {
     let contextCollectors: [ChatContextCollector]
     let memory: AutoManagedChatGPTMemory
+    let functionProvider: ChatFunctionProvider
 
-    init(memory: AutoManagedChatGPTMemory, contextCollectors: ChatContextCollector...) {
+    convenience init(
+        memory: AutoManagedChatGPTMemory,
+        functionProvider: ChatFunctionProvider,
+        contextCollectors: ChatContextCollector...
+    ) {
+        self.init(
+            memory: memory,
+            functionProvider: functionProvider,
+            contextCollectors: contextCollectors
+        )
+    }
+
+    init(
+        memory: AutoManagedChatGPTMemory,
+        functionProvider: ChatFunctionProvider,
+        contextCollectors: [ChatContextCollector]
+    ) {
         self.memory = memory
+        self.functionProvider = functionProvider
         self.contextCollectors = contextCollectors
     }
 
     func updatePromptToMatchContent(systemPrompt: String, content: String) async throws {
+        var content = content
+        let scopes = Self.parseScopes(&content)
+        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)
+        }
         let contextualSystemPrompt = """
         \(language.isEmpty ? "" : "You must always reply in \(language)")
         \(systemPrompt)
 
-        \(
-            contextCollectors
-                .map { $0.generateSystemPrompt(history: oldMessages, content: content) }
-                .joined(separator: "\n")
-        )
+        \(contexts.map(\.systemPrompt).filter { !$0.isEmpty }.joined(separator: "\n\n"))
         """
         await memory.mutateSystemPrompt(contextualSystemPrompt)
+        functionProvider.append(functions: contexts.flatMap(\.functions))
+    }
+}
+
+extension DynamicContextController {
+    static func parseScopes(_ prompt: inout String) -> Set<String> {
+        guard !prompt.isEmpty else { return [] }
+        do {
+            let parser = Parse {
+                "@"
+                Many {
+                    Prefix { $0.isLetter }
+                } separator: {
+                    "+"
+                } terminator: {
+                    " "
+                }
+                Skip {
+                    Many {
+                        " "
+                    }
+                }
+                Rest()
+            }
+            let (scopes, rest) = try parser.parse(prompt)
+            prompt = String(rest)
+            return Set(scopes.map(String.init))
+        } catch {
+            return []
+        }
     }
 }