Support retrieval in context aware prompt to code

Author: intitni

Core/Sources/PromptToCodeService/OpenAIPromptToCodeService.swift

@@ -33,8 +33,8 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
 
         let editor: EditorInformation = XcodeInspector.shared.focusedEditorContent ?? .init(
             editorContent: .init(
-                content: source.allCode,
-                lines: [],
+                content: source.content,
+                lines: source.lines,
                 selections: [source.range],
                 cursorPosition: .outOfScope,
                 lineAnnotations: []

Core/Sources/PromptToCodeService/PromptToCodeServiceType.swift

@@ -19,20 +19,23 @@ public struct PromptToCodeSource {
     public var language: CodeLanguage
     public var documentURL: URL
     public var projectRootURL: URL
-    public var allCode: String
+    public var content: String
+    public var lines: [String]
     public var range: CursorRange
 
     public init(
         language: CodeLanguage,
         documentURL: URL,
         projectRootURL: URL,
-        allCode: String,
+        content: String,
+        lines: [String],
         range: CursorRange
     ) {
         self.language = language
         self.documentURL = documentURL
         self.projectRootURL = projectRootURL
-        self.allCode = allCode
+        self.content = content
+        self.lines = lines
         self.range = range
     }
 }
@@ -74,7 +77,8 @@ extension ContextAwarePromptToCodeService: PromptToCodeServiceType {
                 language: source.language,
                 documentURL: source.documentURL,
                 projectRootURL: source.projectRootURL,
-                allCode: source.allCode,
+                content: source.content,
+                lines: source.lines,
                 range: source.range
             ),
             isDetached: isDetached,
@@ -86,7 +90,7 @@ extension ContextAwarePromptToCodeService: PromptToCodeServiceType {
 
 public struct PromptToCodeServiceFactoryDependencyKey: DependencyKey {
     public static let liveValue: () -> PromptToCodeServiceType = {
-        OpenAIPromptToCodeService()
+        ContextAwarePromptToCodeService()
     }
 
     public static let previewValue: () -> PromptToCodeServiceType = {

Core/Sources/Service/GUI/ChatTabFactory.swift

@@ -92,7 +92,8 @@ enum ChatTabFactory {
                             language: .plaintext,
                             documentURL: .init(fileURLWithPath: "/"),
                             projectRootURL: .init(fileURLWithPath: "/"),
-                            allCode: prompt,
+                            content: prompt,
+                            lines: prompt.breakLines(),
                             range: .outOfScope
                         ),
                         isDetached: true,

Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift

@@ -424,7 +424,8 @@ extension WindowBaseCommandHandler {
                 usesTabsForIndentation: filespace.codeMetadata.usesTabsForIndentation ?? false,
                 documentURL: fileURL,
                 projectRootURL: workspace.projectRootURL,
-                allCode: editor.content,
+                allCode: editor.content, 
+                allLines: editor.lines,
                 isContinuous: isContinuous,
                 commandName: name,
                 defaultPrompt: newPrompt ?? "",

Core/Sources/SuggestionWidget/FeatureReducers/PromptToCode.swift

@@ -57,6 +57,7 @@ public struct PromptToCode: ReducerProtocol {
         public var projectRootURL: URL
         public var documentURL: URL
         public var allCode: String
+        public var allLines: [String]
         public var extraSystemPrompt: String?
         public var generateDescriptionRequirement: Bool?
         public var commandName: String?
@@ -76,6 +77,7 @@ public struct PromptToCode: ReducerProtocol {
             projectRootURL: URL,
             documentURL: URL,
             allCode: String,
+            allLines: [String],
             commandName: String? = nil,
             description: String = "",
             isResponding: Bool = false,
@@ -101,6 +103,7 @@ public struct PromptToCode: ReducerProtocol {
             self.projectRootURL = projectRootURL
             self.documentURL = documentURL
             self.allCode = allCode
+            self.allLines = allLines
             self.extraSystemPrompt = extraSystemPrompt
             self.generateDescriptionRequirement = generateDescriptionRequirement
             self.isAttachedToSelectionRange = isAttachedToSelectionRange
@@ -165,7 +168,8 @@ public struct PromptToCode: ReducerProtocol {
                                 language: copiedState.language,
                                 documentURL: copiedState.documentURL,
                                 projectRootURL: copiedState.projectRootURL,
-                                allCode: copiedState.allCode,
+                                content: copiedState.allCode,
+                                lines: copiedState.allLines,
                                 range: copiedState.selectionRange ?? .outOfScope
                             ),
                             isDetached: !copiedState.isAttachedToSelectionRange,

Core/Sources/SuggestionWidget/FeatureReducers/PromptToCodeGroup.swift

@@ -31,6 +31,7 @@ public struct PromptToCodeGroup: ReducerProtocol {
         public var documentURL: URL
         public var projectRootURL: URL
         public var allCode: String
+        public var allLines: [String]
         public var isContinuous: Bool
         public var commandName: String?
         public var defaultPrompt: String
@@ -46,6 +47,7 @@ public struct PromptToCodeGroup: ReducerProtocol {
             documentURL: URL,
             projectRootURL: URL,
             allCode: String,
+            allLines: [String],
             isContinuous: Bool,
             commandName: String?,
             defaultPrompt: String,
@@ -60,6 +62,7 @@ public struct PromptToCodeGroup: ReducerProtocol {
             self.documentURL = documentURL
             self.projectRootURL = projectRootURL
             self.allCode = allCode
+            self.allLines = allLines
             self.isContinuous = isContinuous
             self.commandName = commandName
             self.defaultPrompt = defaultPrompt
@@ -100,7 +103,8 @@ public struct PromptToCodeGroup: ReducerProtocol {
                     usesTabsForIndentation: s.usesTabsForIndentation,
                     projectRootURL: s.projectRootURL,
                     documentURL: s.documentURL,
-                    allCode: s.allCode,
+                    allCode: s.allCode, 
+                    allLines: s.allLines,
                     commandName: s.commandName,
                     isContinuous: s.isContinuous,
                     selectionRange: s.selectionRange,

Core/Sources/SuggestionWidget/SuggestionPanelContent/PromptToCodePanel.swift

@@ -430,6 +430,7 @@ struct PromptToCodePanel_Preview: PreviewProvider {
             projectRootURL: URL(fileURLWithPath: "path/to/file.txt"),
             documentURL: URL(fileURLWithPath: "path/to/file.txt"),
             allCode: "",
+            allLines: [],
             commandName: "Generate Code",
             description: "Hello world",
             isResponding: false,
@@ -460,6 +461,7 @@ struct PromptToCodePanel_Error_Detached_Preview: PreviewProvider {
             projectRootURL: URL(fileURLWithPath: "path/to/file.txt"),
             documentURL: URL(fileURLWithPath: "path/to/file.txt"),
             allCode: "",
+            allLines: [],
             commandName: "Generate Code",
             description: "Hello world",
             isResponding: false,

Pro

@@ -85,6 +85,7 @@ public struct EditorInformation {
     }
 
     public static func lines(in code: [String], containing range: CursorRange) -> [String] {
+        guard !code.isEmpty else { return [] }
         let startIndex = min(max(0, range.start.line), code.endIndex - 1)
         let endIndex = min(max(startIndex, range.end.line), code.endIndex - 1)
         let selectedLines = code[startIndex...endIndex]

Tool/Sources/SuggestionModel/EditorInformation.swift

@@ -48,6 +48,10 @@ public struct CursorRange: Codable, Hashable, Sendable, Equatable, CustomStringC
         return start == end
     }
 
+    public var isOneLine: Bool {
+        return start.line == end.line
+    }
+    
     public static func == (lhs: CursorRange, rhs: CursorRange) -> Bool {
         return lhs.start == rhs.start && lhs.end == rhs.end
     }