Merge branch 'feature/prompt-to-code-with-line-annotation' into develop

intitni · intitni · commit 4fddb24c9f0e · 2023-09-06T10:02:09.000+08:00
diff --git a/Core/Sources/PromptToCodeService/OpenAIPromptToCodeService.swift b/Core/Sources/PromptToCodeService/OpenAIPromptToCodeService.swift
@@ -2,10 +2,11 @@ import Foundation
 import OpenAIService
 import Preferences
 import SuggestionModel
+import XcodeInspector
 
 public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
     var service: (any ChatGPTServiceType)?
-    
+
     public init() {}
 
     public func stopResponding() {
@@ -14,13 +15,9 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
 
     public func modifyCode(
         code: String,
-        language: CodeLanguage,
-        indentSize: Int,
-        usesTabsForIndentation: Bool,
         requirement: String,
-        projectRootURL: URL,
-        fileURL: URL,
-        allCode: String,
+        source: PromptToCodeSource,
+        isDetached: Bool,
         extraSystemPrompt: String?,
         generateDescriptionRequirement: Bool?
     ) async throws -> AsyncThrowingStream<(code: String, description: String), Error> {
@@ -31,9 +28,25 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
             {
                 return ""
             }
-            return userPreferredLanguage.isEmpty ? "" : "in \(userPreferredLanguage)"
+            return userPreferredLanguage.isEmpty ? "" : " in \(userPreferredLanguage)"
         }()
 
+        let editor: EditorInformation = XcodeInspector.shared.focusedEditorContent ?? .init(
+            editorContent: .init(
+                content: source.allCode,
+                lines: [],
+                selections: [source.range],
+                cursorPosition: .outOfScope,
+                lineAnnotations: []
+            ),
+            selectedContent: code,
+            selectedLines: [],
+            documentURL: source.documentURL,
+            projectURL: source.projectRootURL,
+            relativePath: "",
+            language: source.language
+        )
+
         let rule: String = {
             func generateDescription(index: Int) -> String {
                 let generateDescription = generateDescriptionRequirement ?? UserDefaults.shared
@@ -46,7 +59,7 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
                     """
                     : "\(index). Reply with the result."
             }
-            switch language {
+            switch editor.language {
             case .builtIn(.markdown), .plaintext:
                 if code.isEmpty {
                     return """
@@ -82,20 +95,20 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
         }()
 
         let systemPrompt = {
-            switch language {
+            switch editor.language {
             case .builtIn(.markdown), .plaintext:
                 if code.isEmpty {
                     return """
-                    You are good at writing in \(language.rawValue).
-                    The active file is: \(fileURL.lastPathComponent).
+                    You are good at writing in \(editor.language.rawValue).
+                    The active file is: \(editor.documentURL.lastPathComponent).
                     \(extraSystemPrompt ?? "")
 
                     \(rule)
                     """
                 } else {
                     return """
-                    You are good at writing in \(language.rawValue).
-                    The active file is: \(fileURL.lastPathComponent).
+                    You are good at writing in \(editor.language.rawValue).
+                    The active file is: \(editor.documentURL.lastPathComponent).
                     \(extraSystemPrompt ?? "")
 
                     \(rule)
@@ -104,16 +117,16 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
             default:
                 if code.isEmpty {
                     return """
-                    You are a senior programer in writing in \(language.rawValue).
-                    The active file is: \(fileURL.lastPathComponent).
+                    You are a senior programer in writing in \(editor.language.rawValue).
+                    The active file is: \(editor.documentURL.lastPathComponent).
                     \(extraSystemPrompt ?? "")
 
                     \(rule)
                     """
                 } else {
                     return """
-                    You are a senior programer in writing in \(language.rawValue).
-                    The active file is: \(fileURL.lastPathComponent).
+                    You are a senior programer in writing in \(editor.language.rawValue).
+                    The active file is: \(editor.documentURL.lastPathComponent).
                     \(extraSystemPrompt ?? "")
 
                     \(rule)
@@ -122,31 +135,44 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
             }
         }()
 
+        let annotations = isDetached
+            ? []
+            : extractAnnotations(editorInformation: editor, source: source)
+
         let firstMessage: String? = {
             if code.isEmpty { return nil }
-            switch language {
+            switch editor.language {
             case .builtIn(.markdown), .plaintext:
                 return """
                 ```
                 \(code)
                 ```
+
+                line annotations found:
+                \(annotations.map { "- \($0)" }.joined(separator: "\n"))
                 """
             default:
                 return """
                 ```
                 \(code)
                 ```
+
+                line annotations found:
+                \(annotations.map { "- \($0)" }.joined(separator: "\n"))
                 """
             }
         }()
 
+        let indentation = getCommonLeadingSpaceCount(code)
+
         let secondMessage = """
-        Requirements:###
-        \(requirement)
-        ###
+        I will update the code you just provided.
+        It looks like every line has an indentation of \(indentation) spaces, I will keep that.
+
+        What is your requirement?
         """
 
-        let configuration =  UserPreferenceChatGPTConfiguration()
+        let configuration = UserPreferenceChatGPTConfiguration()
             .overriding(.init(temperature: 0))
         let memory = AutoManagedChatGPTMemory(
             systemPrompt: systemPrompt,
@@ -161,9 +187,10 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
         if let firstMessage {
             await memory.mutateHistory { history in
                 history.append(.init(role: .user, content: firstMessage))
+                history.append(.init(role: .assistant, content: secondMessage))
             }
         }
-        let stream = try await chatGPTService.send(content: secondMessage)
+        let stream = try await chatGPTService.send(content: requirement)
         return .init { continuation in
             Task {
                 var content = ""
@@ -231,5 +258,33 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
 
         return (code, description)
     }
+
+    func getCommonLeadingSpaceCount(_ code: String) -> Int {
+        let lines = code.split(separator: "\n")
+        guard !lines.isEmpty else { return 0 }
+        var commonCount = Int.max
+        for line in lines {
+            let count = line.prefix(while: { $0 == " " }).count
+            commonCount = min(commonCount, count)
+            if commonCount == 0 { break }
+        }
+        return commonCount
+    }
+
+    func extractAnnotations(
+        editorInformation: EditorInformation,
+        source: PromptToCodeSource
+    ) -> [String] {
+        guard let annotations = editorInformation.editorContent?.lineAnnotations else { return [] }
+        return annotations
+            .lazy
+            .filter { annotation in
+                annotation.line >= source.range.start.line + 1
+                    && annotation.line <= source.range.end.line + 1
+            }.map { annotation in
+                let relativeLine = annotation.line - source.range.start.line
+                return "line \(relativeLine): \(annotation.type) \(annotation.message)"
+            }
+    }
 }
 
diff --git a/Core/Sources/PromptToCodeService/PreviewPromptToCodeService.swift b/Core/Sources/PromptToCodeService/PreviewPromptToCodeService.swift
@@ -6,13 +6,9 @@ public final class PreviewPromptToCodeService: PromptToCodeServiceType {
 
     public func modifyCode(
         code: String,
-        language: CodeLanguage,
-        indentSize: Int,
-        usesTabsForIndentation: Bool,
         requirement: String,
-        projectRootURL: URL,
-        fileURL: URL,
-        allCode: String,
+        source: PromptToCodeSource,
+        isDetached: Bool,
         extraSystemPrompt: String?,
         generateDescriptionRequirement: Bool?
     ) async throws -> AsyncThrowingStream<(code: String, description: String), Error> {
diff --git a/Core/Sources/PromptToCodeService/PromptToCodeServiceType.swift b/Core/Sources/PromptToCodeService/PromptToCodeServiceType.swift
@@ -5,20 +5,38 @@ import SuggestionModel
 public protocol PromptToCodeServiceType {
     func modifyCode(
         code: String,
-        language: CodeLanguage,
-        indentSize: Int,
-        usesTabsForIndentation: Bool,
         requirement: String,
-        projectRootURL: URL,
-        fileURL: URL,
-        allCode: String,
+        source: PromptToCodeSource,
+        isDetached: Bool,
         extraSystemPrompt: String?,
         generateDescriptionRequirement: Bool?
     ) async throws -> AsyncThrowingStream<(code: String, description: String), Error>
 
     func stopResponding()
 }
 
+public struct PromptToCodeSource {
+    public var language: CodeLanguage
+    public var documentURL: URL
+    public var projectRootURL: URL
+    public var allCode: String
+    public var range: CursorRange
+
+    public init(
+        language: CodeLanguage,
+        documentURL: URL,
+        projectRootURL: URL,
+        allCode: String,
+        range: CursorRange
+    ) {
+        self.language = language
+        self.documentURL = documentURL
+        self.projectRootURL = projectRootURL
+        self.allCode = allCode
+        self.range = range
+    }
+}
+
 public struct PromptToCodeServiceDependencyKey: DependencyKey {
     public static let liveValue: PromptToCodeServiceType = PreviewPromptToCodeService()
     public static let previewValue: PromptToCodeServiceType = PreviewPromptToCodeService()
diff --git a/Core/Sources/Service/GUI/ChatTabFactory.swift b/Core/Sources/Service/GUI/ChatTabFactory.swift
@@ -87,13 +87,15 @@ enum ChatTabFactory {
 
                     let result = try await service.modifyCode(
                         code: prompt,
-                        language: .plaintext,
-                        indentSize: 4,
-                        usesTabsForIndentation: true,
                         requirement: instruction ?? "Modify content.",
-                        projectRootURL: .init(fileURLWithPath: "/"),
-                        fileURL: .init(fileURLWithPath: "/"),
-                        allCode: prompt,
+                        source: .init(
+                            language: .plaintext,
+                            documentURL: .init(fileURLWithPath: "/"),
+                            projectRootURL: .init(fileURLWithPath: "/"),
+                            allCode: prompt,
+                            range: .outOfScope
+                        ),
+                        isDetached: true,
                         extraSystemPrompt: extraSystemPrompt,
                         generateDescriptionRequirement: false
                     )
diff --git a/Core/Sources/SuggestionWidget/FeatureReducers/PromptToCode.swift b/Core/Sources/SuggestionWidget/FeatureReducers/PromptToCode.swift
@@ -63,7 +63,7 @@ public struct PromptToCode: ReducerProtocol {
         @BindingState public var prompt: String
         @BindingState public var isContinuous: Bool
         @BindingState public var isAttachedToSelectionRange: Bool
-        
+
         public var filename: String { documentURL.lastPathComponent }
         public var canRevert: Bool { history != .empty }
 
@@ -156,13 +156,15 @@ public struct PromptToCode: ReducerProtocol {
                     do {
                         let stream = try await promptToCodeService.modifyCode(
                             code: copiedState.code,
-                            language: copiedState.language,
-                            indentSize: copiedState.indentSize,
-                            usesTabsForIndentation: copiedState.usesTabsForIndentation,
                             requirement: copiedState.prompt,
-                            projectRootURL: copiedState.projectRootURL,
-                            fileURL: copiedState.documentURL,
-                            allCode: copiedState.allCode,
+                            source: .init(
+                                language: copiedState.language,
+                                documentURL: copiedState.documentURL,
+                                projectRootURL: copiedState.projectRootURL,
+                                allCode: copiedState.allCode,
+                                range: copiedState.selectionRange ?? .outOfScope
+                            ),
+                            isDetached: !copiedState.isAttachedToSelectionRange,
                             extraSystemPrompt: copiedState.extraSystemPrompt,
                             generateDescriptionRequirement: copiedState
                                 .generateDescriptionRequirement
