Merge branch 'feature/rag-base-project-scope' into develop

Author: intitni

Core/Package.swift

@@ -210,7 +210,6 @@ let package = Package(
 
                 // context collectors
                 "WebChatContextCollector",
-                "ActiveDocumentChatContextCollector",
                 "SystemInfoChatContextCollector",
 
                 .product(name: "ChatContextCollector", package: "Tool"),
@@ -354,23 +353,6 @@ let package = Package(
             ],
             path: "Sources/ChatContextCollectors/SystemInfoChatContextCollector"
         ),
-
-        .target(
-            name: "ActiveDocumentChatContextCollector",
-            dependencies: [
-                .product(name: "ChatContextCollector", package: "Tool"),
-                .product(name: "OpenAIService", package: "Tool"),
-                .product(name: "Preferences", package: "Tool"),
-                .product(name: "FocusedCodeFinder", package: "Tool"),
-                .product(name: "AppMonitoring", package: "Tool"),
-            ],
-            path: "Sources/ChatContextCollectors/ActiveDocumentChatContextCollector"
-        ),
-
-        .testTarget(
-            name: "ActiveDocumentChatContextCollectorTests",
-            dependencies: ["ActiveDocumentChatContextCollector"]
-        ),
     ]
 )
 

Core/Sources/ChatContextCollectors/SystemInfoChatContextCollector/SystemInfoChatContextCollector.swift

@@ -16,11 +16,14 @@ public final class SystemInfoChatContextCollector: ChatContextCollector {
         scopes: Set<String>,
         content: String,
         configuration: ChatGPTConfiguration
-    ) -> ChatContext? {
+    ) -> ChatContext {
         return .init(
             systemPrompt: """
-            Current Time: \(Self.dateFormatter.string(from: Date())) (You can use it to calculate time in another time zone)
-            """,
+                    Current Time: \(
+                        Self.dateFormatter.string(from: Date())
+                    ) (You can use it to calculate time in another time zone)
+                    """,
+            retrievedContent: [],
             functions: []
         )
     }

Core/Sources/ChatContextCollectors/WebChatContextCollector/WebChatContextCollector.swift

@@ -12,8 +12,8 @@ public final class WebChatContextCollector: ChatContextCollector {
         scopes: Set<String>,
         content: String,
         configuration: ChatGPTConfiguration
-    ) -> ChatContext? {
-        guard scopes.contains("web") || scopes.contains("w") else { return nil }
+    ) -> ChatContext {
+        guard scopes.contains("web") || scopes.contains("w") else { return .empty }
         let links = Self.detectLinks(from: history) + Self.detectLinks(from: content)
         let functions: [(any ChatGPTFunction)?] = [
             SearchFunction(maxTokens: configuration.maxTokens),
@@ -22,6 +22,7 @@ public final class WebChatContextCollector: ChatContextCollector {
         ]
         return .init(
             systemPrompt: "You prefer to answer questions with latest content on the internet.",
+            retrievedContent: [],
             functions: functions.compactMap { $0 }
         )
     }

Core/Sources/ChatGPTChatTab/ChatPanel.swift

@@ -58,7 +58,7 @@ struct ChatPanelMessages: View {
 
                         ChatHistory(chat: chat)
                             .listItemTint(.clear)
-                        
+
                         WithViewStore(chat, observe: \.isReceivingMessage) { viewStore in
                             if viewStore.state {
                                 Spacer(minLength: 12)
@@ -231,35 +231,35 @@ private struct Instruction: View {
         Group {
             Markdown(
                 """
-                You can use scopes to give the bot extra abilities.
+                You can use plugins to perform various tasks.
 
-                | Scope Name | Abilities |
+                | Plugin Name | Description |
                 | --- | --- |
-                | `@file` | Read the metadata of the editing file |
-                | `@code` | Read the code and metadata in the editing file |
-                | `@web` (beta) | Search on Bing or query from a web page |
-                | `@project` | Experimental. Access content of the project |
-
-                To use scopes, you can prefix a message with `@code`.
+                | `/run` | Runs a command under the project root |
+                | `/math` | Solves a math problem in natural language |
+                | `/search` | Searches on Bing and summarizes the results |
+                | `/shortcut(name)` | Runs a shortcut from the Shortcuts.app, with the previous message as input |
+                | `/shortcutInput(name)` | Runs a shortcut and uses its result as a new message |
 
-                You can use shorthand to represent a scope, such as `@c`, and enable multiple scopes with `@c+web`.
+                To use plugins, you can prefix a message with `/pluginName`.
                 """
             )
             .modifier(InstructionModifier())
 
             Markdown(
                 """
-                You can use plugins to perform various tasks.
+                You can use scopes to give the bot extra abilities.
 
-                | Plugin Name | Description |
+                | Scope Name | Abilities |
                 | --- | --- |
-                | `/run` | Runs a command under the project root |
-                | `/math` | Solves a math problem in natural language |
-                | `/search` | Searches on Bing and summarizes the results |
-                | `/shortcut(name)` | Runs a shortcut from the Shortcuts.app, with the previous message as input |
-                | `/shortcutInput(name)` | Runs a shortcut and uses its result as a new message |
+                | `@file` | Read the metadata of the editing file |
+                | `@code` | Read the code and metadata in the editing file |
+                | `@web` (beta) | Search on Bing or query from a web page |
+                | `@project` | Experimental. Access content of the project |
 
-                To use plugins, you can prefix a message with `/pluginName`.
+                To use scopes, you can prefix a message with `@code`.
+
+                You can use shorthand to represent a scope, such as `@c`, and enable multiple scopes with `@c+web`.
                 """
             )
             .modifier(InstructionModifier())
@@ -537,7 +537,7 @@ struct ChatPanelInputArea: View {
         let availableFeatures = plugins + [
             "/exit",
             "@code",
-            "@file",
+            "@project",
             "@web",
         ]
 

Core/Sources/ChatService/AllContextCollector.swift

@@ -6,7 +6,6 @@ import WebChatContextCollector
 import ProChatContextCollectors
 let allContextCollectors: [any ChatContextCollector] = [
     SystemInfoChatContextCollector(),
-    ActiveDocumentChatContextCollector(),
     WebChatContextCollector(),
     ProChatContextCollectors(),
 ]

Core/Sources/ChatService/DynamicContextController.swift

@@ -46,21 +46,42 @@ final class DynamicContextController {
         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,
-                configuration: configuration
-            )
+        let contexts = await withTaskGroup(
+            of: ChatContext.self
+        ) { [scopes, content, configuration] group in
+            for collector in contextCollectors {
+                group.addTask {
+                    await collector.generateContext(
+                        history: oldMessages,
+                        scopes: scopes,
+                        content: content,
+                        configuration: configuration
+                    )
+                }
+            }
+            var contexts = [ChatContext]()
+            for await context in group {
+                contexts.append(context)
+            }
+            return contexts
         }
+        
+        let extraSystemPrompt = contexts
+            .map(\.systemPrompt)
+            .filter { !$0.isEmpty }
+            .joined(separator: "\n")
+        
+        let contextPrompts = contexts
+            .flatMap(\.retrievedContent)
+            .filter { !$0.content.isEmpty }
+            .sorted { $0.priority > $1.priority }
+        
         let contextualSystemPrompt = """
         \(language.isEmpty ? "" : "You must always reply in \(language)")
-        \(systemPrompt)
-
-        \(contexts.map(\.systemPrompt).filter { !$0.isEmpty }.joined(separator: "\n\n"))
+        \(systemPrompt)\(extraSystemPrompt.isEmpty ? "" : "\n\(extraSystemPrompt)")
         """
         await memory.mutateSystemPrompt(contextualSystemPrompt)
+        await memory.mutateRetrievedContent(contextPrompts.map(\.content))
         functionProvider.append(functions: contexts.flatMap(\.functions))
     }
 }

Pro

@@ -15,7 +15,10 @@ let package = Package(
         .library(name: "Logger", targets: ["Logger"]),
         .library(name: "OpenAIService", targets: ["OpenAIService"]),
         .library(name: "ChatTab", targets: ["ChatTab"]),
-        .library(name: "ChatContextCollector", targets: ["ChatContextCollector"]),
+        .library(
+            name: "ChatContextCollector",
+            targets: ["ChatContextCollector", "ActiveDocumentChatContextCollector"]
+        ),
         .library(name: "Environment", targets: ["Environment"]),
         .library(name: "SuggestionModel", targets: ["SuggestionModel"]),
         .library(name: "ASTParser", targets: ["ASTParser"]),
@@ -70,7 +73,7 @@ let package = Package(
         // MARK: - Helpers
 
         .target(name: "XPCShared", dependencies: ["SuggestionModel"]),
-        
+
         .target(name: "Configs"),
 
         .target(name: "Preferences", dependencies: ["Configs", "AIModel"]),
@@ -234,14 +237,6 @@ let package = Package(
             ]
         ),
 
-        .target(
-            name: "ChatContextCollector",
-            dependencies: [
-                "SuggestionModel",
-                "OpenAIService",
-            ]
-        ),
-
         .target(name: "BingSearchService"),
 
         .target(name: "SuggestionService", dependencies: [
@@ -310,6 +305,33 @@ let package = Package(
             )]
         ),
 
+        // MARK: - Chat Context Collector
+
+        .target(
+            name: "ChatContextCollector",
+            dependencies: [
+                "SuggestionModel",
+                "OpenAIService",
+            ]
+        ),
+
+        .target(
+            name: "ActiveDocumentChatContextCollector",
+            dependencies: [
+                "ChatContextCollector",
+                "OpenAIService",
+                "Preferences",
+                "FocusedCodeFinder",
+                "XcodeInspector",
+            ],
+            path: "Sources/ChatContextCollectors/ActiveDocumentChatContextCollector"
+        ),
+
+        .testTarget(
+            name: "ActiveDocumentChatContextCollectorTests",
+            dependencies: ["ActiveDocumentChatContextCollector"]
+        ),
+
         // MARK: - Tests
 
         .testTarget(

Tool/Package.swift

@@ -2,12 +2,32 @@ import Foundation
 import OpenAIService
 
 public struct ChatContext {
+    public struct RetrievedContent {
+        public var content: String
+        public var priority: Int
+
+        public init(content: String, priority: Int) {
+            self.content = content
+            self.priority = priority
+        }
+    }
+
     public var systemPrompt: String
+    public var retrievedContent: [RetrievedContent]
     public var functions: [any ChatGPTFunction]
-    public init(systemPrompt: String, functions: [any ChatGPTFunction]) {
+    public init(
+        systemPrompt: String,
+        retrievedContent: [RetrievedContent],
+        functions: [any ChatGPTFunction]
+    ) {
         self.systemPrompt = systemPrompt
+        self.retrievedContent = retrievedContent
         self.functions = functions
     }
+
+    public static var empty: Self {
+        .init(systemPrompt: "", retrievedContent: [], functions: [])
+    }
 }
 
 public protocol ChatContextCollector {
@@ -16,6 +36,6 @@ public protocol ChatContextCollector {
         scopes: Set<String>,
         content: String,
         configuration: ChatGPTConfiguration
-    ) -> ChatContext?
+    ) async -> ChatContext
 }
 

Tool/Sources/ChatContextCollector/ChatContextCollector.swift

@@ -10,15 +10,15 @@ import XcodeInspector
 public final class ActiveDocumentChatContextCollector: ChatContextCollector {
     public init() {}
 
-    var activeDocumentContext: ActiveDocumentContext?
+    public var activeDocumentContext: ActiveDocumentContext?
 
     public func generateContext(
         history: [ChatMessage],
         scopes: Set<String>,
         content: String,
         configuration: ChatGPTConfiguration
-    ) -> ChatContext? {
-        guard let info = getEditorInformation() else { return nil }
+    ) -> ChatContext {
+        guard let info = getEditorInformation() else { return .empty }
         let context = getActiveDocumentContext(info)
         activeDocumentContext = context
 
@@ -28,10 +28,11 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
                 removedCode.focusedContext = nil
                 return .init(
                     systemPrompt: extractSystemPrompt(removedCode),
+                    retrievedContent: [],
                     functions: []
                 )
             }
-            return nil
+            return .empty
         }
 
         var functions = [any ChatGPTFunction]()
@@ -66,6 +67,7 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
 
         return .init(
             systemPrompt: extractSystemPrompt(context),
+            retrievedContent: [],
             functions: functions
         )
     }