Update system prompt to be contextual

Author: intitni

Core/Sources/ChatService/ChatService.swift

@@ -3,10 +3,22 @@ import Combine
 import Foundation
 import OpenAIService
 
+let defaultSystemPrompt = """
+You are an AI programming assistant.
+You reply should be concise, clear, informative and logical.
+You MUST reply in the format of markdown.
+You MUST embed every code you provide in a markdown code block.
+You MUST add the programming language name at the start of the markdown code block.
+If you are asked to help perform a task, think step-by-step.
+"""
+
 public final class ChatService: ObservableObject {
     public let chatGPTService: any ChatGPTServiceType
     let pluginController: ChatPluginController
+    let contextController: DynamicContextController
     var cancellable = Set<AnyCancellable>()
+    var systemPrompt = defaultSystemPrompt
+    var extraSystemPrompt = ""
 
     public init<T: ChatGPTServiceType>(chatGPTService: T) {
         self.chatGPTService = chatGPTService
@@ -16,6 +28,7 @@ public final class ChatService: ObservableObject {
             TerminalChatPlugin.self,
             AITerminalChatPlugin.self
         )
+        contextController = DynamicContextController(chatGPTService: chatGPTService)
 
         chatGPTService.objectWillChange.sink { [weak self] _ in
             self?.objectWillChange.send()
@@ -25,7 +38,11 @@ public final class ChatService: ObservableObject {
     public func send(content: String) async throws {
         let handledInPlugin = try await pluginController.handleContent(content)
         if handledInPlugin { return }
-        
+        try await contextController.updatePromptToMatchContent(systemPrompt: """
+        \(systemPrompt)
+        \(extraSystemPrompt)
+        """)
+
         _ = try await chatGPTService.send(content: content, summary: nil)
     }
 
@@ -51,11 +68,17 @@ public final class ChatService: ObservableObject {
         }
     }
 
-    public func mutateSystemPrompt(_ newPrompt: String) async {
-        await chatGPTService.mutateSystemPrompt(newPrompt)
+    /// Setting it to `nil` to reset the system prompt
+    public func mutateSystemPrompt(_ newPrompt: String?) {
+        systemPrompt = newPrompt ?? defaultSystemPrompt
     }
 
+    public func mutateExtraSystemPrompt(_ newPrompt: String) {
+        extraSystemPrompt = newPrompt
+    }
+
     public func mutateHistory(_ mutator: @escaping (inout [ChatMessage]) -> Void) async {
         await chatGPTService.mutateHistory(mutator)
     }
 }
+

Core/Sources/ChatService/DynamicContextController.swift

@@ -1,24 +1,89 @@
 import Foundation
 import OpenAIService
+import SuggestionModel
 import XcodeInspector
+import Preferences
 
 final class DynamicContextController {
     let chatGPTService: any ChatGPTServiceType
-    
+
     init(chatGPTService: any ChatGPTServiceType) {
         self.chatGPTService = chatGPTService
     }
-    
+
     func updatePromptToMatchContent(systemPrompt: String) async throws {
-        
+        let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
+        let content = getEditorInformation()
+        let relativePath = content.documentURL.path
+            .replacingOccurrences(of: content.projectURL.path, with: "")
+        let selectionRange = content.editorContent?.selections.first ?? .outOfScope
+        let contextualSystemPrompt = """
+        \(language.isEmpty ? "" : "You must always reply in \(language)")
+        \(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: Same as selection range end
+        Selected Code (start from line \(selectionRange.end.line)):```\(content.language.rawValue)
+        \(content.selectedContent)
+        ```
+
+        Line Annotations:
+        \(content.editorContent?.lineAnnotations.map { "- \($0)" }.joined(separator: "\n") ?? "N/A")
+        ###
+        """
+        print(contextualSystemPrompt)
+        await chatGPTService.mutateSystemPrompt(contextualSystemPrompt)
     }
 }
 
 extension DynamicContextController {
-    func getEditorInformation() -> Any? {
-        guard let editor = XcodeInspector.shared.focusedEditor else { return nil }
-        let content = editor.content
-        
-        return nil
+    struct Information {
+        let editorContent: SourceEditor.Content?
+        let selectedContent: String
+        let documentURL: URL
+        let projectURL: URL
+        let language: CodeLanguage
+    }
+
+    func getEditorInformation() -> Information {
+        let editorContent = XcodeInspector.shared.focusedEditor?.content
+        let documentURL = XcodeInspector.shared.activeDocumentURL
+        let projectURL = XcodeInspector.shared.activeProjectURL
+        let language = languageIdentifierFromFileURL(documentURL)
+
+        if let editorContent, let range = editorContent.selections.first {
+            let startIndex = min(
+                max(0, range.start.line),
+                editorContent.lines.endIndex - 1
+            )
+            let endIndex = min(
+                max(startIndex, range.end.line),
+                editorContent.lines.endIndex - 1
+            )
+            let selectedContent = editorContent.lines[startIndex...endIndex]
+            return .init(
+                editorContent: editorContent,
+                selectedContent: selectedContent.joined(),
+                documentURL: documentURL,
+                projectURL: projectURL,
+                language: language
+            )
+        }
+
+        return .init(
+            editorContent: editorContent,
+            selectedContent: "",
+            documentURL: documentURL,
+            projectURL: projectURL,
+            language: language
+        )
     }
 }
+

Core/Sources/Service/GUI/WidgetDataSource.swift

@@ -42,13 +42,7 @@ final class WidgetDataSource {
     func createChatIfNeeded(for url: URL) -> ChatService {
         let build = {
             let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
-            let systemPrompt = """
-            \(language.isEmpty ? "" : "You must always reply in \(language)")
-            You are a senior programmer, you will answer my questions concisely. If you are replying with code, embed the code in a code block in markdown.
-            
-            You don't have any code in advance, ask me to provide it when needed.
-            """
-            let service = ChatService(chatGPTService: ChatGPTService(systemPrompt: systemPrompt))
+            let service = ChatService(chatGPTService: ChatGPTService())
             let provider = ChatProvider(
                 service: service,
                 fileURL: url,

Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift

@@ -397,40 +397,17 @@ extension WindowBaseCommandHandler {
         defer { presenter.markAsProcessing(false) }
 
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
-        let codeLanguage = languageIdentifierFromFileURL(fileURL)
 
         let code = {
             guard let selection = editor.selections.last,
                   selection.start != selection.end else { return "" }
             return editor.selectedCode(in: selection)
         }()
 
-        var systemPrompt = specifiedSystemPrompt ?? {
-            if code.isEmpty {
-                return """
-                \(language.isEmpty ? "" : "You must always reply in \(language)")
-                You are a senior programmer, you will answer my questions concisely. If you are replying with code, embed the code in a code block in markdown.
-                
-                You don't have any code in advance, ask me to provide it when needed.
-                """
-            }
-            return """
-            \(language.isEmpty ? "" : "You must always reply in \(language)")
-            You are a senior programmer, you will answer my questions concisely about the code below, or modify it according to my requests. When you receive a modification request, reply with the modified code in a code block.
-            ```\(codeLanguage.rawValue)
-            \(code)
-            ```
-            """
-        }()
-
-        if let extraSystemPrompt {
-            systemPrompt += "\n\(extraSystemPrompt)"
-        }
-
         let chat = WidgetDataSource.shared.createChatIfNeeded(for: fileURL)
-
-        await chat.mutateSystemPrompt(systemPrompt)
+        
+        chat.mutateSystemPrompt(specifiedSystemPrompt)
+        chat.mutateExtraSystemPrompt(extraSystemPrompt ?? "")
 
         Task {
             let customCommandPrefix = {
@@ -484,20 +461,10 @@ extension WindowBaseCommandHandler {
         let focusedElementURI = try await Environment.fetchFocusedElementURI()
         let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
 
-        var systemPrompt = specifiedSystemPrompt ?? """
-        \(language.isEmpty ? "" : "You must always reply in \(language)")
-        You are a senior programmer, you will answer my questions concisely. If you are replying with code, embed the code in a code block in markdown.
-        
-        You don't have any code in advance, ask me to provide it when needed.
-        """
-
-        if let extraSystemPrompt {
-            systemPrompt += "\n\(extraSystemPrompt)"
-        }
-
         let chat = WidgetDataSource.shared.createChatIfNeeded(for: focusedElementURI)
 
-        await chat.mutateSystemPrompt(systemPrompt)
+        chat.mutateSystemPrompt(specifiedSystemPrompt)
+        chat.mutateExtraSystemPrompt(extraSystemPrompt ?? "")
 
         Task {
             let customCommandPrefix = {

ExtensionService/ServiceDelegate.swift

@@ -17,3 +17,4 @@ class ServiceDelegate: NSObject, NSXPCListenerDelegate {
         return true
     }
 }
+