Add command chat with selection

Author: intitni

Copilot for Xcode.xcodeproj/project.pbxproj

@@ -46,6 +46,7 @@
 		C882175A294187E100A22FD3 /* Client in Frameworks */ = {isa = PBXBuildFile; productRef = C8821759294187E100A22FD3 /* Client */; };
 		C882175C294187EF00A22FD3 /* Client in Frameworks */ = {isa = PBXBuildFile; productRef = C882175B294187EF00A22FD3 /* Client */; };
 		C8C8B60929AFA35F00034BEE /* CopilotForXcodeExtensionService.app in Embed XPCService */ = {isa = PBXBuildFile; fileRef = C861E60E2994F6070056CB02 /* CopilotForXcodeExtensionService.app */; settings = {ATTRIBUTES = (RemoveHeadersOnCopy, ); }; };
+		C8DCF00029CE11D500FDDDD7 /* ChatWithSelection.swift in Sources */ = {isa = PBXBuildFile; fileRef = C8DCEFFF29CE11D500FDDDD7 /* ChatWithSelection.swift */; };
 		C8EE079D29CC21300043B6D9 /* AccountView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C8EE079C29CC21300043B6D9 /* AccountView.swift */; };
 		C8EE079F29CC25C20043B6D9 /* OpenAIView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C8EE079E29CC25C20043B6D9 /* OpenAIView.swift */; };
 		C8EE07A129CC9ED30043B6D9 /* ExplainSelectionCommand.swift in Sources */ = {isa = PBXBuildFile; fileRef = C8EE07A029CC9ED30043B6D9 /* ExplainSelectionCommand.swift */; };
@@ -191,6 +192,7 @@
 		C87F3E61293DD004008523E8 /* Styles.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Styles.swift; sourceTree = "<group>"; };
 		C887BC832965D96000931567 /* DEVELOPMENT.md */ = {isa = PBXFileReference; lastKnownFileType = net.daringfireball.markdown; path = DEVELOPMENT.md; sourceTree = "<group>"; };
 		C8CD828229B88006008D044D /* TestPlan.xctestplan */ = {isa = PBXFileReference; lastKnownFileType = text; path = TestPlan.xctestplan; sourceTree = "<group>"; };
+		C8DCEFFF29CE11D500FDDDD7 /* ChatWithSelection.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ChatWithSelection.swift; sourceTree = "<group>"; };
 		C8EE079C29CC21300043B6D9 /* AccountView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AccountView.swift; sourceTree = "<group>"; };
 		C8EE079E29CC25C20043B6D9 /* OpenAIView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OpenAIView.swift; sourceTree = "<group>"; };
 		C8EE07A029CC9ED30043B6D9 /* ExplainSelectionCommand.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ExplainSelectionCommand.swift; sourceTree = "<group>"; };
@@ -257,6 +259,7 @@
 				C8009C022941C576007AA7E8 /* RealtimeSuggestionCommand.swift */,
 				C800DBB0294C624D00B04CAC /* PrefetchSuggestionsCommand.swift */,
 				C8EE07A029CC9ED30043B6D9 /* ExplainSelectionCommand.swift */,
+				C8DCEFFF29CE11D500FDDDD7 /* ChatWithSelection.swift */,
 				C81458972939EFDC00135263 /* Info.plist */,
 				C81458982939EFDC00135263 /* EditorExtension.entitlements */,
 			);
@@ -519,6 +522,7 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				C8DCF00029CE11D500FDDDD7 /* ChatWithSelection.swift in Sources */,
 				C81458942939EFDC00135263 /* SourceEditorExtension.swift in Sources */,
 				C8520301293C4D9000460097 /* Helpers.swift in Sources */,
 				C8009BFF2941C551007AA7E8 /* ToggleRealtimeSuggestionsCommand.swift in Sources */,

Core/Sources/Client/AsyncXPCService.swift

@@ -178,6 +178,14 @@ public struct AsyncXPCService {
             { $0.explainSelection }
         )
     }
+    
+    public func chatWithSelection(editorContent: EditorContent) async throws -> UpdatedContent? {
+        try await suggestionRequest(
+            connection,
+            editorContent,
+            { $0.chatWithSelection }
+        )
+    }
 }
 
 struct AutoFinishContinuation<T> {

Core/Sources/Service/SuggestionCommandHandler/CommentBaseCommandHandler.swift

@@ -175,6 +175,10 @@ struct CommentBaseCommandHandler: SuggestionCommandHandler {
     func explainSelection(editor: EditorContent) async throws -> UpdatedContent? {
         throw NotSupportedInCommentMode()
     }
+    
+    func chatWithSelection(editor: EditorContent) async throws -> UpdatedContent? {
+        throw NotSupportedInCommentMode()
+    }
 }
 
 // MARK: - Unsupported

Core/Sources/Service/SuggestionCommandHandler/SuggestionCommandHandler.swift

@@ -18,4 +18,6 @@ protocol SuggestionCommandHandler {
     func generateRealtimeSuggestions(editor: EditorContent) async throws -> UpdatedContent?
     @ServiceActor
     func explainSelection(editor: EditorContent) async throws -> UpdatedContent?
+    @ServiceActor
+    func chatWithSelection(editor: EditorContent) async throws -> UpdatedContent?
 }

Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift

@@ -194,15 +194,14 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
         let fileURL = try await Environment.fetchCurrentFileURL()
         let endpoint = UserDefaults.shared.value(for: \.chatGPTEndpoint)
         let model = UserDefaults.shared.value(for: \.chatGPTModel)
-        let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
+        var language = UserDefaults.shared.value(for: \.chatGPTLanguage)
+        if language.isEmpty { language = "English" }
         guard let selection = editor.selections.last else { return }
 
         let service = ChatGPTService(
             systemPrompt: """
-            You are a code explanation engine, you can only explain the code concisely, do not interpret or translate it. Reply in \(
-                language
-                    .isEmpty ? language : "English"
-            )
+            You are a code explanation engine, you can only explain the code concisely, do not interpret or translate it
+            Reply in \(language)
             """,
             apiKey: UserDefaults.shared.value(for: \.openAIAPIKey),
             endpoint: endpoint.isEmpty ? nil : endpoint,
@@ -214,15 +213,82 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
         let code = editor.selectedCode(in: selection)
         Task {
             try? await service.send(
-                content: removeContinousSpaces(from: code),
+                content: removeContinuousSpaces(from: code),
                 summary: "Explain selected code from `\(selection.start.line + 1):\(selection.start.character + 1)` to `\(selection.end.line + 1):\(selection.end.character + 1)`."
             )
         }
 
         presenter.presentChatGPTConversation(service, fileURL: fileURL)
     }
+
+    func chatWithSelection(editor: EditorContent) async throws -> UpdatedContent? {
+        Task {
+            do {
+                try await _chatWithSelection(editor: editor)
+            } catch {
+                presenter.presentError(error)
+            }
+        }
+        return nil
+    }
+
+    private func _chatWithSelection(editor: EditorContent) async throws {
+        presenter.markAsProcessing(true)
+        defer { presenter.markAsProcessing(false) }
+
+        let fileURL = try await Environment.fetchCurrentFileURL()
+        let endpoint = UserDefaults.shared.value(for: \.chatGPTEndpoint)
+        let model = UserDefaults.shared.value(for: \.chatGPTModel)
+        var language = UserDefaults.shared.value(for: \.chatGPTLanguage)
+        if language.isEmpty { language = "English" }
+
+        let code = {
+            guard let selection = editor.selections.last,
+                  selection.start != selection.end else { return "" }
+            return editor.selectedCode(in: selection)
+        }()
+
+        let prompt = {
+            if code.isEmpty {
+                return """
+                You are a senior programmer, you will answer my questions concisely in \(language)
+                """
+            }
+            return """
+            You are a senior programmer, you will answer my questions concisely in \(
+                language
+            ) about the code
+            ```
+            \(removeContinuousSpaces(from: code))
+            ```
+            """
+        }()
+
+        let service = ChatGPTService(
+            systemPrompt: prompt,
+            apiKey: UserDefaults.shared.value(for: \.openAIAPIKey),
+            endpoint: endpoint.isEmpty ? nil : endpoint,
+            model: model.isEmpty ? nil : model,
+            temperature: 1,
+            maxToken: UserDefaults.shared.value(for: \.chatGPTMaxToken)
+        )
+
+        Task {
+            if !code.isEmpty, let selection = editor.selections.last {
+                await service.mutateHistory { history in
+                    history.append(.init(
+                        role: .user,
+                        content: "",
+                        summary: "Chat about selected code from `\(selection.start.line + 1):\(selection.start.character + 1)` to `\(selection.end.line + 1):\(selection.end.character)`.\nThe code will persisted in the conversation."
+                    ))
+                }
+            }
+        }
+
+        presenter.presentChatGPTConversation(service, fileURL: fileURL)
+    }
 }
 
-func removeContinousSpaces(from string: String) -> String {
+func removeContinuousSpaces(from string: String) -> String {
     return string.replacingOccurrences(of: " +", with: " ", options: .regularExpression)
 }

Core/Sources/Service/XPCService.swift

@@ -223,6 +223,15 @@ public class XPCService: NSObject, XPCServiceProtocol {
             try await handler.explainSelection(editor: editor)
         }
     }
+    
+    public func chatWithSelection(
+        editorContent: Data,
+        withReply reply: @escaping (Data?, Error?) -> Void
+    ) {
+        replyWithUpdatedContent(editorContent: editorContent, withReply: reply) { handler, editor in
+            try await handler.chatWithSelection(editor: editor)
+        }
+    }
 
     // MARK: - Settings
 

Core/Sources/XPCShared/XPCServiceProtocol.swift

@@ -40,6 +40,10 @@ public protocol XPCServiceProtocol {
         editorContent: Data,
         withReply reply: @escaping (Data?, Error?) -> Void
     )
+    func chatWithSelection(
+        editorContent: Data,
+        withReply reply: @escaping (Data?, Error?) -> Void
+    )
 
     func toggleRealtimeSuggestion(withReply reply: @escaping (Error?) -> Void)
 

EditorExtension/ChatWithSelection.swift

@@ -0,0 +1,29 @@
+import Client
+import CopilotModel
+import Foundation
+import XcodeKit
+
+class ChatWithSelectionCommand: NSObject, XCSourceEditorCommand, CommandType {
+    var name: String { "Chat With Selection" }
+
+    func perform(
+        with invocation: XCSourceEditorCommandInvocation,
+        completionHandler: @escaping (Error?) -> Void
+    ) {
+        Task {
+            do {
+                let service = try getService()
+                if let content = try await service.chatWithSelection(
+                    editorContent: .init(invocation)
+                ) {
+                    invocation.accept(content)
+                }
+                completionHandler(nil)
+            } catch is CancellationError {
+                completionHandler(nil)
+            } catch {
+                completionHandler(error)
+            }
+        }
+    }
+}

EditorExtension/SourceEditorExtension.swift

@@ -13,6 +13,7 @@ class SourceEditorExtension: NSObject, XCSourceEditorExtension {
             ToggleRealtimeSuggestionsCommand(),
             RealtimeSuggestionsCommand(),
             PrefetchSuggestionsCommand(),
+            ChatWithSelectionCommand(),
             ExplainSelectionCommand(),
         ].map(makeCommandDefinition)
     }